#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import shutil
from pathlib import Path

import torch
import torch.distributed as dist
import torch.distributed.checkpoint as dcp
from safetensors.torch import save_file


CONFIG = {
    "architectures": ["LaoyaoMoeForCausalLM"],
    "auto_map": {
        "AutoConfig": "configuration_laoyao_moe.LaoyaoMoeConfig",
        "AutoModelForCausalLM": "modeling_laoyao_moe.LaoyaoMoeForCausalLM",
    },
    "model_type": "laoyao_moe",
    "vocab_size": 154856,
    "hidden_size": 1536,
    "num_hidden_layers": 12,
    "num_attention_heads": 24,
    "num_key_value_heads": 4,
    "intermediate_size": 4608,
    "moe_intermediate_size": 6144,
    "num_experts": 12,
    "num_experts_per_tok": 4,
    "moe_expert_capacity_factor": 1.25,
    "moe_layers": [2, 4, 6, 8, 10],
    "max_position_embeddings": 8192,
    "rope_theta": 10000.0,
    "rms_norm_eps": 1e-5,
    "tie_word_embeddings": True,
    "use_cache": False,
    "torch_dtype": "bfloat16",
}


def init_dist() -> None:
    if not dist.is_initialized():
        dist.init_process_group("gloo", init_method="tcp://127.0.0.1:29691", rank=0, world_size=1)


def load_dcp_tensors(checkpoint_dir: Path) -> dict[str, torch.Tensor]:
    init_dist()
    metadata = dcp.FileSystemReader(str(checkpoint_dir)).read_metadata()
    state = {}
    for key, meta in metadata.state_dict_metadata.items():
        if key.startswith(("optimizer.", "opt_param_scheduler", "rng_state", "iteration")):
            continue
        if not hasattr(meta, "size"):
            continue
        state[key] = torch.empty(tuple(meta.size), dtype=meta.properties.dtype)
    print(f"[load] tensors={len(state)} checkpoint={checkpoint_dir}", flush=True)
    dcp.load(state, checkpoint_id=str(checkpoint_dir))
    return state


def convert_state(src: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
    dst = {
        "model.embed_tokens.weight": src["embedding.word_embeddings.weight"].contiguous(),
        "model.norm.weight": src["decoder.final_layernorm.weight"].contiguous(),
    }
    for layer in range(CONFIG["num_hidden_layers"]):
        sp = f"decoder.layers.{layer}"
        dp = f"model.layers.{layer}"
        dst[f"{dp}.input_layernorm.weight"] = src[f"{sp}.self_attention.linear_qkv.layer_norm_weight"].contiguous()
        dst[f"{dp}.self_attn.qkv_proj.weight"] = src[f"{sp}.self_attention.linear_qkv.weight"].contiguous()
        dst[f"{dp}.self_attn.o_proj.weight"] = src[f"{sp}.self_attention.linear_proj.weight"].contiguous()
        if layer in CONFIG["moe_layers"]:
            dst[f"{dp}.pre_mlp_layernorm.weight"] = src[f"{sp}.pre_mlp_layernorm.weight"].contiguous()
            dst[f"{dp}.mlp.router.weight"] = src[f"{sp}.mlp.router.weight"].contiguous()
            dst[f"{dp}.mlp.experts_fc1"] = src[f"{sp}.mlp.experts.experts.linear_fc1.weight"].contiguous()
            dst[f"{dp}.mlp.experts_fc2"] = src[f"{sp}.mlp.experts.experts.linear_fc2.weight"].contiguous()
        else:
            dst[f"{dp}.pre_mlp_layernorm.weight"] = src[f"{sp}.mlp.linear_fc1.layer_norm_weight"].contiguous()
            dst[f"{dp}.mlp.fc1.weight"] = src[f"{sp}.mlp.linear_fc1.weight"].contiguous()
            dst[f"{dp}.mlp.fc2.weight"] = src[f"{sp}.mlp.linear_fc2.weight"].contiguous()
    return dst


def copy_tokenizer(tokenizer_dir: Path, output_dir: Path) -> None:
    for path in tokenizer_dir.iterdir():
        if path.is_file():
            shutil.copy2(path, output_dir / path.name)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint-dir", required=True)
    parser.add_argument("--tokenizer-dir", default="tokenizer/glm5.2")
    parser.add_argument("--output-dir", required=True)
    args = parser.parse_args()

    checkpoint_dir = Path(args.checkpoint_dir)
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    src = load_dcp_tensors(checkpoint_dir)
    dst = convert_state(src)
    print(f"[save] hf_tensors={len(dst)} output={output_dir}", flush=True)
    save_file(dst, str(output_dir / "model.safetensors"), metadata={"format": "pt"})

    copy_tokenizer(Path(args.tokenizer_dir), output_dir)
    (output_dir / "config.json").write_text(json.dumps(CONFIG, indent=2, ensure_ascii=False) + "\n")
    shutil.copy2(Path(__file__).with_name("configuration_laoyao_moe.py"), output_dir / "configuration_laoyao_moe.py")
    shutil.copy2(Path(__file__).with_name("modeling_laoyao_moe.py"), output_dir / "modeling_laoyao_moe.py")
    print("[done]", flush=True)


if __name__ == "__main__":
    main()
