#!/usr/bin/env python3
"""
Frox AI Morph 1.1 — Export to HuggingFace / vLLM-loadable format

Morph is a custom architecture, not a HuggingFace AutoModel. To serve it
with vLLM in production (see the backend architecture doc), it needs:
  1. A HuggingFace-style directory: config.json, model.safetensors,
     tokenizer files, generation_config.json
  2. A registered `AutoModelForCausalLM` mapping OR a vLLM model plugin
     that knows how to read Morph's config and weight names

This script produces (1) — the HF-format directory — which is the input
vLLM's `--model` flag expects. It does NOT write the vLLM plugin itself
(that's a `vllm.model_executor.models` registration, versioned against a
specific vLLM release), but it lays out weights in a way that a plugin's
`load_weights()` can consume directly since parameter names match the
saved state dict 1:1.

Usage:
    python scripts/export_hf.py --model ./frox-morph-1-1-output/sft_final \
        --output ./frox-morph-1-1-hf
"""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

import torch

from model.architecture.morph_model import MorphForCausalLM
from utils.common import print_banner


def export_to_hf_format(model_path: str, output_path: str, save_dtype: str = "float16"):
    print(f"📂 Loading Morph 1.1 from {model_path}...")
    model = MorphForCausalLM.from_saved(model_path, device="cpu")
    cfg = model.config

    out = Path(output_path)
    out.mkdir(parents=True, exist_ok=True)

    # ── config.json (HF-style, with Morph-specific fields preserved) ──
    hf_config = {
        "architectures": ["MorphForCausalLM"],
        "model_type": "morph",
        "vocab_size": cfg.total_vocab_size,
        "hidden_size": cfg.hidden_size,
        "intermediate_size": cfg.intermediate_size,
        "num_hidden_layers": cfg.num_hidden_layers,
        "num_attention_heads": cfg.num_attention_heads,
        "num_key_value_heads": cfg.num_key_value_heads,
        "head_dim": cfg.head_dim,
        "max_position_embeddings": cfg.max_position_embeddings,
        "rope_theta": cfg.rope_theta,
        "rope_scaling": {
            "type": "yarn",
            "factor": cfg.rope_scaling_factor,
            "original_max_position_embeddings": 2048,
        },
        "rms_norm_eps": cfg.rms_norm_eps,
        "hidden_act": cfg.hidden_act,
        "tie_word_embeddings": cfg.tie_word_embeddings,
        "torch_dtype": save_dtype,
        "pad_token_id": cfg.pad_token_id,
        "bos_token_id": cfg.bos_token_id,
        "eos_token_id": cfg.eos_token_id,
        # Morph-specific — a vLLM plugin reads these to configure attention
        "qk_norm": cfg.qk_norm,
        "use_sliding_window": cfg.use_sliding_window,
        "sliding_window_size": cfg.sliding_window_size,
        "transformers_version": "4.47.0",
        "morph_version": "1.1.0",
    }
    with open(out / "config.json", "w") as f:
        json.dump(hf_config, f, indent=2)
    print(f"  ✓ config.json")

    # ── generation_config.json ──
    gen_config = {
        "bos_token_id": cfg.bos_token_id,
        "eos_token_id": cfg.eos_token_id,
        "pad_token_id": cfg.pad_token_id,
        "do_sample": True,
        "temperature": 0.7,
        "top_p": 0.9,
        "top_k": 50,
        "max_new_tokens": 4096,
        "repetition_penalty": 1.1,
    }
    with open(out / "generation_config.json", "w") as f:
        json.dump(gen_config, f, indent=2)
    print(f"  ✓ generation_config.json")

    # ── model.safetensors ──
    try:
        from safetensors.torch import save_file
        dtype = getattr(torch, save_dtype)
        state_dict = {k: v.to(dtype).contiguous() for k, v in model.state_dict().items()}

        # If weights are tied, safetensors refuses to save duplicate storage —
        # clone the lm_head weight so both tensors have independent storage.
        if cfg.tie_word_embeddings and "lm_head.weight" in state_dict:
            state_dict["lm_head.weight"] = state_dict["lm_head.weight"].clone()

        save_file(state_dict, str(out / "model.safetensors"), metadata={"format": "pt"})
        print(f"  ✓ model.safetensors ({sum(v.numel() for v in state_dict.values()):,} params)")
    except ImportError:
        print("  ⚠ safetensors not installed — saving as model.pt instead")
        torch.save(model.state_dict(), out / "model.pt")

    # ── tokenizer ──
    tokenizer_src = Path(model_path) / "tokenizer"
    if tokenizer_src.exists():
        import shutil
        shutil.copytree(tokenizer_src, out / "tokenizer_files", dirs_exist_ok=True)
        for f in (out / "tokenizer_files").iterdir():
            f.rename(out / f.name)
        (out / "tokenizer_files").rmdir()
        print(f"  ✓ tokenizer files")
    else:
        print(f"  ⚠ No tokenizer found at {tokenizer_src} — copy it manually before serving")

    print(f"\n✅ Exported to {output_path}")
    print(f"\nNext steps:")
    print(f"  1. Register the Morph architecture with vLLM (model_executor/models/morph.py)")
    print(f"  2. Serve: vllm serve {output_path} --dtype {save_dtype} "
          f"--max-model-len {cfg.max_position_embeddings}")


def main():
    parser = argparse.ArgumentParser(description="Export Morph 1.1 to HF/vLLM format")
    parser.add_argument("--model", type=str, required=True)
    parser.add_argument("--output", type=str, required=True)
    parser.add_argument("--dtype", choices=["float16", "bfloat16", "float32"], default="float16")
    args = parser.parse_args()

    print_banner()
    export_to_hf_format(args.model, args.output, args.dtype)


if __name__ == "__main__":
    main()