import argparse
import sys
from typing import cast

from kernbench.benches.registry import list_all, resolve
from kernbench.cli.report import format_report
from kernbench.common.types import SimEngine
from kernbench.probes.probe import cmd_probe
from kernbench.runtime_api.bench_runner import run_bench
from kernbench.runtime_api.types import DeviceSelector, resolve_device
from kernbench.sim_engine.engine import GraphEngine
from kernbench.topology.builder import resolve_topology
from kernbench.topology.types import TopologyGraph


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

    runp = sub.add_parser("run", help="Run a benchmark")
    runp.add_argument("--topology", required=True)
    runp.add_argument(
        "--bench", required=True,
        help="Bench name (kebab-case) or numeric index from 'kernbench list'",
    )
    runp.add_argument(
        "--device", default=None, help="Target device: 'all' or 'sip:<N>' (default: all)"
    )
    runp.add_argument(
        "--verify-data", action="store_true", default=False,
        help="Enable Phase 2 data verification (ADR-0020)",
    )
    runp.set_defaults(_handler=cmd_run)

    listp = sub.add_parser("list", help="List registered benches")
    listp.set_defaults(_handler=cmd_list)

    probep = sub.add_parser("probe", help="Probe latency and BW for predefined traffic patterns")
    probep.add_argument("--topology", required=True)
    probep.add_argument("--case", default="all", help="Case name or 'all' (default: all)")
    probep.set_defaults(_handler=cmd_probe)

    webp = sub.add_parser("web", help="Launch topology viewer in browser")
    webp.add_argument("--port", type=int, default=8765, help="HTTP port (default: 8765)")
    webp.add_argument("--no-open", action="store_true", help="Don't auto-open browser")
    webp.set_defaults(_handler=cmd_web)

    return p


def engine_factory(
    topology: object, device: DeviceSelector, *, enable_data: bool = False,
) -> SimEngine:
    topo_obj = cast(TopologyGraph, getattr(topology, "topology_obj", topology))
    return GraphEngine(topo_obj, enable_data=enable_data)


def cmd_web(args) -> int:
    from kernbench.web.server import serve
    serve(port=args.port, open_browser=not args.no_open)
    return 0


def cmd_list(args) -> int:
    specs = list_all()
    print(f"{'#':>3}  {'NAME':<22}  DESCRIPTION")
    print("-" * 80)
    for s in specs:
        print(f"{s.index:>3}  {s.name:<22}  {s.description}")
    return 0


def cmd_run(args) -> int:
    print("> Running benchmark with:", args)

    topo = resolve_topology(args.topology)
    spec_entry = resolve(args.bench)
    device = resolve_device(args.device)
    verify_data = getattr(args, "verify_data", False)

    def _factory(topology, device):
        return engine_factory(topology, device, enable_data=verify_data)

    result = run_bench(
        topology=topo, bench_fn=spec_entry.run, device=device, engine_factory=_factory,
    )

    topo_obj = getattr(topo, "topology_obj", topo)
    spec = getattr(topo_obj, "spec", None)
    if result.traces:
        print(format_report(result.traces, title=spec_entry.name, spec=spec))
    print(result.summary_text())

    # Phase 2 diagnostic summary (ADR-0020). The actual Phase 2 replay
    # already runs inside engine.wait() → _flush_data_phase(). We only
    # print the summary here; no redundant re-execution.
    if verify_data and result.engine is not None:
        op_log = result.engine.op_log
        if op_log:
            n_gemm = sum(1 for r in op_log if r.op_kind == "gemm")
            n_math = sum(1 for r in op_log if r.op_kind == "math")
            print(f"[data] Phase 2 complete: {len(op_log)} ops ({n_gemm} gemm, {n_math} math)")
        else:
            print("[data] No op_log recorded — skipping Phase 2")

    return 0 if result.completion.ok else 1


def main(argv=None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)
    return int(args._handler(args))


if __name__ == "__main__":
    sys.exit(main())
