"""Command line interface.

Examples
--------
    python -m kronllm.cli train --exp E2
    python -m kronllm.cli train --exp E2 --steps 300 --run-dir runs/smoke
    python -m kronllm.cli generate --run runs/E2 --prompt "Once upon a time"
    python -m kronllm.cli eval --run runs/E2
    python -m kronllm.cli analyze --run runs/E2
"""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path
from typing import Any, Dict, List, Optional

import torch


# ---------------------------------------------------------------------------
def _load_model(run_dir: str | Path, checkpoint: str = "best.pt"):
    from kronllm.config import load_config
    from kronllm.model import KronTransformer
    from kronllm.quantization import dequantize_state_dict

    run_dir = Path(run_dir)
    path = run_dir / checkpoint
    if not path.exists():
        path = run_dir / "last.pt"
    if not path.exists():
        raise FileNotFoundError(f"no checkpoint in {run_dir}")
    ck = torch.load(path, map_location="cpu", weights_only=False)
    cfg = ck.get("cfg") or load_config(run_dir / "config.yaml")
    if ck.get("quantized"):
        state = dequantize_state_dict(ck["state"], int(ck["bits"]))
    else:
        state = ck["model"]
    model = KronTransformer(cfg)
    model.load_state_dict(state)
    model.eval()
    return model, cfg, ck


def _device(name: Optional[str] = None) -> torch.device:
    if name:
        return torch.device(name)
    return torch.device("cuda" if torch.cuda.is_available() else "cpu")


# ---------------------------------------------------------------------------
def cmd_train(args: argparse.Namespace) -> None:
    from kronllm.config import experiment_config, load_config
    from kronllm.train import train

    if args.exp:
        cfg = experiment_config(args.exp, args.config, args.override)
    else:
        cfg = load_config(args.config, args.override)

    if args.steps is not None:
        cfg["training"]["max_steps"] = args.steps
    if args.epochs is not None:
        cfg["training"]["epochs"] = args.epochs
    if args.batch_size is not None:
        cfg["training"]["batch_size"] = args.batch_size
    if args.lr is not None:
        cfg["training"]["lr"] = args.lr
    if args.device is not None:
        cfg["training"]["device"] = args.device
    if args.no_amp:
        cfg["training"]["amp"] = "none"
    if args.data_dir is not None:
        cfg["data"]["data_dir"] = args.data_dir
    if args.wandb:
        cfg["logging"]["wandb"] = True

    run_dir = Path(args.run_dir or (Path(cfg["training"]["out_dir"])
                                    / (args.exp or "custom")))
    results = train(cfg, run_dir, resume=args.resume)
    print(json.dumps({k: results[k] for k in
                      ("experiment", "n_params", "val_ppl", "best_val_ppl",
                       "val_bpb", "time_per_epoch_sec", "peak_vram_gb")}, indent=2))


def cmd_generate(args: argparse.Namespace) -> None:
    from kronllm.generate import generate_texts
    from kronllm.tokenizer import load_tokenizer

    model, cfg, _ = _load_model(args.run, args.checkpoint)
    device = _device(args.device)
    model.to(device)
    texts = generate_texts(
        model, [args.prompt], max_new_tokens=args.max_new,
        temperature=args.temperature, top_k=args.top_k,
        n_samples=args.samples, seed=args.seed, device=device,
        tokenizer=load_tokenizer(cfg))
    for text in texts:
        print("-" * 70)
        print(text)


def cmd_eval(args: argparse.Namespace) -> None:
    from kronllm.data import open_dataset
    from kronllm.diagnostics import (comparison_report, sigma_theta_report)
    from kronllm.train import diagnostics, evaluate

    model, cfg, ck = _load_model(args.run, args.checkpoint)
    device = _device(args.device)
    model.to(device)
    ds = open_dataset(cfg, "val")
    out: Dict[str, Any] = {"checkpoint": str(args.run), "step": ck.get("step"),
                           **comparison_report(model, int(cfg["data"]["seq_len"]))}
    out.update(evaluate(model, ds, 32, device, args.batches, cfg["training"]["amp"]))
    dg = diagnostics(model, ds, 32, device, spectral=True,
                     amp=cfg["training"]["amp"])
    out.update({k: v for k, v in dg.items() if isinstance(v, (int, float))})
    print(json.dumps(out, indent=2, default=float))


def cmd_analyze(args: argparse.Namespace) -> None:
    scripts_dir = Path(__file__).resolve().parents[1] / "scripts"
    if str(scripts_dir) not in sys.path:
        sys.path.insert(0, str(scripts_dir))
    import analyze as analyze_mod

    argv = ["--run", str(args.run), "--checkpoint", args.checkpoint,
            "--samples", str(args.samples)]
    if args.skip_generation:
        argv.append("--skip-generation")
    old = sys.argv
    try:
        sys.argv = ["analyze.py"] + argv
        analyze_mod.main()
    finally:
        sys.argv = old


