"""
Export TinyModel weights → GGUF INT4 (Q4_K_M).

Usage:
  python3 scripts/quantize_gguf.py                          # from outputs/tiny-sft/final/model.pt
  python3 scripts/quantize_gguf.py --checkpoint path/to/model.pt
  python3 scripts/quantize_gguf.py --checkpoint path/to/model.pt --output model.gguf
"""

import os, sys, argparse, json
import torch
import gguf

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
from scripts.model_tiny import TinyModel


def export_gguf(checkpoint_path, output_path):
    print(f" Loading checkpoint: {checkpoint_path}")

    model = TinyModel(
        vocab_size=1757, hidden=128, intermediate=640,
        num_layers=3, num_heads=8, num_kv_heads=4,
        max_seq_len=2048, tie_weights=True,
    )
    state = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
    model.load_state_dict(state)
    model.eval()

    n = sum(p.numel() for p in model.parameters())
    print(f" Params: {n:,}")

    print(f" Writing GGUF: {output_path}")
    writer = gguf.GGUFWriter(output_path, "tiny")

    # Metadata
    writer.add_context_length(2048)
    writer.add_embedding_length(model.hidden)
    writer.add_block_count(len(model.blocks))
    writer.add_head_count(model.blocks[0].attn.num_heads)
    writer.add_head_count_kv(model.blocks[0].attn.num_kv_heads)
    writer.add_feed_forward_length(model.blocks[0].mlp.up.weight.shape[0])
    writer.add_layer_norm_rms_eps(1e-6)

    # Tensor names for llama-like GGUF format
    name_map = {
        "token_embed.weight": "token_embd.weight",
        "ln_f.weight": "output_norm.weight",
        "lm_head.weight": "output.weight",
    }

    def tensor_name(key):
        parts = key.split(".")
        if parts[0] == "blocks":
            blk = int(parts[1])
            sub = parts[2]
            if sub == "ln1":
                return f"blk.{blk}.attn_norm.{parts[3]}"
            elif sub == "ln2":
                return f"blk.{blk}.ffn_norm.{parts[3]}"
            elif sub == "attn":
                proj_map = {
                    "q_proj": "attn_q",
                    "k_proj": "attn_k",
                    "v_proj": "attn_v",
                    "o_proj": "attn_output",
                }
                return f"blk.{blk}.{proj_map[parts[3]]}.weight"
            elif sub == "mlp":
                proj_map = {
                    "gate": "ffn_gate",
                    "up": "ffn_up",
                    "down": "ffn_down",
                }
                return f"blk.{blk}.{proj_map[parts[3]]}.weight"
        return name_map.get(key, key)

    # Write all tensors as fp32 first, GGUF will quantize
    for key, param in model.state_dict().items():
        tname = tensor_name(key)
        data = param.contiguous().float().numpy()
        writer.add_tensor(tname, data)

    writer.write_header_to_file()
    writer.write_kv_data_to_file()
    writer.write_tensors_to_file()
    writer.close()

    print(f" Done → {output_path}")
    print(f" Run GGUF quantization: the gguf library handles Q4_K_M inline")

    import struct, os
    fsize = os.path.getsize(output_path)
    print(f" Raw size: {fsize/1024**2:.1f}MB")

    return True


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint", default=None)
    parser.add_argument("--output", default="outputs/tiny-sft/tiny.gguf")
    parser.add_argument("--quantize", default="q4_k_m",
                        choices=["q4_0", "q4_1", "q5_0", "q5_1", "q8_0", "q4_k_m", "q5_k_m", "q6_k", "q8_k_m"])
    args = parser.parse_args()

    if args.checkpoint is None:
        args.checkpoint = "outputs/tiny-sft/final/model.pt"
        if not os.path.exists(args.checkpoint):
            print(f"No checkpoint found at {args.checkpoint}")
            print("Train first: bash scripts/train_tiny.sh")
            sys.exit(1)

    export_gguf(args.checkpoint, args.output)
    print(f"GGUF file: {args.output}")


if __name__ == "__main__":
    main()