"""Run the S/R v2 training matrix or one locked validation comparison.

This is a cross-platform orchestrator around ``run_backtest_snapshot.py``. It
contains no backtest logic; every arm still runs through the production-parity
Python harness.
"""

from __future__ import annotations

import argparse
import subprocess
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
RUNNER = ROOT / "scripts" / "run_backtest_snapshot.py"
COMPARE = ROOT / "scripts" / "compare_sr_variants.py"
TRAINING_ARMS = (
    "production_control",
    "rewrite_legacy_primary",
    "soft_zones_legacy_primary",
    "confirmed_rounds_legacy_primary",
    "gate_v2_legacy_primary",
)
LOCKABLE_ARMS = TRAINING_ARMS[1:]


def _add_common(parser: argparse.ArgumentParser) -> None:
    parser.add_argument(
        "--snapshot",
        default="backtest_snapshots/prod.sqlite",
        help="Local SQLite snapshot path.",
    )
    parser.add_argument("--workers", type=int, default=7)


def _args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    commands = parser.add_subparsers(dest="command", required=True)
    train = commands.add_parser("train", help="Run all arms before 2024-07-01.")
    _add_common(train)
    validate = commands.add_parser(
        "validate",
        help="Run production control and one locked arm from 2024-07-01.",
    )
    _add_common(validate)
    validate.add_argument("--locked-arm", required=True, choices=LOCKABLE_ARMS)
    return parser.parse_args()


def _run_arm(
    arm: str,
    snapshot: str,
    workers: int,
    *,
    entry_flag: str,
    entry_date: str,
    output: Path,
) -> None:
    print(f"Running S/R arm: {arm}", flush=True)
    subprocess.run(
        [
            sys.executable,
            str(RUNNER),
            snapshot,
            "--workers", str(workers),
            "--allow-spawn",
            "--sr-variant", arm,
            entry_flag, entry_date,
            "--sr-audit",
            "--out", str(output),
        ],
        cwd=ROOT,
        check=True,
    )


def _train(args: argparse.Namespace) -> None:
    for arm in TRAINING_ARMS:
        _run_arm(
            arm,
            args.snapshot,
            args.workers,
            entry_flag="--entry-end",
            entry_date="2024-06-30",
            output=ROOT / "reports" / f"backtest-sr-v2-train-{arm}.json",
        )
    print("Training matrix complete. Lock one arm before running validation.")


def _validate(args: argparse.Namespace) -> None:
    reports: dict[str, Path] = {}
    for arm in ("production_control", args.locked_arm):
        output = ROOT / "reports" / f"backtest-sr-v2-validation-{arm}.json"
        reports[arm] = output
        _run_arm(
            arm,
            args.snapshot,
            args.workers,
            entry_flag="--entry-start",
            entry_date="2024-07-01",
            output=output,
        )
    subprocess.run(
        [
            sys.executable,
            str(COMPARE),
            str(reports["production_control"]),
            str(reports[args.locked_arm]),
            "--out-csv", str(ROOT / "reports" / "sr-v2-validation-cohorts.csv"),
            "--out-json", str(ROOT / "reports" / "sr-v2-validation-comparison.json"),
        ],
        cwd=ROOT,
        check=True,
    )


def main() -> None:
    args = _args()
    if args.command == "train":
        _train(args)
    else:
        _validate(args)


if __name__ == "__main__":
    main()