# ---------------------------------------------------------------------------
def cmd_quantize(args: argparse.Namespace) -> None:
    """Quantize a checkpoint to int8/int4 and report the quality delta."""
    from kronllm.data import open_dataset
    from kronllm.quantization import quantize_checkpoint
    from kronllm.tokenizer import bits_per_byte, load_tokenizer
    from kronllm.train import evaluate

    run_dir = Path(args.run)
    model, cfg, _ = _load_model(run_dir, args.checkpoint)
    device = _device(args.device)
    model.to(device)
    dataset = open_dataset(cfg, "val")
    bpt = load_tokenizer(cfg).bytes_per_token
    amp = cfg["training"]["amp"]

    base = evaluate(model, dataset, args.batch_size, device, args.batches, amp)
    base_bpb = bits_per_byte(base["val_loss"], bpt)
    print(f"[fp32] loss {base['val_loss']:.4f} | bpb {base_bpb:.4f}")

    src = run_dir / args.checkpoint
    if not src.exists():
        src = run_dir / "last.pt"
    reports = []
    for bits in args.bits:
        destination = run_dir / f"model_int{bits}.pt"
        report = quantize_checkpoint(src, destination, bits=bits)
        quantized_model, _, _ = _load_model(run_dir, destination.name)
        quantized_model.to(device)
        result = evaluate(quantized_model, dataset, args.batch_size, device,
                          args.batches, amp)
        result_bpb = bits_per_byte(result["val_loss"], bpt)
        report.update({"val_loss": result["val_loss"], "val_bpb": result_bpb,
                       "delta_bpb": result_bpb - base_bpb})
        print(f"[int{bits}] loss {result['val_loss']:.4f} | bpb {result_bpb:.4f} "
              f"| delta {report['delta_bpb']:+.4f} | "
              f"{report['size_after_mb']:.2f} MB ({report['compression']}x)")
        reports.append(report)
        del quantized_model
        if device.type == "cuda":
            torch.cuda.empty_cache()

    payload = {"checkpoint": str(src), "fp32_loss": base["val_loss"],
               "fp32_bpb": base_bpb, "reports": reports}
    out = run_dir / "quantization.json"
    out.write_text(json.dumps(payload, indent=2), encoding="utf-8")
    print(f"wrote {out}")


# ---------------------------------------------------------------------------
def cmd_budget(args: argparse.Namespace) -> None:
    """Print the parameter budget of the E-series (all models with ``--all``)."""
    from .config import EXPERIMENTS, experiment_config
    from .model import KronTransformer

    limit = int(args.limit)
    if args.all:
        ids = list(EXPERIMENTS)
    else:
        ids = [exp_id for exp_id in EXPERIMENTS if exp_id.startswith("E")]
    print(f"{'E':>4} {'total':>10} {'matrix':>10} {'<=1M':>6}  description")
    worst = 0
    for exp_id in ids:
        model = KronTransformer(experiment_config(exp_id))
        worst = max(worst, model.n_params)
        ok = "ok" if model.n_params <= limit else "FAIL"
        print(f"{exp_id:>4} {model.n_params:>10,} {model.n_params_matrix:>10,} "
              f"{ok:>6}  {EXPERIMENTS[exp_id]['desc']}")
    print(f"\nmax {worst:,} / {limit:,} ({worst / limit:.1%})")
    if worst > limit:
        raise SystemExit(1)


def build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(prog="kronllm")
    sub = p.add_subparsers(dest="command", required=True)

    tr = sub.add_parser("train", help="train one model / experiment")
    tr.add_argument("--exp", default=None, help="E0..E12")
    tr.add_argument("--config", default=None)
    tr.add_argument("--run-dir", default=None)
    tr.add_argument("--resume", default=None)
    tr.add_argument("--steps", type=int, default=None)
    tr.add_argument("--epochs", type=int, default=None)
    tr.add_argument("--batch-size", type=int, default=None)
    tr.add_argument("--lr", type=float, default=None)
    tr.add_argument("--device", default=None)
    tr.add_argument("--data-dir", default=None)
    tr.add_argument("--no-amp", action="store_true")
    tr.add_argument("--wandb", action="store_true")
    tr.add_argument("--override", action="append", default=[])
    tr.set_defaults(func=cmd_train)

    ge = sub.add_parser("generate", help="sample from a checkpoint")
    ge.add_argument("--run", required=True)
    ge.add_argument("--checkpoint", default="best.pt")
    ge.add_argument("--prompt", default="Once upon a time")
    ge.add_argument("--max-new", type=int, default=200)
    ge.add_argument("--temperature", type=float, default=1.0)
    ge.add_argument("--top-k", type=int, default=8)
    ge.add_argument("--samples", type=int, default=1)
    ge.add_argument("--seed", type=int, default=None)
    ge.add_argument("--device", default=None)
    ge.set_defaults(func=cmd_generate)

    ev = sub.add_parser("eval", help="evaluate a checkpoint")
    ev.add_argument("--run", required=True)
    ev.add_argument("--checkpoint", default="best.pt")
    ev.add_argument("--batches", type=int, default=64)
    ev.add_argument("--device", default=None)
    ev.set_defaults(func=cmd_eval)

    an = sub.add_parser("analyze", help="spectra / heatmaps / samples")
    an.add_argument("--run", required=True)
    an.add_argument("--checkpoint", default="best.pt")
    an.add_argument("--samples", type=int, default=10)
    an.add_argument("--skip-generation", action="store_true")
    an.set_defaults(func=cmd_analyze)

    qt = sub.add_parser("quantize", help="int8/int4 checkpoint quantization")
    qt.add_argument("--run", required=True)
    qt.add_argument("--checkpoint", default="best.pt")
    qt.add_argument("--bits", type=int, nargs="+", default=[8, 4])
    qt.add_argument("--batches", type=int, default=32)
    qt.add_argument("--batch-size", type=int, default=32)
    qt.add_argument("--device", default=None)
    qt.set_defaults(func=cmd_quantize)

    bd = sub.add_parser("budget", help="print the parameter budget of E0..E13")
    bd.add_argument("--limit", type=int, default=1_000_000)
    bd.add_argument("--all", action="store_true",
                    help="also print Q1-Q3, D1 and S1 (deliberately above 1M)")
    bd.set_defaults(func=cmd_budget)
    return p


def main(argv: Optional[List[str]] = None) -> None:
    args = build_parser().parse_args(argv)
    args.func(args)


if __name__ == "__main__":
    main()