from __future__ import annotations

import argparse
from dataclasses import asdict

from .controller import TrainerConfig, fit_controller
from .benchmarks import run_benchmark_suite
from .gui import launch_runtime_gui
from .smoke import run_tinygrad_gate_demo


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

    demo = sub.add_parser("demo", help="run the tinygrad gate-controller smoke demo")
    demo.add_argument("--steps", type=int, default=80)
    demo.add_argument("--batch-size", type=int, default=64)
    demo.add_argument("--seed", type=int, default=0)
    demo.add_argument("--target-accuracy", type=float, default=0.99)

    fit = sub.add_parser("fit", help="fit an ntkmirror controller on a frozen causal LM")
    fit.add_argument("--model", default="OpenPeerAI/OpenPeerLLM")
    fit.add_argument("--fallback-model", default="sshleifer/tiny-gpt2")
    fit.add_argument("--train-jsonl")
    fit.add_argument("--out", default="runs/openpeer_controller.pt")
    fit.add_argument("--gates", type=int, default=512)
    fit.add_argument("--steps", type=int, default=40)
    fit.add_argument("--demo-mode", action="store_true")

    bench = sub.add_parser("bench", help="run local gate-controller benchmarks and generate charts")
    bench.add_argument("--steps", type=int, nargs="+", default=[20, 40, 80, 120])
    bench.add_argument("--batch-size", type=int, default=64)
    bench.add_argument("--seed", type=int, default=0)
    bench.add_argument("--output-dir", default="artifacts/benchmarks")
    bench.add_argument("--target-accuracy", type=float, default=0.99)

    gui = sub.add_parser("gui", help="launch the runtime GUI")

    return parser


def main(argv: list[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)

    if args.command == "demo":
        result = run_tinygrad_gate_demo(steps=args.steps, batch_size=args.batch_size, seed=args.seed, target_accuracy=args.target_accuracy)
        print("tinygrad gate demo")
        print(f"initial_accuracy={result.initial_accuracy:.3f}")
        print(f"final_accuracy={result.final_accuracy:.3f}")
        print(f"final_loss={result.final_loss:.4f}")
        print(f"target_accuracy={result.target_accuracy:.3f}")
        print(f"reached_target={result.reached_target}")
        print(f"trained_steps={result.trained_steps}")
        print(f"learned_gate_sample={result.learned_gate_sample}")
        return 0

    if args.command == "bench":
        result = run_benchmark_suite(
            step_counts=args.steps,
            batch_size=args.batch_size,
            seed=args.seed,
            output_dir=args.output_dir,
            target_accuracy=args.target_accuracy,
        )
        print(f"benchmarks_csv={result.csv_path}")
        print(f"charts={result.chart_paths}")
        return 0

    if args.command == "gui":
        return launch_runtime_gui()

    config = TrainerConfig(
        model_name=args.model,
        fallback_model_name=args.fallback_model,
        train_jsonl=args.train_jsonl,
        out_path=args.out,
        gates=args.gates,
        steps=args.steps,
        demo_mode=args.demo_mode,
    )
    output_path = fit_controller(config)
    print(f"saved_controller={output_path}")
    print(f"config={asdict(config)}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())