#!/usr/bin/env python3
"""Generate concise loop context from GPU RANS verifier reports."""

from __future__ import annotations

import argparse
import json
import math
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Mapping

ROOT = Path(__file__).resolve().parents[1]
DEFAULT_CONTEXT = ROOT / ".loop/diagnostic-context.md"
DEFAULT_BASELINE = ROOT / ".loop/diagnostic-baseline.json"
REPORT_GLOBS = (
    "tmp/**/verifier_report.json",
    "tmp/**/report.json",
)
TMP_REPORT_GLOBS = (
    "*gpu*/report.json",
    "*gpu*/verifier_report.json",
    "worker_gpu_rans_solver*/report.json",
    "judge_gpu_rans_solver*/report.json",
    "*gpu*rans*/report.json",
    "*gpu*rans*/verifier_report.json",
)


def load_json(path: Path) -> Any | None:
    try:
        return json.loads(path.read_text())
    except Exception:
        return None


def is_verifier_report(value: Any) -> bool:
    if not isinstance(value, Mapping):
        return False
    return any(key in value for key in ("verifier_evidence", "first_divergence_summary", "artifact_comparisons", "differential_trace"))


def is_gpu_report(value: Mapping[str, Any], path: Path) -> bool:
    backend = value.get("backend") if isinstance(value.get("backend"), Mapping) else {}
    requested = str(backend.get("requested") or "").lower()
    selected = str(backend.get("selected") or "").lower()
    return requested == "gpu" or selected == "gpu" or "gpu" in path.parent.name.lower()


def discover_latest_report(root: Path) -> tuple[Path | None, Mapping[str, Any] | None]:
    candidates: list[Path] = []
    for pattern in REPORT_GLOBS:
        candidates.extend(root.glob(pattern))
    tmp_root = Path("/tmp")
    if tmp_root.exists():
        for pattern in TMP_REPORT_GLOBS:
            candidates.extend(tmp_root.glob(pattern))
    unique = sorted({path.resolve() for path in candidates if path.is_file()}, key=lambda path: path.stat().st_mtime, reverse=True)
    reports: list[tuple[Path, Mapping[str, Any]]] = []
    for path in unique:
        data = load_json(path)
        if is_verifier_report(data):
            reports.append((path, data))  # type: ignore[arg-type]
    if not reports:
        return None, None
    gpu_reports = [(path, data) for path, data in reports if is_gpu_report(data, path)]
    return (gpu_reports or reports)[0]


def get_path(value: Mapping[str, Any], path: str) -> Any:
    cursor: Any = value
    for part in path.split("."):
        if not isinstance(cursor, Mapping):
            return None
        cursor = cursor.get(part)
    return cursor


def finite_number(value: Any) -> float | None:
    try:
        number = float(value)
    except (TypeError, ValueError):
        return None
    return number if math.isfinite(number) else None


def fmt(value: Any) -> str:
    number = finite_number(value)
    if number is None:
        if value is True:
            return "true"
        if value is False:
            return "false"
        if value is None:
            return "-"
        return str(value)
    if number == 0.0:
        return "0"
    if abs(number) >= 1e4 or abs(number) < 1e-3:
        return f"{number:.3e}"
    return f"{number:.6g}"


def status_word(value: Any) -> str:
    if value is True:
        return "passed"
    if value is False:
        return "failed"
    return str(value or "unknown")


def first_nested_key(value: Any, target: str, path: str = "") -> tuple[str, Mapping[str, Any]] | None:
    if isinstance(value, Mapping):
        for key, item in value.items():
            next_path = f"{path}.{key}" if path else str(key)
            if key == target and isinstance(item, Mapping):
                return next_path, item
            found = first_nested_key(item, target, next_path)
            if found is not None:
                return found
    elif isinstance(value, list):
        for index, item in enumerate(value):
            found = first_nested_key(item, target, f"{path}[{index}]")
            if found is not None:
                return found
    return None


