"""Merge a numbered training checkpoint into a servable bf16 model, like the end of `ccslm.train`.

  python -m ccslm.export --run out/runs/A --list                  # checkpoints with step, loss, eval loss
  python -m ccslm.export --run out/runs/A --step 100              # -> out/runs/A/merged-step-00100
  python -m ccslm.export --run out/runs/A --step latest --out /somewhere/model

The output has the base's tokenizer files byte for byte and its own config.json, so the pipeline's
eval, speed and package stages take it as they take `merged` (`--step` on the pipeline).
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any

from ccslm.common import base_dir as pinned
from ccslm.common import log, read_json


def snapshots(run: Path) -> list[Path]:
    root = run / "checkpoints"
    return sorted(p for p in root.glob("step-*") if p.is_dir() and not p.name.endswith(".tmp"))


def find(run: Path, step: str) -> Path:
    found = snapshots(run)
    if not found:
        raise SystemExit(f"no checkpoints under {run / 'checkpoints'}")
    if step == "latest":
        return found[-1]
    where = run / "checkpoints" / f"step-{int(step):05d}"
    if not where.is_dir():
        raise SystemExit(f"no checkpoint {where.name}; have {', '.join(p.name for p in found)}")
    return where


def merged_dir(run: Path, step: int) -> Path:
    return run / f"merged-step-{step:05d}"


def export_snapshot(run: Path, snapshot: Path, out: Path, device: str = "cpu") -> Path:
    import torch
    from transformers import AutoModelForCausalLM

    from ccslm.train import export

    config = read_json(run / "config.json", {})
    state = read_json(snapshot / "state.json", {})
    method = state.get("method") or config.get("method") or "lora"
    start = Path(config["start_from"]) if config.get("start_from") else pinned()
    if method == "lora":
        from peft import PeftModel

        model = AutoModelForCausalLM.from_pretrained(str(start), dtype=torch.bfloat16)
        model = PeftModel.from_pretrained(model, str(snapshot))
    else:
        model = AutoModelForCausalLM.from_pretrained(str(snapshot), dtype=torch.bfloat16)
    model.to(device)
    export(model, method, start, pinned(), out)
    log(f"step {state.get('step')} ({method}) merged into {out}")
    return out


def main() -> None:
    parser = argparse.ArgumentParser(description="merge a numbered checkpoint into a servable model")
    parser.add_argument("--run", required=True, help="the training run folder, e.g. out/runs/A")
    parser.add_argument("--step", default="latest", help="step number or latest")
    parser.add_argument("--out", default="", help="default: <run>/merged-step-NNNNN")
    parser.add_argument("--device", default="cpu")
    parser.add_argument("--list", action="store_true")
    args = parser.parse_args()
    run = Path(args.run)
    if args.list:
        for snap in snapshots(run):
            state: dict[str, Any] = read_json(snap / "state.json", {})
            merged = merged_dir(run, int(state.get("step", 0)))
            print(json.dumps({"checkpoint": snap.name, **{k: state.get(k) for k in ("step", "loss", "eval_loss", "seen_tokens")}, "merged": merged.is_dir()}))
        return
    snapshot = find(run, args.step)
    step = int(read_json(snapshot / "state.json", {}).get("step") or snapshot.name.split("-")[-1])
    export_snapshot(run, snapshot, Path(args.out) if args.out else merged_dir(run, step), args.device)


if __name__ == "__main__":
    main()