#!/usr/bin/env python3

"""Run the branch-triage-shape correctness lane over a local git fixture.

Ensures (or generates) the branch triage fixture, probes Oak for branch review

or batch triage JSON support, scores measured rows against BRANCH_MANIFEST.json,

and emits task-loop lane rows. Missing Oak capability yields explicit skip rows

(returncode 77) with null metrics — never a fabricated pass.

"""

from __future__ import annotations

import argparse

import json

import shutil

import subprocess

import sys

import tempfile

from datetime import datetime, timezone

from pathlib import Path

from typing import Any

ROOT = Path(__file__).resolve().parents[1]

sys.path.insert(0, str(ROOT / "scripts"))

from oakbench import environment as oakbench_environment

from oakbench import platform_clock

from oakbench.branch_triage_shape import (

    MANIFEST_FILENAME,

    TRIAGE_METRIC_KEYS,

    load_manifest,

    score_manifest_against_actuals,

    triage_metrics_row_fields,

    unmeasured_triage_metrics,

    validate_manifest,

    validate_triage_metrics,

)

from oakbench.fixture_registry import load_fixture_registry

from oakbench.results import ResultsStore

from oakbench.rows import SKIP_RETURNCODE, validate_rows

GENERATOR = ROOT / "scripts" / "make_branch_triage_fixture.py"

DEFAULT_WORKDIR = Path(tempfile.gettempdir()) / "oak-branch-triage-fixture"

DEFAULT_RESULTS = ROOT / "results" / "branch-triage-fixture"

PROFILE = "branch-triage-fixture"

FIXTURE_ID = "bench-branch-triage-shape-v2"

SCENARIO = "branch_triage_shape"

MEASUREMENT_SOURCE = "direct_cli_timed_subprocess"

DEFAULT_SEED = "oakbench-branch-triage-v2"

DEFAULT_BRANCH_COUNT = 6

def utc_bench_id() -> str:

    return datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")

def ensure_fixture(workdir: Path, seed: str, branch_count: int) -> Path:

    fixture_root = workdir / f"{seed}-n{branch_count}"

    manifest_path = fixture_root / MANIFEST_FILENAME

    if not manifest_path.is_file():

        proc = subprocess.run(

            [

                sys.executable,

                str(GENERATOR),

                "--out",

                str(fixture_root),

                "--seed",

                seed,

                "--branch-count",

                str(branch_count),

            ],

            capture_output=True,

            text=True,

        )

        if proc.returncode != 0:

            raise RuntimeError(proc.stderr or proc.stdout or "fixture generation failed")

    errors = validate_manifest(load_manifest(manifest_path))

    if errors:

        raise RuntimeError("; ".join(errors))

    return fixture_root

def oak_supports_review_triage(oak_bin: Path, repo: Path) -> tuple[bool, str]:

    proc = subprocess.run(

        [str(oak_bin), "branch", "review", "main", "--merge-preview", "--json"],

        cwd=repo,

        capture_output=True,

        text=True,

        env=oakbench_environment.base_env(),

    )

    if proc.returncode != 0:

        return False, "branch_review_unavailable"

    try:

        payload = json.loads(proc.stdout)

    except json.JSONDecodeError:

        return False, "branch_review_invalid_json"

    triage = payload.get("triage", payload)

    if not isinstance(triage, dict):

        return False, "branch_review_missing_triage"

    if "recommended_action" not in triage:

        return False, "branch_review_missing_recommended_action"

    return True, "branch_review"

def oak_supports_batch_triage(oak_bin: Path, repo: Path) -> tuple[bool, str]:

    proc = subprocess.run(

        [str(oak_bin), "branch", "triage", "--against", "main", "--status", "open", "--json"],

        cwd=repo,

        capture_output=True,

        text=True,

        env=oakbench_environment.base_env(),

    )

    if proc.returncode != 0:

        return False, "branch_triage_unavailable"

    try:

        payload = json.loads(proc.stdout)

    except json.JSONDecodeError:

        return False, "branch_triage_invalid_json"

    if not isinstance(payload, dict):

        return False, "branch_triage_invalid_payload"

    return True, "branch_triage"