def artifact_checks(report: Mapping[str, Any]) -> list[dict[str, Any]]:
    out: list[dict[str, Any]] = []
    comparisons = report.get("artifact_comparisons", {}) if isinstance(report.get("artifact_comparisons"), Mapping) else {}
    for family in ("pressure_inputs", "matrix_operator", "solver"):
        family_report = comparisons.get(family) if isinstance(comparisons.get(family), Mapping) else {}
        for check in family_report.get("checks", []) if isinstance(family_report.get("checks"), list) else []:
            if not isinstance(check, Mapping):
                continue
            largest = check.get("largest_difference") if isinstance(check.get("largest_difference"), Mapping) else {}
            location = largest.get("location") if isinstance(largest.get("location"), Mapping) else {}
            out.append(
                {
                    "family": family,
                    "name": check.get("name"),
                    "allclose": check.get("allclose"),
                    "reason": check.get("reason"),
                    "max_abs": check.get("max_abs"),
                    "mean_abs": check.get("mean_abs"),
                    "rms_abs": check.get("rms_abs"),
                    "location": location,
                }
            )
    out.sort(key=lambda item: (item.get("allclose") is True, item["family"], str(item.get("name"))))
    return out


def field_checks(report: Mapping[str, Any]) -> list[dict[str, Any]]:
    out: list[dict[str, Any]] = []
    modes = report.get("modes", {}) if isinstance(report.get("modes"), Mapping) else {}
    for mode_name in ("run_one", "split"):
        mode = modes.get(mode_name) if isinstance(modes.get(mode_name), Mapping) else {}
        comparisons = mode.get("comparisons", {}) if isinstance(mode.get("comparisons"), Mapping) else {}
        for field, data in comparisons.items():
            if not isinstance(data, Mapping):
                continue
            out.append(
                {
                    "mode": mode_name,
                    "field": field,
                    "allclose": data.get("allclose"),
                    "max_abs": data.get("max_abs"),
                    "mean_abs": data.get("mean_abs"),
                    "rms_abs": data.get("rms_abs"),
                    "location": data.get("location"),
                }
            )
    out.sort(key=lambda item: (item.get("allclose") is True, item["mode"], str(item.get("field"))))
    return out

def differential_trace_checks(report: Mapping[str, Any]) -> list[dict[str, Any]]:
    trace = report.get("differential_trace", {}) if isinstance(report.get("differential_trace"), Mapping) else {}
    checks = trace.get("checkpoints", []) if isinstance(trace.get("checkpoints"), list) else []
    out: list[dict[str, Any]] = []
    for check in checks:
        if not isinstance(check, Mapping):
            continue
        largest = check.get("largest_difference") if isinstance(check.get("largest_difference"), Mapping) else {}
        location = largest.get("location") if isinstance(largest.get("location"), Mapping) else {}
        metadata = check.get("metadata") if isinstance(check.get("metadata"), Mapping) else {}
        diff = check.get("difference_stats") if isinstance(check.get("difference_stats"), Mapping) else {}
        out.append(
            {
                "name": check.get("name"),
                "status": check.get("status"),
                "reason": check.get("reason"),
                "family": metadata.get("family"),
                "lifecycle_phase": metadata.get("lifecycle_phase"),
                "substitution_supported": metadata.get("substitution_supported"),
                "max_abs": diff.get("max_abs"),
                "mean_abs": diff.get("mean_abs"),
                "rms_abs": diff.get("rms_abs"),
                "location": location,
            }
        )
    out.sort(key=lambda item: (item.get("status") == "passed", str(item.get("name"))))
    return out


