#!/usr/bin/env python3
"""SPIRA training CLI.

Usage:
    python3 -m spira.train --data data/train.bin \
        --out artifacts/spira_model.npz \
        [--snn-tokens 600000] [--am-tokens 1200000] [--epochs 2] [--quick]

Writes:
    <out>                  : model checkpoint (npz)
    <out>.train_log.jsonl  : tick-level progress log
    <out>.history.json     : final training history
"""
from __future__ import annotations

import argparse
import json
import sys
import time

from .model import SPIRA, SPIRAConfig


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--data", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--snn-tokens", type=int, default=600_000)
    ap.add_argument("--am-tokens", type=int, default=1_200_000)
    ap.add_argument("--epochs", type=int, default=2)
    ap.add_argument("--vocab", type=int, default=4096)
    ap.add_argument("--n-res", type=int, default=512)
    ap.add_argument("--target-avg-len", type=float, default=4.5)
    ap.add_argument("--quick", action="store_true", help="tiny smoke-test run")
    args = ap.parse_args()

    cfg = SPIRAConfig(
        vocab_size=args.vocab,
        n_res=args.n_res,
        snn_tokens=args.snn_tokens,
        am_tokens=args.am_tokens,
        epochs=args.epochs,
        target_avg_len=args.target_avg_len,
    )
    if args.quick:
        cfg.snn_tokens, cfg.am_tokens, cfg.epochs = 20_000, 20_000, 1

    data = open(args.data, "rb").read()
    t0 = time.perf_counter()

    log_path = args.out + ".train_log.jsonl"
    log_f = open(log_path, "w")

    def log(msg: str) -> None:
        print(msg, flush=True)
        log_f.write(json.dumps({"t": round(time.perf_counter() - t0, 1), "msg": msg}) + "\n")
        log_f.flush()

    def tick_log(d: dict) -> None:
        log_f.write(json.dumps({"t": round(time.perf_counter() - t0, 1), **d}) + "\n")
        log_f.flush()

    model = SPIRA.build(data, cfg, log=log)
    model.prepare()
    # free the patch list — the registry carries everything we need
    model._patches_train = None

    log("[spira] training starts")
    history = model.train(
        log_fn=lambda msg: tick_log({"msg": msg}) or print(msg, flush=True),
        ckpt_path=args.out,
    )
    history["build_meta"] = model.meta
    history["config"] = {k: v for k, v in vars(cfg).items()}
    history["wall_clock_s"] = time.perf_counter() - t0
    with open(args.out + ".history.json", "w") as f:
        json.dump(history, f, indent=2)
    model.save(args.out)
    log(f"[spira] DONE in {history['wall_clock_s']:.0f}s -> {args.out}")
    log_f.close()


if __name__ == "__main__":
    main()