def review_branch(oak_bin: Path, repo: Path, branch: str) -> dict[str, Any]:

    proc = subprocess.run(

        [str(oak_bin), "branch", "review", branch, "--merge-preview", "--json"],

        cwd=repo,

        capture_output=True,

        text=True,

        check=False,

        env=oakbench_environment.base_env(),

    )

    if proc.returncode != 0:

        return {}

    try:

        return json.loads(proc.stdout)

    except json.JSONDecodeError:

        return {}

def batch_triage(oak_bin: Path, repo: Path) -> dict[str, dict[str, Any]]:

    proc = subprocess.run(

        [str(oak_bin), "branch", "triage", "--against", "main", "--status", "open", "--json"],

        cwd=repo,

        capture_output=True,

        text=True,

        check=False,

        env=oakbench_environment.base_env(),

    )

    if proc.returncode != 0:

        return {}

    try:

        payload = json.loads(proc.stdout)

    except json.JSONDecodeError:

        return {}

    rows = payload.get("branches") or payload.get("rows") or []

    actuals: dict[str, dict[str, Any]] = {}

    if isinstance(rows, list):

        for row in rows:

            if isinstance(row, dict) and row.get("branch"):

                actuals[str(row["branch"])] = row

    return actuals

def base_row(

    *,

    bench_id: str,

    operation: str,

    run_index: int,

    elapsed_ms: float,

    branch_count: int,

    registry_fields: dict[str, Any],

) -> dict[str, Any]:

    return {

        "bench_id": bench_id,

        "profile": PROFILE,

        "scenario": SCENARIO,

        "operation": operation,

        "run": run_index,

        "subject": "oak",

        "subject_kind": "oak",

        "elapsed_ms": round(elapsed_ms, 3),

        "returncode": 0,

        "command": [],

        "measurement_source": MEASUREMENT_SOURCE,

        "branch_count": branch_count,

        **registry_fields,

    }

def skip_row(

    meta: dict[str, Any],

    operation: str,

    run_index: int,

    reason: str,

    elapsed_ms: float = 0.0,

) -> dict[str, Any]:

    return {

        **meta,

        "operation": operation,

        "run": run_index,

        "elapsed_ms": round(elapsed_ms, 3),

        "returncode": SKIP_RETURNCODE,

        "skipped": True,

        "skip_reason": reason,

        "command": [],

        **triage_metrics_row_fields(unmeasured_triage_metrics(branches_expected=int(meta.get("branch_count") or 0))),

    }

def scoring_failures(scored: dict[str, Any]) -> list[dict[str, Any]]:

    return [

        row

        for row in scored["branch_results"]

        if not row.get("measured")

        or not row.get("correct")

        or row.get("false_close")

        or row.get("false_vcs_merge_safe")

        or row.get("false_merge_allowed")

    ]

