#!/usr/bin/env python3
"""Automated free model probing, admission verification, and catalog synchronization for OMP."""

from __future__ import annotations

import argparse
import json
import os
import re
import subprocess
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

try:
    import yaml
except ModuleNotFoundError:  # pragma: no cover - fresh-install convenience
    sys.exit("PyYAML is required: pip install pyyaml")

PKG_ROOT = Path(__file__).resolve().parent.parent
CATALOG_DIR = PKG_ROOT / "templates" / "catalog"
CATALOG_PATH = CATALOG_DIR if CATALOG_DIR.is_dir() else (PKG_ROOT / "catalog.yml")
TEST_CONFIG_PATH = PKG_ROOT / "tests" / "test_router.py"

@dataclass
class ProbeResult:
    selector: str
    runs: int
    successes: int
    mean_ttft_s: float
    mean_tps: float
    status: str  # "VERIFIED", "SLOW", "QUARANTINE"
    error_summary: str = ""


def run_command(args: list[str], timeout: int = 120) -> tuple[int, str, str]:
    """Execute a subprocess command with capture."""
    try:
        proc = subprocess.run(
            args,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            timeout=timeout,
        )
        return proc.returncode, proc.stdout, proc.stderr
    except subprocess.TimeoutExpired:
        return 124, "", f"Timed out after {timeout}s"
    except Exception as e:
        return 1, "", str(e)


def discover_free_models(zen_only: bool = False, openrouter_only: bool = False) -> list[str]:
    """Discover available free models from opencode and omp models list."""
    discovered: set[str] = set()

    # 1. Query omp models --json
    code, out, _ = run_command(["omp", "models", "--json"], timeout=20)
    if code == 0 and out.strip():
        try:
            data = json.loads(out)
            items = data if isinstance(data, list) else (data.get("models") or data.get("data") or [])
            for item in items:
                if not isinstance(item, dict):
                    continue
                mid = item.get("id") or ""
                prov = item.get("provider") or ""
                full_id = mid if "/" in mid and mid.startswith(prov + "/") else f"{prov}/{mid}"

                if not openrouter_only and prov == "opencode-zen":
                    # All opencode-zen free models
                    if "free" in mid or mid in {
                        "big-pickle",
                        "x-preview-f-free",
                        "mimo-v2.5-free",
                        "hy3-free",
                        "laguna-s-2.1-free",
                        "muse-spark-1.2-contributor-free",
                        "nemotron-3-ultra-free",
                        "nemotron-3.5-lightning-free",
                        "deepseek-v4-flash-free",
                    }:
                        discovered.add(full_id)

                if not zen_only and (prov in {"openrouter", "openrouter-1", "openrouter-2"} or mid.startswith("openrouter/")):
                    if ":free" in mid or mid.endswith("/free") or "free" in mid:
                        # Use openrouter-1/2 if available, else openrouter
                        discovered.add(full_id)
        except Exception as e:
            sys.stderr.write(f"Warning: failed to parse omp models JSON discovery: {e}\n")
    # 2. Query opencode models if opencode is installed
    if not openrouter_only:
        code, out, _ = run_command(["opencode", "models"], timeout=10)
        if code == 0 and out.strip():
            for line in out.splitlines():
                line = line.strip()
                if line.startswith("opencode/") and "free" in line:
                    model_id = line.replace("opencode/", "")
                    discovered.add(f"opencode-zen/{model_id}")

    # Fallback default known pool
    if not discovered:
        if not openrouter_only:
            discovered.update(
                [
                    "opencode-zen/nemotron-3.5-lightning-free",
                    "opencode-zen/muse-spark-1.2-contributor-free",
                    "opencode-zen/hy3-free",
                    "opencode-zen/laguna-s-2.1-free",
                    "opencode-zen/x-preview-f-free",
                ]
            )
        if not zen_only:
            discovered.update(
                [
                    "openrouter-1/openrouter/free",
                ]
            )

    return sorted(discovered)


