"""Extract model weights from a Lightning checkpoint into a standalone model.pt."""
from __future__ import annotations

import argparse
import sys
from pathlib import Path

# Ensure project root is on sys.path so `src` package is importable
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

import torch
from src.model.config import BoccaccioConfig
from src.model.transformer import BoccaccioForCausalLM


def main() -> None:
    parser = argparse.ArgumentParser(description="Extract model.pt from Lightning checkpoint")
    parser.add_argument("--ckpt", required=True, help="Path to Lightning .ckpt file")
    parser.add_argument("--output", required=True, help="Path to save model.pt")
    parser.add_argument("--variant", default="model", help="Model variant in configs/model.yaml")
    args = parser.parse_args()

    config = BoccaccioConfig.from_yaml("configs/model.yaml", variant=args.variant)

    ckpt = torch.load(args.ckpt, map_location="cpu", weights_only=False)
    state_dict = ckpt.get("state_dict", ckpt)

    # Remap prefixes from Lightning checkpoint to BoccaccioForCausalLM
    # Possible prefixes: "model._orig_mod.X" (torch.compile) or "model.X" (standard)
    cleaned = {}
    for key, value in state_dict.items():
        new_key = key
        if new_key.startswith("model._orig_mod."):
            new_key = new_key[len("model._orig_mod."):]
        elif new_key.startswith("model."):
            new_key = new_key[len("model."):]
        elif new_key.startswith("_orig_mod."):
            new_key = new_key[len("_orig_mod."):]
        cleaned[new_key] = value

    model = BoccaccioForCausalLM(config)
    missing, unexpected = model.load_state_dict(cleaned, strict=False)
    if missing:
        print(f"  WARNING: {len(missing)} missing keys: {missing[:3]}")
    if unexpected:
        print(f"  WARNING: {len(unexpected)} unexpected keys: {unexpected[:3]}")

    torch.save(model.state_dict(), args.output)
    print(f"model.pt saved to {args.output} (step {ckpt.get('global_step', '?')})")


if __name__ == "__main__":
    main()