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
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.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.ln_semantic.weight"] = state_dict.pop(
                    "transformer.ln_semantic.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.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.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.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
    optim_state_dict = None
    # TODO: this uses shard not scatter (means we need a copy for all optmizers), cause scatter is bugged?
    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()
    # TODO: maybe rank0_only=True and offload_to_cpu=True
    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, original_osd
    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_dual_model_checkpoint(
    out_dir,
    codec_model,
    discriminator_model,
    generator_optimizer,
    discriminator_optimizer,
    best_val_loss,
    current_val_loss,
    step_save_iters,
    time_since_last_loss,
    codec_args=None,
    discriminator_args=None,
    iter_num=None,
    n_hours=None,
    debug_val_only=False,
    save_best_ckpt=True,
    save_periodic_ckpt=True,
):
    """
    Save checkpoint for dual model training (codec + discriminator).
    Properly handles FSDP and DDP state dict extraction.
    """
    if debug_val_only:
        return min(current_val_loss, best_val_loss)
    
    dist_barrier()
    t_save = time.time()
    
    # Check if models are FSDP wrapped
    codec_is_fsdp = isinstance(codec_model, FSDP)
    disc_is_fsdp = isinstance(discriminator_model, FSDP)
    
    if codec_is_fsdp or disc_is_fsdp:
        # For FSDP models, use proper state dict configuration
        if codec_is_fsdp:
            FSDP.set_state_dict_type(
                codec_model,
                StateDictType.FULL_STATE_DICT,
                FullStateDictConfig(rank0_only=True, offload_to_cpu=True),
                FullOptimStateDictConfig(rank0_only=True, offload_to_cpu=True),
            )
        if disc_is_fsdp:
            FSDP.set_state_dict_type(
                discriminator_model,
                StateDictType.FULL_STATE_DICT,
                FullStateDictConfig(rank0_only=True, offload_to_cpu=True),
                FullOptimStateDictConfig(rank0_only=True, offload_to_cpu=True),
            )
    
    # Extract model state dicts (handles both FSDP and DDP/regular models)
    if codec_is_fsdp:
        codec_state = codec_model.state_dict()
    else:
        # For DDP, unwrap with .module, for regular models use directly
        codec_state = getattr(codec_model, 'module', codec_model).state_dict()
    
    if disc_is_fsdp:
        disc_state = discriminator_model.state_dict()  
    else:
        # For DDP, unwrap with .module, for regular models use directly
        disc_state = getattr(discriminator_model, 'module', discriminator_model).state_dict()
    
    # Extract optimizer state dicts
    if codec_is_fsdp:
        gen_optim_state = FSDP.optim_state_dict(codec_model, generator_optimizer)
    else:
        gen_optim_state = generator_optimizer.state_dict()
        
    if disc_is_fsdp:
        disc_optim_state = FSDP.optim_state_dict(discriminator_model, discriminator_optimizer)
    else:
        disc_optim_state = discriminator_optimizer.state_dict()
    
    if is_master():  # only write state dicts on rank 0
        assert codec_state is not None, "codec model state dict is None"
        assert disc_state is not None, "discriminator model state dict is None"
        assert gen_optim_state is not None, "generator optimizer state dict is None"
        assert disc_optim_state is not None, "discriminator optimizer state dict is None"
        
        checkpoint_dict = {
            "iter_num": iter_num,
            "best_val_loss": current_val_loss,
            "codec_model": codec_state,
            "discriminator_model": disc_state,
            "generator_optimizer": gen_optim_state,
            "discriminator_optimizer": disc_optim_state,
            "codec_args": codec_args,
            "discriminator_args": discriminator_args,
            "n_hours": n_hours,
        }

        print_with_time(f"saving dual model checkpoint to {out_dir}")
        
        # Save best checkpoint
        if current_val_loss < best_val_loss and save_best_ckpt:
            checkpoint_path = os.path.join(out_dir, "best_checkpoint.pt")
            torch.save(checkpoint_dict, checkpoint_path)
            print_with_time(f"saved best checkpoint to {checkpoint_path}")
        
        # Save periodic checkpoint
        if save_periodic_ckpt and iter_num % step_save_iters == 0 and iter_num > 0:
            checkpoint_path = os.path.join(out_dir, f"checkpoint_iter_{iter_num}.pt")
            torch.save(checkpoint_dict, checkpoint_path)  
            print_with_time(f"saved checkpoint to {checkpoint_path}")
    
    estimation_time = time.time() - t_save
    eval_time_pct = np.clip(estimation_time / time_since_last_loss * 100, 0, 100)
    del codec_state, disc_state, gen_optim_state, disc_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"checkpoint saving took {estimation_time:.1f} seconds. ({eval_time_pct:.1f}% of loop)")
    return best_val_loss