def probe_single_model(selector: str, runs: int = 3, max_tokens: int = 64) -> ProbeResult:
    """Probe a single model using omp bench with JSON output."""
    cmd = [
        "omp",
        "bench",
        selector,
        f"--runs={runs}",
        "--par=1",
        f"--max-tokens={max_tokens}",
        "--json",
    ]
    code, out, err = run_command(cmd, timeout=runs * 45 + 15)

    if code != 0 and not out.strip():
        # Complete failure before bench produced JSON
        err_msg = (err or out).strip()[:100].replace("\n", " ")
        return ProbeResult(
            selector=selector,
            runs=runs,
            successes=0,
            mean_ttft_s=0.0,
            mean_tps=0.0,
            status="QUARANTINE",
            error_summary=err_msg or f"Exit {code}",
        )

    try:
        data = json.loads(out)
        models = data.get("models", [])
        if not models:
            return ProbeResult(
                selector=selector,
                runs=runs,
                successes=0,
                mean_ttft_s=0.0,
                mean_tps=0.0,
                status="QUARANTINE",
                error_summary="No models in bench output",
            )

        model_data = models[0]
        results = model_data.get("results", [])
        stats = model_data.get("stats") or {}

        successes = sum(1 for r in results if r.get("ok"))
        errors = [r.get("error", "").strip() for r in results if not r.get("ok") and r.get("error")]

        mean_ttft_s = 0.0
        mean_tps = 0.0

        if stats:
            ttft_ms = stats.get("ttftMs", {}).get("mean", 0)
            mean_ttft_s = round(ttft_ms / 1000.0, 2)
            mean_tps = round(stats.get("tokensPerSecond", {}).get("mean", 0), 1)

        # Strict admission rule: requires at least 3 serial runs AND 100% success
        if successes == runs and runs >= 3:
            if mean_tps >= 15.0 and mean_ttft_s <= 5.0:
                status = "VERIFIED"
            else:
                status = "SLOW"
            error_summary = ""
        elif successes == runs and runs < 3:
            status = "PROBED"
            error_summary = f"{runs} runs (need >= 3 to verify)"
        else:
            status = "QUARANTINE"
            first_err = errors[0] if errors else "unknown failure (no bench result)"
            # Compact error summary
            if "429" in first_err or "Rate limit" in first_err:
                error_summary = "429 Rate limited"
            elif "400" in first_err or "unavailable" in first_err:
                error_summary = "400 Model unavailable"
            elif "503" in first_err or "500" in first_err:
                error_summary = "5xx Server error"
            else:
                error_summary = first_err.splitlines()[0][:45]

        return ProbeResult(
            selector=selector,
            runs=runs,
            successes=successes,
            mean_ttft_s=mean_ttft_s,
            mean_tps=mean_tps,
            status=status,
            error_summary=error_summary,
        )

    except Exception as e:
        return ProbeResult(
            selector=selector,
            runs=runs,
            successes=0,
            mean_ttft_s=0.0,
            mean_tps=0.0,
            status="QUARANTINE",
            error_summary=f"Parse error: {e}",
        )


def format_results_table(results: list[ProbeResult]) -> str:
    """Format probe results into an ASCII status table."""
    lines = [
        "",
        "┌" + "─" * 46 + "┬" + "─" * 12 + "┬" + "─" * 10 + "┬" + "─" * 10 + "┬" + "─" * 24 + "┐",
        f"│ {'Model Selector':<44} │ {'Status':<10} │ {'TTFT (s)':<8} │ {'TPS':<8} │ {'Notes / Errors':<22} │",
        "├" + "─" * 46 + "┼" + "─" * 12 + "┼" + "─" * 10 + "┼" + "─" * 10 + "┼" + "─" * 24 + "┤",
    ]

    for r in results:
        if r.status == "VERIFIED":
            status_icon = "✅ VERIFIED"
        elif r.status == "PROBED":
            status_icon = "ℹ️  PROBED"
        elif r.status == "SLOW":
            status_icon = "⚠️  SLOW"
        else:
            status_icon = "🚫 QUARANTINE"
        ttft_str = f"{r.mean_ttft_s:.2f}s" if r.successes > 0 else "-"
        tps_str = f"{r.mean_tps:.1f}/s" if r.successes > 0 else "-"
        note = r.error_summary if r.error_summary else f"{r.successes}/{r.runs} passed"
        lines.append(f"│ {r.selector:<44} │ {status_icon:<10} │ {ttft_str:<8} │ {tps_str:<8} │ {note:<22} │")

    lines.append("└" + "─" * 46 + "┴" + "─" * 12 + "┴" + "─" * 10 + "┴" + "─" * 10 + "┴" + "─" * 24 + "┘")
    lines.append("")
    return "\n".join(lines)


