"""Unified CLI: `python -m ai_mouse {train,eval,balabit-adapter}`

Subcommands dispatch to the underlying modules. This is the recommended
top-level entry; you can also call `python -m ai_mouse.eval` etc. directly.
"""
from __future__ import annotations

import argparse
import logging
import sys
from pathlib import Path


def _train_main(args: argparse.Namespace) -> int:
    from tools.trainer import train

    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
    train(
        data_path=Path(args.data),
        output_dir=Path(args.output),
        epochs=args.epochs,
        batch_size=args.batch_size,
        lr=args.lr,
        seq_len=args.seq_len,
        resume_from=Path(args.resume_from) if args.resume_from else None,
    )
    return 0


def _eval_main(args: argparse.Namespace) -> int:
    from tools.eval.__main__ import main as eval_main
    # Reconstruct argv for the sub-CLI
    argv = [
        "--model-dir", args.model_dir,
        "--reference", str(args.reference),
        "--n-samples", str(args.n_samples),
        "--n-reference", str(args.n_reference),
        "--output", str(args.output),
        "--tag", args.tag,
        "--seed", str(args.seed),
    ]
    return eval_main(argv)


def _balabit_main(args: argparse.Namespace) -> int:
    from tools.data_adapters.balabit import main as bal_main
    argv = [
        "--input", str(args.input),
        "--output", str(args.output),
        "--window-ms", str(args.window_ms),
        "--min-dist", str(args.min_dist),
        "--min-events", str(args.min_events),
        "--max-span-ms", str(args.max_span_ms),
        "--max-gap-ms", str(args.max_gap_ms),
    ]
    if args.overwrite:
        argv.append("--overwrite")
    return bal_main(argv)


def main() -> int:
    p = argparse.ArgumentParser(prog="ai_mouse", description="AI Mouse trajectory toolkit")
    sub = p.add_subparsers(dest="cmd", required=True)

    # train
    pt = sub.add_parser("train", help="Train (or fine-tune) the Flow Matching model")
    pt.add_argument("--data", required=True, help="Path to traces.jsonl")
    pt.add_argument("--output", required=True, help="Output checkpoint dir")
    pt.add_argument("--epochs", type=int, default=200)
    pt.add_argument("--batch-size", type=int, default=64)
    pt.add_argument("--lr", type=float, default=3e-4)
    pt.add_argument("--seq-len", type=int, default=64)
    pt.add_argument("--resume-from", default=None, help="Checkpoint dir to resume from (for fine-tune)")
    pt.set_defaults(func=_train_main)

    # eval
    pe = sub.add_parser("eval", help="Generate evaluation report")
    pe.add_argument("--model-dir", required=True)
    pe.add_argument("--reference", type=Path, required=True)
    pe.add_argument("--n-samples", type=int, default=200)
    pe.add_argument("--n-reference", type=int, default=1000)
    pe.add_argument("--output", type=Path, required=True)
    pe.add_argument("--tag", default="eval")
    pe.add_argument("--seed", type=int, default=0)
    pe.set_defaults(func=_eval_main)

    # balabit-adapter
    pb = sub.add_parser("balabit-adapter", help="Convert Balabit dataset to traces.jsonl")
    pb.add_argument("--input", type=Path, required=True)
    pb.add_argument("--output", type=Path, default=Path("data/pretrain_traces.jsonl"))
    pb.add_argument("--window-ms", type=int, default=1200)
    pb.add_argument("--min-dist", type=int, default=50)
    pb.add_argument("--min-events", type=int, default=5)
    pb.add_argument("--max-span-ms", type=int, default=5000)
    pb.add_argument("--max-gap-ms", type=int, default=200)
    pb.add_argument("--overwrite", action="store_true")
    pb.set_defaults(func=_balabit_main)

    args = p.parse_args()
    return args.func(args)


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