from contextlib import contextmanager
import datetime
import importlib
import json
import logging
import os
import shutil
import time
import threading

from collections import OrderedDict
import numpy as np
import torch
from torch.distributed import barrier, is_initialized
import torch.distributed.checkpoint as dcp
from torch.distributed.checkpoint.state_dict import (
    get_state_dict,
    set_state_dict,
    get_model_state_dict,
    get_optimizer_state_dict,
    StateDictOptions,
)
from torch.distributed.fsdp import (
    FullyShardedDataParallel as FSDP,
    FullStateDictConfig,
    FullOptimStateDictConfig,
    StateDictType,
)


if (
    importlib.util.find_spec("torch.nn.attention") is not None
    and importlib.util.find_spec("torch.nn.attention.sdpa_kernel") is not None
):
    TORCH_IS_NIGHTLY = True
else:
    TORCH_IS_NIGHTLY = False


def is_ddp():
    return int(os.environ.get("RANK", -1)) != -1


def is_master():
    if is_ddp():
        return int(os.environ["RANK"]) == 0
    return True


def is_main_gpu_on_node():
    if is_ddp():
        return int(os.environ["LOCAL_RANK"]) == 0
    return True


def print_with_time(content):
    """Print the content with the current time."""
    print(f"[{datetime.datetime.now().strftime('%Y-%m-%d_%H:%M:%S')}]: {content}")


def print_with_time_master(content):
    if is_master():
        print_with_time(content)


def _write_json(info, out_fp):
    if is_master():
        with open(out_fp, "w") as f:
            json.dump(info, f)


def _read_json(fp):
    with open(fp, "r") as f:
        info = json.load(f)
    return info


def hash_string_to_number(s, max_val=1023):
    return hash(s) % (max_val + 1)


@contextmanager
def suppress_logging(highest_level=logging.CRITICAL):
    previous_level = logging.root.manager.disable
    logging.disable(highest_level)
    try:
        yield
    finally:
        logging.disable(previous_level)


def dist_barrier():
    if is_initialized():
        barrier()


def _compare_dicts(
    dict1,
    dict2,
    ignore_set="compatibility",
    preload_strict=True,
):
    if isinstance(ignore_set, str) and ignore_set == "compatibility":
        ignore_set = set(
            [
                "rope_theta",
                "dropout",
                "attention_type",
                "coarse_infer_token",
                "coarse_pad_token",
                "semantic_pad_token",
                "semantic_infer_token",
                "use_text_loss",
                "text_pad_token",
            ]
        )
    if is_master():
        missing_keys = set(dict1.keys()) - set(dict2.keys()) - ignore_set
        if len(missing_keys) > 0:
            if preload_strict:
                raise ValueError(f"checkpoint missing config keys: {missing_keys}")
            else:
                print_with_time_master(f"warning: checkpoint missing config keys: {missing_keys}")
        extra_keys = set(dict2.keys()) - set(dict1.keys()) - ignore_set
        if len(extra_keys) > 0:
            if preload_strict:
                raise ValueError(f"checkpoint has extra config keys: {extra_keys}")
            else:
                print_with_time_master(f"warning: checkpoint has extra config keys: {extra_keys}")
        for k in set(dict1.keys()) - ignore_set:
            if dict1.get(k) != dict2.get(k):
                if preload_strict:
                    raise ValueError(f"checkpoint config key mismatch on {k}")
                else:
                    print_with_time_master(f"warning: checkpoint config key mismatch on {k}")


def verify_preload_model_args(model_args, preload_checkpoint, preload_strict=True):
    checkpoint_model_args = _read_json(preload_checkpoint.rstrip("/") + ".json")["model_args"]
    _compare_dicts(model_args, checkpoint_model_args, preload_strict=preload_strict)


def _compare_partial_content(file1, file2, num_samples=10):
    # fast way to approximately compare two large files (should be fine for state dicts)
    print_with_time_master(
        "careful, using approximation for checkpoint loading. could be wrong in principle"
    )
    stat1 = os.stat(file1)
    stat2 = os.stat(file2)
    if stat1.st_size != stat2.st_size:
        return False

    interval = stat1.st_size // num_samples
    with open(file1, "rb") as f1, open(file2, "rb") as f2:
        for i in range(num_samples):
            offset = i * interval
            f1.seek(offset)
            f2.seek(offset)
            if f1.read(1024) != f2.read(1024):  # Read 1 KB at each sample point
                return False
    return True