def run_lane(args: argparse.Namespace) -> list[dict[str, Any]]:

    registry = load_fixture_registry()

    spec = registry.get(FIXTURE_ID)

    registry_fields = spec.row_fields() if spec else {

        "fixture_id": FIXTURE_ID,

        "fixture_generator": "scripts/make_branch_triage_fixture.py",

        "fixture_version": "v1",

        "fixture_seed": args.seed,

        "fixture_manifest_sha256": None,

    }

    bench_id = args.bench_id or utc_bench_id()

    meta = base_row(

        bench_id=bench_id,

        operation="triage.fixture.ensure",

        run_index=args.run,

        elapsed_ms=0.0,

        branch_count=args.branch_count,

        registry_fields=registry_fields,

    )

    start = platform_clock.monotonic_ms()

    try:

        fixture_root = ensure_fixture(args.workdir, args.seed, args.branch_count)

    except RuntimeError as exc:

        return [

            skip_row(meta, "triage.fixture.ensure", args.run, f"fixture_generation_failed:{exc}",

                     platform_clock.monotonic_ms() - start)

        ]

    repo = fixture_root / "repo"

    manifest = load_manifest(fixture_root / MANIFEST_FILENAME)

    rows: list[dict[str, Any]] = [

        {

            **meta,

            "elapsed_ms": round(platform_clock.monotonic_ms() - start, 3),

            "fixture_root": str(fixture_root),

            "branches_expected": len(manifest.get("branches") or []),

        }

    ]

    oak_bin = Path(args.oak_bin) if args.oak_bin else shutil.which("oak")

    if oak_bin is None or not Path(oak_bin).exists():

        rows.append(

            skip_row(meta, "triage.score", args.run, "oak_binary_unavailable",

                     platform_clock.monotonic_ms() - start)

        )

        return rows

    oak_path = Path(oak_bin)

    probe_start = platform_clock.monotonic_ms()

    batch_ok, batch_reason = oak_supports_batch_triage(oak_path, repo)

    review_ok, review_reason = oak_supports_review_triage(oak_path, repo)

    if not batch_ok and not review_ok:

        rows.append(

            skip_row(

                meta,

                "triage.score",

                args.run,

                f"branch_triage_capability_missing:{batch_reason}:{review_reason}",

                platform_clock.monotonic_ms() - probe_start,

            )

        )

        return rows

    actuals: dict[str, dict[str, Any]] = {}

    if batch_ok:

        actuals = batch_triage(oak_path, repo)

    if not actuals and review_ok:

        for branch in manifest.get("branches") or []:

            name = str(branch.get("name") or "")

            if name:

                actuals[name] = review_branch(oak_path, repo, name)

    score_start = platform_clock.monotonic_ms()

    scored = score_manifest_against_actuals(manifest, actuals)

    metrics = scored["metrics"]

    metric_fields = triage_metrics_row_fields(metrics)

    metric_errors = validate_triage_metrics(metric_fields)

    failures = scoring_failures(scored)

    rows.append(

        {

            **meta,

            "operation": "triage.score",

            "elapsed_ms": round(platform_clock.monotonic_ms() - score_start, 3),

            "returncode": 1 if metric_errors or failures else 0,

            "command": [str(oak_path), "branch", "triage_or_review", "--json"],

            "triage_provider": batch_reason if batch_ok else review_reason,

            "branch_results": scored["branch_results"],

            **metric_fields,

        }

    )

    return rows

def parse_args(argv: list[str] | None = None) -> argparse.Namespace:

    parser = argparse.ArgumentParser(description="Branch triage fixture correctness lane")

    parser.add_argument("--workdir", type=Path, default=DEFAULT_WORKDIR)

    parser.add_argument("--results", type=Path, default=DEFAULT_RESULTS)

    parser.add_argument("--seed", default=DEFAULT_SEED)

    parser.add_argument("--branch-count", type=int, default=DEFAULT_BRANCH_COUNT)

    parser.add_argument("--quick", action="store_true")

    parser.add_argument("--medium", action="store_true")

    parser.add_argument("--run", type=int, default=0)

    parser.add_argument("--bench-id")

    parser.add_argument("--oak-bin")

    parser.add_argument("--dry-run", action="store_true")

    return parser.parse_args(argv)

def main(argv: list[str] | None = None) -> int:

    args = parse_args(argv)

    if args.quick:

        args.branch_count = DEFAULT_BRANCH_COUNT

    if args.medium:

        args.branch_count = 60

    rows = run_lane(args)

    contract_errors = validate_rows(rows, "task-loop")

    for row in rows:

        metric_payload = {key: row.get(key) for key in TRIAGE_METRIC_KEYS}

        contract_errors.extend(validate_triage_metrics(metric_payload))

    if contract_errors:

        print("\n".join(contract_errors), file=sys.stderr)

        return 1

    if args.dry_run:

        print(json.dumps(rows, indent=2, sort_keys=True))

        return 0

    store = ResultsStore(args.results, lane="task-loop")

    bench_id = str(rows[0]["bench_id"])

    store.write(bench_id, rows)

    failures = [row for row in rows if row.get("returncode") not in (0, SKIP_RETURNCODE)]

    return 1 if failures else 0

if __name__ == "__main__":

    raise SystemExit(main())