def extract_metrics(report: Mapping[str, Any]) -> dict[str, Any]:
    metrics: dict[str, Any] = {
        "status.passed": report.get("status") == "passed",
        "verifier.passed": get_path(report, "verifier_evidence.passed") is True,
    }
    first = report.get("first_divergence_summary") if isinstance(report.get("first_divergence_summary"), Mapping) else {}
    if first:
        metrics["first.target"] = first.get("first_target")
        metrics["first.family"] = first.get("artifact_family")
        metrics["first.stage_group"] = first.get("stage_group")
    trace = report.get("differential_trace") if isinstance(report.get("differential_trace"), Mapping) else {}
    if trace:
        first_trace = trace.get("first_divergence") if isinstance(trace.get("first_divergence"), Mapping) else {}
        metrics["trace.status"] = trace.get("status")
        metrics["trace.first_checkpoint"] = first_trace.get("name") if first_trace else None
        metrics["trace.failed_count"] = finite_number(trace.get("failed_count")) or 0
        metrics["trace.missing_count"] = finite_number(trace.get("missing_count")) or 0
        if first_trace:
            for key in ("max_abs", "mean_abs", "rms_abs"):
                number = finite_number(first_trace.get(key))
                if number is not None:
                    metrics[f"trace.first.{key}"] = number
    for check in artifact_checks(report):
        base = f"artifact.{check['family']}.{check['name']}"
        metrics[f"{base}.allclose"] = check.get("allclose") is True
        for key in ("max_abs", "mean_abs", "rms_abs"):
            number = finite_number(check.get(key))
            if number is not None:
                metrics[f"{base}.{key}"] = number
    for check in field_checks(report):
        base = f"field.{check['mode']}.{check['field']}"
        metrics[f"{base}.allclose"] = check.get("allclose") is True
        for key in ("max_abs", "mean_abs", "rms_abs"):
            number = finite_number(check.get(key))
            if number is not None:
                metrics[f"{base}.{key}"] = number
    preconditioner = first_nested_key(report, "preconditioner_diagnostic")
    if preconditioner is not None:
        _, data = preconditioner
        for key in (
            "gpu_dilu_vs_reference_delta_l2",
            "diagonal_vs_reference_delta_l2",
        ):
            number = finite_number(data.get(key))
            if number is not None:
                metrics[f"preconditioner.{key}"] = number
        for path, metric_name in (
            ("gpu_dilu_vs_openfoam_reference.max_abs", "preconditioner.gpu_dilu_vs_reference.max_abs"),
            ("gpu_dilu_vs_openfoam_reference.rms_abs", "preconditioner.gpu_dilu_vs_reference.rms_abs"),
            ("diagonal_vs_openfoam_reference.max_abs", "preconditioner.diagonal_vs_reference.max_abs"),
            ("diagonal_vs_openfoam_reference.rms_abs", "preconditioner.diagonal_vs_reference.rms_abs"),
        ):
            number = finite_number(get_path(data, path))
            if number is not None:
                metrics[metric_name] = number
    return metrics


def classify_delta(current: Mapping[str, Any], baseline: Mapping[str, Any] | None) -> list[dict[str, Any]]:
    if not baseline:
        return [{"metric": name, "status": "newly_available", "current": value, "baseline": None} for name, value in sorted(current.items())]
    out: list[dict[str, Any]] = []
    previous = baseline.get("metrics", {}) if isinstance(baseline.get("metrics"), Mapping) else {}
    for name, value in sorted(current.items()):
        old = previous.get(name)
        status = "newly_available"
        if old is not None:
            if isinstance(value, bool) and isinstance(old, bool):
                if value == old:
                    status = "unchanged"
                elif value and not old:
                    status = "improved"
                else:
                    status = "regressed"
            else:
                now_num = finite_number(value)
                old_num = finite_number(old)
                if now_num is not None and old_num is not None:
                    tolerance = max(1e-15, abs(old_num) * 1e-9)
                    if abs(now_num - old_num) <= tolerance:
                        status = "unchanged"
                    elif now_num < old_num:
                        status = "improved"
                    else:
                        status = "regressed"
                else:
                    status = "unchanged" if value == old else "changed"
        out.append({"metric": name, "status": status, "current": value, "baseline": old})
    return out


def metric_priority(item: Mapping[str, Any]) -> tuple[int, str]:
    status_order = {"regressed": 0, "improved": 1, "newly_available": 2, "changed": 3, "unchanged": 4}
    return status_order.get(str(item.get("status")), 9), str(item.get("metric"))


def render_location(location: Any) -> str:
    if not isinstance(location, Mapping):
        return "-"
    entity = location.get("entity_kind") or "array"
    index = location.get("entity_index")
    component = location.get("component_index")
    if component is None:
        return f"{entity}[{index}]"
    return f"{entity}[{index}] component={component}"