def load_old_state_dict(
    model_args,
    preload_checkpoint,
    model,
    local_cache_dir,
    preload_strict=True,
    use_mmap=False,
):
    if local_cache_dir is not None:
        local_ckpt_fp = os.path.join(local_cache_dir, "ckpt.pt")
        print_with_time_master("verifying model args...")
        if is_master():
            if os.path.exists(local_ckpt_fp) and _compare_partial_content(
                preload_checkpoint, local_ckpt_fp
            ):
                checkpoint_model_args = torch.load(local_ckpt_fp, mmap=use_mmap)["model_args"]
            else:
                checkpoint_model_args = torch.load(preload_checkpoint, mmap=use_mmap)["model_args"]
            _compare_dicts(model_args, checkpoint_model_args, preload_strict=preload_strict)
            del checkpoint_model_args
        dist_barrier()
        if is_main_gpu_on_node():
            os.makedirs(local_cache_dir, exist_ok=True)
            if not os.path.exists(local_ckpt_fp) or not _compare_partial_content(
                preload_checkpoint, local_ckpt_fp
            ):
                print_with_time_master("copying checkpoint file to local cachedir...")
                # If the file is large, consider using a buffered copy
                with open(preload_checkpoint, "rb") as src, open(local_ckpt_fp, "wb") as dst:
                    shutil.copyfileobj(src, dst, length=1024 * 1024)  # 1MB buffer
    else:
        local_ckpt_fp = preload_checkpoint
    dist_barrier()
    for n_gpu in range(torch.cuda.device_count()):
        print_with_time_master(f"loading model state_dict on gpu {n_gpu}")
        if n_gpu == int(os.environ["LOCAL_RANK"]):
            state_dict = torch.load(local_ckpt_fp, mmap=use_mmap)["model"]
            model_state_dict = model.state_dict()
            for k in model_state_dict.keys():
                if (
                    k in model_state_dict
                    and k in state_dict
                    and model_state_dict[k].shape != state_dict[k].shape
                ):
                    print_with_time_master(f"warning: state dict shape mismatch on {k}")
                    del state_dict[k]
                if k not in state_dict:
                    print_with_time_master(f"warning: state dict missing key {k}")

            # handle old heads
            if "transformer.wte_text.weight" in state_dict:
                state_dict["input_modules.text_input.wte_text.weight"] = state_dict.pop(
                    "transformer.wte_text.weight"
                )
            if "transformer.ln_text.weight" in state_dict:
                state_dict["input_modules.text_input.ln_text.weight"] = state_dict.pop(
                    "transformer.ln_text.weight"
                )
            if "transformer.wte_semantic.0.weight" in state_dict:
                state_dict["input_modules.semantic_input_0.wte_semantic.weight"] = state_dict.pop(
                    "transformer.wte_semantic.0.weight"
                )
            if "transformer.ln_semantic.weight" in state_dict:
                state_dict["input_modules.semantic_input_0.ln_semantic.weight"] = state_dict.pop(
                    "transformer.ln_semantic.weight"
                )

            if "input_modules.semantic_input.wte_semantic.weight" in state_dict:
                state_dict["input_modules.semantic_input_0.wte_semantic.weight"] = state_dict.pop(
                    "input_modules.semantic_input.wte_semantic.weight"
                )
            if "input_modules.semantic_input.ln_semantic.weight" in state_dict:
                state_dict["input_modules.semantic_input_0.ln_semantic.weight"] = state_dict.pop(
                    "input_modules.semantic_input.ln_semantic.weight"
                )
            if "output_modules.semantic_output.lm_head.weight" in state_dict:
                state_dict["output_modules.semantic_output_0.lm_head.weight"] = state_dict.pop(
                    "output_modules.semantic_output.lm_head.weight"
                )

            if "use_text_loss" not in model_args:  # old model, copy over old heads
                if "lm_heads.0.weight" in state_dict:
                    state_dict["output_modules.semantic_output_0.lm_head.weight"] = state_dict[
                        "lm_heads.0.weight"
                    ]
            elif model_args.get("use_text_loss"):
                if "lm_heads.0.weight" in state_dict:
                    state_dict["output_modules.text_output.lm_head.weight"] = state_dict.pop(
                        "lm_heads.0.weight"
                    )
                if "lm_heads.1.weight" in state_dict:
                    state_dict["output_modules.semantic_output_0.lm_head.weight"] = state_dict.pop(
                        "lm_heads.1.weight"
                    )
            else:
                if "lm_heads.0.weight" in state_dict:
                    state_dict["output_modules.semantic_output_0.lm_head.weight"] = state_dict.pop(
                        "lm_heads.0.weight"
                    )

            # can't do assign=True, has some really weird effects in initialization
            model.load_state_dict(state_dict, strict=preload_strict)
            del state_dict, model_state_dict
        dist_barrier()