def run_audit() -> int:
    """Run test suite and fallback chain verification."""
    print("=== OMP Runtime Audit ===")

    # 1. Run pytest suite
    print("1. Running pytest test_omp_config.py...")
    code, out, err = run_command([sys.executable, "-m", "pytest", "-q", str(TEST_CONFIG_PATH)], timeout=30)
    print(out.strip() or err.strip())
    if code != 0:
        print("❌ Test suite failed!")
        return 1
    print("\n2. Checking fallback chains...")
    code, out, err = run_command(["omp", "config", "get", "retry.fallbackChains"], timeout=10)
    if code != 0:
        print(f"❌ Failed to get fallback chains (exit {code}): {err}")
        return 1

    try:
        chains = json.loads(out)
    except Exception as e:
        print(f"❌ Failed to parse fallback chains JSON: {e}\nRaw output:\n{out[:300]}")
        return 1

    required_roles = ("smol", "commit", "tiny", "task", "default", "plan", "slow")
    missing_roles = [r for r in required_roles if r not in chains]
    if missing_roles:
        print(f"❌ Missing required roles in fallback chains: {missing_roles}")
        return 1

    empty_chains = [r for r in required_roles if not chains.get(r)]
    if empty_chains:
        print(f"❌ Empty fallback chains detected: {empty_chains}")
        return 1

    for role in required_roles:
        chain = chains[role]
        tail = chain[-1]
        print(f"   • {role:<8} ({len(chain)} hops) -> tail: {tail}")

    print("\n✅ Audit complete.")
    return 0


def main() -> int:
    parser = argparse.ArgumentParser(description="OMP Automated Free Model Prober & Auditor")
    parser.add_argument("--model", "-m", help="Probe specific model selector")
    parser.add_argument("--runs", "-r", type=int, default=3, help="Serial benchmark iterations (default: 3)")
    parser.add_argument("--max-tokens", type=int, default=64, help="Max tokens per run (default: 64)")
    parser.add_argument("--zen-only", action="store_true", help="Probe only OpenCode Zen free models")
    parser.add_argument("--openrouter-only", action="store_true", help="Probe only OpenRouter free models")
    parser.add_argument("--json", action="store_true", help="Output JSON results")
    parser.add_argument("--audit", action="store_true", help="Run full regression audit suite")

    args = parser.parse_args()

    if args.audit:
        return run_audit()

    if args.model:
        targets = [args.model]
    else:
        print("🔍 Discovering available free models...", file=sys.stderr)
        targets = discover_free_models(zen_only=args.zen_only, openrouter_only=args.openrouter_only)
        print(f"Found {len(targets)} candidates to probe: {', '.join(targets)}", file=sys.stderr)

    results: list[ProbeResult] = []
    print(f"\n🚀 Probing {len(targets)} models (runs={args.runs}, par=1, tokens={args.max_tokens})...\n", file=sys.stderr)

    for selector in targets:
        print(f"  • Probing {selector} ...", flush=True, end=" ", file=sys.stderr)
        res = probe_single_model(selector, runs=args.runs, max_tokens=args.max_tokens)
        status_label = "✅ VERIFIED" if res.status == "VERIFIED" else ("⚠️  SLOW" if res.status == "SLOW" else "🚫 QUARANTINE")
        perf_info = f"({res.mean_ttft_s:.2f}s, {res.mean_tps:.1f} tok/s)" if res.successes > 0 else f"({res.error_summary})"
        print(f"{status_label} {perf_info}", file=sys.stderr)
        results.append(res)

    if args.json:
        data = [
            {
                "selector": r.selector,
                "status": r.status,
                "successes": r.successes,
                "runs": r.runs,
                "mean_ttft_s": r.mean_ttft_s,
                "mean_tps": r.mean_tps,
                "error": r.error_summary,
            }
            for r in results
        ]
        print(json.dumps(data, indent=2))
    else:
        print(format_results_table(results))

        verified = [r for r in results if r.status == "VERIFIED"]
        if verified:
            print(f"🎉 {len(verified)} verified models available for catalog.yml admission:")
            for v in verified:
                print(f"   - {v.selector} (TTFT {v.mean_ttft_s}s, TPS {v.mean_tps}/s)")
        else:
            print("⚠️ No models met the strict admission criteria (3/3 ok, TPS >= 15).")

    return 0


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