#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Convert Qwen2.5-VL distributed shards (model_world_size_*_rank_*.pt) into a
Hugging Face-compatible folder with .safetensors and index.
"""

import os
import re
import json
import argparse
import shutil
from typing import Dict, List

import torch
from transformers import AutoConfig, Qwen2_5_VLForConditionalGeneration


# ============================================================
# CLI
# ============================================================
def parse_args():
    p = argparse.ArgumentParser(description="Merge Qwen2.5-VL shards to HF safetensors.")
    p.add_argument("--ckpt_dir", required=True)
    p.add_argument("--config_dir", required=True)
    p.add_argument("--out_dir", required=True)
    p.add_argument("--max_shard_size", default="2000MB")
    p.add_argument("--trust_ckpt", action="store_true")
    p.add_argument("--convert_jinja", action="store_true")
    p.add_argument("--strict", action="store_true")
    p.add_argument("--print_samples", type=int, default=12)
    return p.parse_args()


# ============================================================
# DTensor unwrap helper
# ============================================================
def unwrap_dtensor_recursive(x):
    """Offline-safe unwrap: convert DTensor to its local tensor (no distributed comm)."""
    import torch
    try:
        from torch.distributed.tensor import DTensor
    except Exception:
        DTensor = None

    if DTensor is not None and isinstance(x, DTensor):
        try:
            return x.to_local()  # safe, no comm
        except Exception:
            if hasattr(x, "_local_tensor"):
                return x._local_tensor
            raise

    if isinstance(x, torch.nn.Parameter):
        return unwrap_dtensor_recursive(x.data)
    if isinstance(x, torch.Tensor):
        return x
    if isinstance(x, dict):
        return {k: unwrap_dtensor_recursive(v) for k, v in x.items()}
    if isinstance(x, (list, tuple)):
        t = [unwrap_dtensor_recursive(v) for v in x]
        return type(x)(t)
    return x


# ============================================================
# Loading helpers
# ============================================================
def safe_load(path: str, trust_ckpt: bool):
    """Load a shard, force weights_only=False, unwrap DTensors, and return clean state_dict."""
    import torch
    import torch.distributed.tensor

    try:
        obj = torch.load(path, map_location="cpu", weights_only=False)
    except Exception as e:
        if not trust_ckpt:
            raise RuntimeError(f"❌ torch.load failed for {path}: {e}")
        print(f"[safe_load] retry with weights_only=False (trusted ckpt) for {path}")
        obj = torch.load(path, map_location="cpu", weights_only=False)

    obj = unwrap_dtensor_recursive(obj)

    def strip_prefix(k: str) -> str:
        return k.replace("module.", "").replace("_orig_mod.", "")

    if isinstance(obj, dict):
        if "state_dict" in obj and isinstance(obj["state_dict"], dict):
            sd = obj["state_dict"]
        elif "module" in obj and isinstance(obj["module"], dict):
            sd = obj["module"]
        else:
            sd = obj
    else:
        sd = obj

    return {strip_prefix(k): unwrap_dtensor_recursive(v) for k, v in sd.items()}


def find_shards(ckpt_dir: str) -> List[str]:
    files = sorted(f for f in os.listdir(ckpt_dir) if re.match(r"model_world_size_\d+_rank_\d+\.pt$", f))
    if not files:
        raise FileNotFoundError("No shards found: model_world_size_*_rank_*.pt")
    return [os.path.join(ckpt_dir, f) for f in files]


# ============================================================
# Merge logic
# ============================================================
def shapes_match_except_dim(a: torch.Tensor, b: torch.Tensor, dim: int) -> bool:
    if a.ndim != b.ndim:
        return False
    for i in range(a.ndim):
        if i == dim:
            continue
        if a.shape[i] != b.shape[i]:
            return False
    return True


def try_concat(tensors: List[torch.Tensor]) -> torch.Tensor:
    if len(tensors) == 1:
        return tensors[0]
    ref = tensors[0]
    if all(shapes_match_except_dim(ref, t, 0) for t in tensors[1:]):
        return torch.cat(tensors, dim=0)
    if ref.ndim >= 2 and all(shapes_match_except_dim(ref, t, 1) for t in tensors[1:]):
        return torch.cat(tensors, dim=1)
    if ref.ndim >= 1 and all(shapes_match_except_dim(ref, t, ref.ndim - 1) for t in tensors[1:]):
        return torch.cat(tensors, dim=-1)
    raise RuntimeError("No valid concat dimension among (0,1,-1)")


def merge_shards(shard_paths: List[str], trust_ckpt: bool) -> Dict[str, torch.Tensor]:
    print("[info] shards:", [os.path.basename(p) for p in shard_paths])
    parts = [safe_load(p, trust_ckpt) for p in shard_paths]
    keys = list(parts[0].keys())
    merged: Dict[str, torch.Tensor] = {}
    n = len(parts)

    # ===== hints for concat along dim=0 =====
    SHARD_DIM0_HINTS = (
        # language blocks
        ".self_attn.q_proj.", ".self_attn.k_proj.", ".self_attn.v_proj.", ".self_attn.o_proj.",
        ".mlp.gate_proj.", ".mlp.up_proj.", ".mlp.down_proj.",
        "lm_head.weight", "embed_tokens.weight",
        ".input_layernorm.weight", ".post_attention_layernorm.weight", ".norm.weight",
        # visual blocks
        "model.visual.blocks.", "model.visual.merger.", "model.visual.patch_embed."
    )

    def prefer_concat_dim0(name: str) -> bool:
        return any(h in name for h in SHARD_DIM0_HINTS)

    for k in keys:
        tensors = [sd[k] for sd in parts if k in sd]
        if not tensors:
            continue
        if len(tensors) != n:
            merged[k] = tensors[0]
            continue

        same_shape = all(t.shape == tensors[0].shape for t in tensors[1:])
        if same_shape:
            if tensors[0].ndim >= 1 and prefer_concat_dim0(k):
                merged[k] = torch.cat(tensors, dim=0)
            else:
                merged[k] = tensors[0]
            continue

        try:
            merged[k] = try_concat(tensors)
        except Exception:
            merged[k] = tensors[0]

    print(f"[merge] merged parameters: {len(merged)}")
    return merged


# ============================================================
# Utility
# ============================================================
def ensure_config_ok(config_dir: str):
    if not os.path.isdir(config_dir):
        raise FileNotFoundError(f"config_dir not found: {config_dir}")
    cfg = os.path.join(config_dir, "config.json")
    if not os.path.exists(cfg):
        raise FileNotFoundError(f"config.json missing in {config_dir}")


def maybe_convert_jinja(config_dir: str, out_dir: str):
    src = os.path.join(config_dir, "chat_template.jinja")
    dst = os.path.join(out_dir, "chat_template.json")
    if os.path.exists(src) and not os.path.exists(dst):
        with open(src, "r", encoding="utf-8") as f:
            jinja = f.read()
        with open(dst, "w", encoding="utf-8") as f:
            json.dump({"chat_template": jinja}, f, ensure_ascii=False, indent=2)
        print("[convert] chat_template.jinja -> chat_template.json")


def copy_side_files(config_dir: str, out_dir: str):
    names = [
        "config.json", "tokenizer.json", "tokenizer_config.json",
        "merges.txt", "vocab.json", "special_tokens_map.json",
        "added_tokens.json", "preprocessor_config.json",
        "generation_config.json", "video_preprocessor_config.json",
    ]
    for name in names:
        src = os.path.join(config_dir, name)
        dst = os.path.join(out_dir, name)
        if os.path.exists(src) and not os.path.exists(dst):
            try:
                shutil.copy2(src, dst)
                print(f"[copy] {name}")
            except Exception as e:
                print(f"[copy] skip {name}: {e}")


# ============================================================
# Main
# ============================================================
def main():
    args = parse_args()
    ensure_config_ok(args.config_dir)
    os.makedirs(args.out_dir, exist_ok=True)

    shard_paths = find_shards(args.ckpt_dir)
    merged = merge_shards(shard_paths, args.trust_ckpt)

    merged = unwrap_dtensor_recursive(merged)
    for k, v in list(merged.items()):
        if isinstance(v, torch.Tensor) and not v.is_contiguous():
            merged[k] = v.contiguous()

    cfg = AutoConfig.from_pretrained(args.config_dir)
    model = Qwen2_5_VLForConditionalGeneration(config=cfg)

    missing, unexpected = model.load_state_dict(merged, strict=args.strict)
    print(f"[load_state_dict] Missing: {len(missing)} | Unexpected: {len(unexpected)}")
    if args.print_samples and missing:
        print("  sample missing:", missing[:args.print_samples])
    if args.print_samples and unexpected:
        print("  sample unexpected:", unexpected[:args.print_samples])

    model.save_pretrained(args.out_dir, safe_serialization=True, max_shard_size=args.max_shard_size)
    print(f"[save] ✅ wrote safetensors to: {args.out_dir}")

    if args.convert_jinja:
        maybe_convert_jinja(args.config_dir, args.out_dir)
    copy_side_files(args.config_dir, args.out_dir)
    print("[done] all good.")


if __name__ == "__main__":
    main()