def load_old_optimizer_state_dict(
    model,
    optimizer,
    ckpt_fp,
):
    assert os.path.exists(ckpt_fp)
    dist_barrier()
    iter_num = 0
    n_tokens = 0
    best_val_loss = 1e9

    # Check if using FSDP2 or FSDP1
    is_fsdp2 = not isinstance(model, FSDP)

    if is_fsdp2:
        # FSDP2: Load checkpoint on all ranks and use set_optimizer_state_dict
        checkpoint = torch.load(ckpt_fp, weights_only=False)
        optim_state_dict = checkpoint["optimizer"]
        iter_num = checkpoint["iter_num"]
        n_tokens = checkpoint.get("n_tokens", 0)
        best_val_loss = checkpoint["best_val_loss"]

        # Use FSDP2 API to load optimizer state
        # Note: full_state_dict=True tells it the loaded state is a full (gathered) state dict
        from torch.distributed.checkpoint.state_dict import set_optimizer_state_dict

        set_optimizer_state_dict(
            model,
            optimizer,
            optim_state_dict=optim_state_dict,
            options=StateDictOptions(strict=False, full_state_dict=True),
        )
        del checkpoint, optim_state_dict
    else:
        # FSDP1: Use old API with sequential loading per GPU
        optim_state_dict = None
        for n_gpu in range(torch.cuda.device_count()):
            print_with_time_master(f"loading optimizer state_dict on gpu {n_gpu}")
            if n_gpu == int(os.environ["LOCAL_RANK"]):
                checkpoint = torch.load(ckpt_fp, weights_only=False)
                optim_state_dict = FSDP.optim_state_dict_to_load(
                    model, optimizer, checkpoint["optimizer"]
                )
                iter_num = checkpoint["iter_num"]
                n_tokens = checkpoint.get("n_tokens", 0)
                best_val_loss = checkpoint["best_val_loss"]
                del checkpoint
            dist_barrier()
        assert optim_state_dict is not None
        optimizer.load_state_dict(optim_state_dict)
        del optim_state_dict

    return iter_num, n_tokens, best_val_loss


# DCP save info:
#  dcp save a directory eg 'last_ckpt' which needs to be loaded with dcp
#  additionally we save a json file at eg 'last_ckpt.json' to contain model args etc for verification


def load_checkpoint(
    preload_checkpoint,
    preload_optimizer,
    model,
    optimizer,
):
    dist_barrier()
    iter_num = 0
    n_tokens = 0
    best_val_loss = 1e9
    print_with_time_master(f"loading checkpoint from {preload_checkpoint}")
    model_state_dict, optimizer_state_dict = get_state_dict(model, optimizer)
    dist_barrier()
    state_dict = {"model": model_state_dict}
    if preload_optimizer:
        state_dict["optimizer"] = optimizer_state_dict
    dcp.load(
        state_dict=state_dict,
        checkpoint_id=preload_checkpoint,
    )
    dist_barrier()
    # sets our state dicts on the model and optimizer, now that we've loaded
    set_state_dict(
        model,
        (),
        model_state_dict=model_state_dict,
        optim_state_dict=optimizer_state_dict if preload_optimizer else {},
    )
    dist_barrier()
    if preload_optimizer:
        info = _read_json(preload_checkpoint.rstrip("/") + ".json")
        iter_num = info["iter_num"]
        n_tokens = info.get("n_tokens", 0)
        best_val_loss = info["best_val_loss"]
    dist_barrier()
    print_with_time_master("done loading checkpoint")
    return iter_num, n_tokens, best_val_loss


def _write_dcp_checkpoint(checkpoint, checkpoint_info, out_dir, ckpt_name):
    _write_json(checkpoint_info, os.path.join(out_dir, f"{ckpt_name}.json"))
    dist_barrier()
    if TORCH_IS_NIGHTLY:
        dcp.save(checkpoint, checkpoint_id=os.path.join(out_dir, ckpt_name))
    else:
        fs_storage_writer = torch.distributed.checkpoint.FileSystemWriter(
            os.path.join(out_dir, ckpt_name)
        )
        dcp.save(checkpoint, fs_storage_writer)
    dist_barrier()