def render_context(report_path: Path | None, report: Mapping[str, Any] | None, baseline: Mapping[str, Any] | None) -> str:
    generated = datetime.now(timezone.utc).isoformat()
    if report is None or report_path is None:
        return "\n".join(
            [
                "# GPU RANS Loop Diagnostic Context",
                "",
                f"Generated: {generated}",
                "",
                "No verifier report was found under repo tmp/ or /tmp GPU RANS work directories.",
                "Next action: run `scripts/verify_gpu_rans_solver.sh --work <work> --report <work>/report.json` or the focused verifier, then rerun this hook.",
                "",
            ]
        )

    first = report.get("first_divergence_summary") if isinstance(report.get("first_divergence_summary"), Mapping) else {}
    artifacts = report.get("intermediate_artifacts", {}) if isinstance(report.get("intermediate_artifacts"), Mapping) else {}
    families = artifacts.get("families", []) if isinstance(artifacts.get("families"), list) else []
    metrics = extract_metrics(report)
    deltas = classify_delta(metrics, baseline)
    checks = artifact_checks(report)
    fields = field_checks(report)
    trace = report.get("differential_trace", {}) if isinstance(report.get("differential_trace"), Mapping) else {}
    trace_first = trace.get("first_divergence") if isinstance(trace.get("first_divergence"), Mapping) else {}
    trace_substitution = trace.get("substitution") if isinstance(trace.get("substitution"), Mapping) else {}
    trace_checks = differential_trace_checks(report)
    preconditioner = first_nested_key(report, "preconditioner_diagnostic")
    solver_trace = first_nested_key(report, "linear_solver_trace")

    lines = [
        "# GPU RANS Loop Diagnostic Context",
        "",
        f"Generated: {generated}",
        f"Latest report: `{report_path}`",
        f"Report status: `{report.get('status')}`",
        f"Verifier evidence passed: `{get_path(report, 'verifier_evidence.passed')}`",
        "",
        "## First divergence",
        "",
        f"- Target: `{first.get('first_target') if first else None}`",
        f"- Artifact family: `{first.get('artifact_family') if first else None}`",
        f"- Stage group: `{first.get('stage_group') if first else None}`",
        f"- Evidence path: `{first.get('evidence_path') if first else None}`",
        f"- Field/reason: `{first.get('field') if first else None}` / `{first.get('reason') if first else None}`",
        "",
        "## Differential trace",
        "",
        f"- Status: `{trace.get('status') if trace else None}`",
        f"- First checkpoint: `{trace_first.get('name') if trace_first else None}`",
        f"- Lifecycle: `{get_path(trace_first, 'metadata.lifecycle_phase') if trace_first else None}`",
        f"- max_abs / rms_abs: `{fmt(trace_first.get('max_abs') if trace_first else None)}` / `{fmt(trace_first.get('rms_abs') if trace_first else None)}`",
        f"- Substitution requested/loaded: `{trace_substitution.get('requested') if trace_substitution else []}` / `{trace_substitution.get('loaded') if trace_substitution else []}`",
        f"- Source artifacts: reference=`{get_path(trace, 'roles.reference.path') if trace else None}`, candidate=`{get_path(trace, 'roles.candidate.path') if trace else None}`",
        "",
        "| Checkpoint | Status | Lifecycle | max_abs | rms_abs | Location | Substitute? |",
        "|---|---:|---|---:|---:|---|---:|",
    ]
    for check in trace_checks[:12]:
        lines.append(
            "| {name} | {status} | {phase} | {max_abs} | {rms_abs} | {location} | {substitute} |".format(
                name=check.get("name"),
                status=check.get("status"),
                phase=check.get("lifecycle_phase"),
                max_abs=fmt(check.get("max_abs")),
                rms_abs=fmt(check.get("rms_abs")),
                location=render_location(check.get("location")),
                substitute=status_word(check.get("substitution_supported")),
            )
        )
    lines.extend(["", "## Solver phase evidence", "", "| Family | Status | Why |", "|---|---:|---|"])
    for family in families:
        if not isinstance(family, Mapping):
            continue
        lines.append(f"| {family.get('name')} | {family.get('status')} | {family.get('why')} |")

    lines.extend(["", "## Numeric artifact comparisons", "", "| Family | Check | Status | max_abs | mean_abs | rms_abs | Location |", "|---|---|---:|---:|---:|---:|---|"])
    for check in checks[:16]:
        lines.append(
            "| {family} | {name} | {status} | {max_abs} | {mean_abs} | {rms_abs} | {location} |".format(
                family=check.get("family"),
                name=check.get("name"),
                status=status_word(check.get("allclose")),
                max_abs=fmt(check.get("max_abs")),
                mean_abs=fmt(check.get("mean_abs")),
                rms_abs=fmt(check.get("rms_abs")),
                location=render_location(check.get("location")),
            )
        )

    lines.extend(["", "## Field comparison symptoms", "", "| Mode | Field | Status | max_abs | mean_abs | rms_abs | Location |", "|---|---|---:|---:|---:|---:|---|"])
    for check in fields[:12]:
        lines.append(
            "| {mode} | {field} | {status} | {max_abs} | {mean_abs} | {rms_abs} | {location} |".format(
                mode=check.get("mode"),
                field=check.get("field"),
                status=status_word(check.get("allclose")),
                max_abs=fmt(check.get("max_abs")),
                mean_abs=fmt(check.get("mean_abs")),
                rms_abs=fmt(check.get("rms_abs")),
                location=render_location(check.get("location")),
            )
        )

    lines.extend(["", "## Linear solver and preconditioner trace", ""])
    if preconditioner is None:
        lines.append("No preconditioner diagnostic artifact found in the latest report.")
    else:
        path, data = preconditioner
        lines.extend(
            [
                f"Preconditioner evidence path: `{path}`",
                f"- GPU DILU vs OpenFOAM reference max_abs: `{fmt(get_path(data, 'gpu_dilu_vs_openfoam_reference.max_abs'))}`",
                f"- GPU DILU vs OpenFOAM reference rms_abs: `{fmt(get_path(data, 'gpu_dilu_vs_openfoam_reference.rms_abs'))}`",
                f"- Diagonal/current vs OpenFOAM reference max_abs: `{fmt(get_path(data, 'diagonal_vs_openfoam_reference.max_abs'))}`",
                f"- Residual entering preconditioner recorded: `{data.get('residual_entering_preconditioner') is not None}`",
            ]
        )
    if solver_trace is not None:
        path, data = solver_trace
        lines.append(f"Solver trace path: `{path}`")
        for point in data.get("trace_points", []) if isinstance(data.get("trace_points"), list) else []:
            if isinstance(point, Mapping):
                lines.append(f"- `{point.get('name')}`: {point.get('step')}")

    lines.extend(["", "## Delta versus retained baseline", "", "| Metric | Status | Current | Baseline |", "|---|---:|---:|---:|"])
    for item in sorted(deltas, key=metric_priority)[:24]:
        lines.append(f"| `{item.get('metric')}` | {item.get('status')} | {fmt(item.get('current'))} | {fmt(item.get('baseline'))} |")

    lines.extend(
        [
            "",
            "## Next target hint",
            "",
            f"Focus first on `{first.get('first_target') if first else 'unknown'}`. Treat downstream field symptoms as unreliable until that artifact or missing evidence closes.",
            "",
        ]
    )
    return "\n".join(lines)


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--root", type=Path, default=ROOT)
    parser.add_argument("--report", type=Path, default=None, help="Explicit report path; otherwise discover newest verifier report")
    parser.add_argument("--context", type=Path, default=DEFAULT_CONTEXT)
    parser.add_argument("--baseline", type=Path, default=DEFAULT_BASELINE)
    return parser.parse_args(argv)


def main(argv: list[str] | None = None) -> int:
    args = parse_args(argv)
    root = args.root.resolve()
    if args.report is not None:
        report_path = args.report.resolve()
        loaded = load_json(report_path)
        report = loaded if is_verifier_report(loaded) else None
    else:
        report_path, report = discover_latest_report(root)

    baseline = load_json(args.baseline) if args.baseline.exists() else None
    baseline_mapping = baseline if isinstance(baseline, Mapping) else None
    args.context.parent.mkdir(parents=True, exist_ok=True)
    args.context.write_text(render_context(report_path, report, baseline_mapping), encoding="utf-8")

    if report is not None and report_path is not None:
        current = {
            "updated_at": datetime.now(timezone.utc).isoformat(),
            "report": str(report_path),
            "metrics": extract_metrics(report),
            "first_divergence_summary": report.get("first_divergence_summary"),
            "differential_trace": report.get("differential_trace"),
        }
        args.baseline.parent.mkdir(parents=True, exist_ok=True)
        args.baseline.write_text(json.dumps(current, indent=2, sort_keys=True) + "\n", encoding="utf-8")
        print(f"updated diagnostic context: {args.context} from {report_path}")
    else:
        print(f"updated diagnostic context without report: {args.context}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())