def save_checkpoint(
    out_dir,
    model,
    optimizer,
    best_val_loss,
    current_val_loss,
    step_save_iters,
    time_since_last_loss,
    model_args=None,
    iter_num=None,
    n_tokens=None,
    debug_val_only=False,
):
    if debug_val_only:
        return min(current_val_loss, best_val_loss)
    dist_barrier()
    t_save = time.time()
    model_state, optim_state = get_state_dict(model, optimizer)
    print_with_time_master(f"saving checkpoint to {out_dir}")
    checkpoint_info = {
        "model_args": model_args,
        "iter_num": iter_num,
        "n_tokens": n_tokens,
        "best_val_loss": current_val_loss,
    }
    checkpoint = {
        "model": model_state,
        "optimizer": optim_state,
    }
    if current_val_loss < best_val_loss:
        _write_dcp_checkpoint(checkpoint, checkpoint_info, out_dir, "best_ckpt")
    _write_dcp_checkpoint(checkpoint, checkpoint_info, out_dir, "last_ckpt")
    _write_dcp_checkpoint({"model": model_state}, checkpoint_info, out_dir, "last_ckpt_infer")
    if iter_num % step_save_iters == 0:
        _write_dcp_checkpoint(checkpoint, checkpoint_info, out_dir, f"step_{iter_num}_ckpt")
    estimation_time = time.time() - t_save
    eval_time_pct = np.clip(estimation_time / time_since_last_loss * 100, 0, 100)
    del model_state, optim_state
    torch.cuda.empty_cache()
    if current_val_loss < best_val_loss:
        best_val_loss = current_val_loss
    dist_barrier()
    print_with_time_master(f"saving took {estimation_time:.1f} seconds. ({eval_time_pct:.1f}% of loop)")
    return best_val_loss


def save_old_checkpoint(
    out_dir,
    model,
    optimizer,
    best_val_loss,
    current_val_loss,
    step_save_iters,
    time_since_last_loss,
    model_args=None,
    iter_num=None,
    n_tokens=None,
    debug_val_only=False,
    save_best_ckpt=True,
    save_last_ckpt=True,
):
    if debug_val_only:
        return min(current_val_loss, best_val_loss)
    dist_barrier()
    t_save = time.time()

    # Check if using FSDP2 (no FSDP wrapper) or FSDP1 (has FSDP wrapper)
    is_fsdp2 = not isinstance(model, FSDP)

    if is_fsdp2:
        # FSDP2: Use new state_dict APIs
        options = StateDictOptions(
            full_state_dict=True,
            cpu_offload=True,
        )
        model_state = get_model_state_dict(model, options=options)
        optim_state = get_optimizer_state_dict(model, optimizer, options=options)
    else:
        # FSDP1: Use old APIs
        FSDP.set_state_dict_type(
            model,
            StateDictType.FULL_STATE_DICT,
            FullStateDictConfig(rank0_only=True, offload_to_cpu=True),
            FullOptimStateDictConfig(rank0_only=True, offload_to_cpu=True),
        )
        model_state = model.state_dict()
        original_osd = optimizer.state_dict()
        optim_state = FSDP.optim_state_dict(model, optimizer, optim_state_dict=original_osd)
    if is_master():  # only write state dicts on rank 0
        assert model_state is not None, "some sort of distributed bug"
        assert optim_state is not None, "some sort of distributed bug"
        checkpoint = {
            "model": model_state,
            "optimizer": optim_state,
            "model_args": model_args,
            "iter_num": iter_num,
            "n_tokens": n_tokens,
            "best_val_loss": current_val_loss,
        }

        print_with_time(f"saving checkpoint to {out_dir}")
        if current_val_loss < best_val_loss and save_best_ckpt:
            torch.save(checkpoint, os.path.join(out_dir, "best_ckpt.pt"))

            def copy_checkpoint():
                shutil.copy(os.path.join(out_dir, "best_ckpt.pt"), os.path.join(out_dir, "last_ckpt.pt"))

            # Start a new thread to copy best_ckpt.pt to last_ckpt.pt -- safe a bit of time
            thread = threading.Thread(target=copy_checkpoint)
            thread.start()
        elif save_last_ckpt:
            torch.save(checkpoint, os.path.join(out_dir, "last_ckpt.pt"))
        infer_sd = {k: checkpoint[k] for k in ["model", "model_args", "best_val_loss"]}
        infer_sd["model"] = OrderedDict({k: v.to(torch.bfloat16) for k, v in infer_sd["model"].items()})
        # probably always want to save the last infer checkpoint
        torch.save(infer_sd, os.path.join(out_dir, "last_ckpt_infer.pt"))
        if iter_num % step_save_iters == 0 and iter_num > 0:
            torch.save(
                checkpoint,
                os.path.join(out_dir, f"step_{iter_num}_ckpt.pt"),
            )
            torch.save(infer_sd, os.path.join(out_dir, f"step_{iter_num}_infer.pt"))
    estimation_time = time.time() - t_save
    eval_time_pct = np.clip(estimation_time / time_since_last_loss * 100, 0, 100)
    del model_state, optim_state
    torch.cuda.empty_cache()
    if current_val_loss < best_val_loss:
        best_val_loss = current_val_loss
    dist_barrier()
    print_with_time_master(f"saving took {estimation_time:.1f} seconds. ({eval_time_pct:.1f}% of loop)")
    return best_val_loss