#!/usr/bin/env python3
"""Artifact-first benchmark for Brick's declarative architecture policy gate."""

from __future__ import annotations

import argparse
import datetime as dt
import json
import os
import pathlib
import platform
import shutil
import statistics
import subprocess
import sys
import time

import runner

ROOT = pathlib.Path(__file__).resolve().parents[2]
PLAN = ROOT / "tests" / "benchmarks" / "plans" / "policy.json"


def load_plan() -> dict:
    plan = json.loads(PLAN.read_text(encoding="utf-8"))
    if plan.get("schema_version") != 1 or plan.get("name") != "policy":
        raise ValueError("invalid policy benchmark plan identity")
    if int(plan.get("repeats", 0)) < 3:
        raise ValueError("policy benchmark requires at least three repeats")
    if not (ROOT / plan.get("project", "")).is_dir():
        raise ValueError("policy benchmark project is missing")
    for case in plan.get("cases", []):
        if not (ROOT / case.get("policy", "")).is_file():
            raise ValueError(f"policy benchmark fixture is missing: {case.get('name')!r}")
        if type(case.get("expected_passed")) is not bool:
            raise ValueError(f"policy expected_passed must be boolean: {case.get('name')!r}")
        if not isinstance(case.get("expected_matches"), int):
            raise ValueError(f"policy expected_matches must be integer: {case.get('name')!r}")
    if not plan.get("cases"):
        raise ValueError("policy benchmark has no cases")
    return plan


def run_logged(command: list[str], stdout_path: pathlib.Path, stderr_path: pathlib.Path) -> None:
    completed = subprocess.run(command, cwd=ROOT, text=True, capture_output=True, check=False)
    stdout_path.write_text(completed.stdout, encoding="utf-8")
    stderr_path.write_text(completed.stderr, encoding="utf-8")
    if completed.returncode:
        raise RuntimeError(completed.stderr or completed.stdout)


def run_measured(
    command: list[str], stdout_path: pathlib.Path, stderr_path: pathlib.Path
) -> tuple[int, float, int]:
    if not hasattr(os, "wait4"):
        raise RuntimeError("RSS measurement requires os.wait4")
    started = time.perf_counter()
    with stdout_path.open("wb") as stdout, stderr_path.open("wb") as stderr:
        process = subprocess.Popen(command, cwd=ROOT, stdout=stdout, stderr=stderr)
        _, status, usage = os.wait4(process.pid, 0)
        process.returncode = os.waitstatus_to_exitcode(status)
    elapsed_ms = round((time.perf_counter() - started) * 1000, 3)
    max_rss_kib = int(usage.ru_maxrss)
    if sys.platform == "darwin":
        max_rss_kib //= 1024
    return process.returncode, elapsed_ms, max_rss_kib


def run(out_dir: pathlib.Path | None = None) -> pathlib.Path:
    plan = load_plan()
    stamp = dt.datetime.now(dt.UTC).strftime("%Y%m%dT%H%M%SZ")
    artifact = (out_dir or ROOT / "build" / "benchmarks" / "policy" / stamp).resolve()
    artifact.mkdir(parents=True, exist_ok=False)
    binary = ROOT / "target" / "debug" / "brick"
    prepare = ["cargo", "build", "--quiet"]
    run_logged(prepare, artifact / "prepare.stdout", artifact / "prepare.stderr")
    repository = artifact / "repository"
    shutil.copytree(ROOT / plan["project"], repository)
    manifest = {
        "schema_version": 1,
        "started_at": stamp,
        "plan": "policy",
        "root": str(ROOT),
        "provenance": runner.git_provenance(ROOT),
        "host": {"platform": platform.platform(), "python": platform.python_version()},
        "prepare_command": prepare,
    }
    (artifact / "run_manifest.json").write_text(
        json.dumps(manifest, indent=2) + "\n", encoding="utf-8"
    )
    (artifact / "plan.resolved.json").write_text(
        json.dumps(plan, indent=2) + "\n", encoding="utf-8"
    )

    samples: list[dict] = []
    signatures: dict[str, str] = {}
    for repeat in range(plan["repeats"]):
        graph = artifact / f"graph-{repeat}.json"
        run_logged(
            [
                str(binary),
                "build",
                str(repository),
                "--json",
                str(graph),
                "--html",
                str(artifact / f"building-{repeat}.html"),
            ],
            artifact / f"build-{repeat}.stdout",
            artifact / f"build-{repeat}.stderr",
        )
        for case in plan["cases"]:
            name = case["name"]
            result_path = artifact / f"{name}-result-{repeat}.json"
            command = [
                str(binary),
                "check",
                str(graph),
                "--min-score",
                "0",
                "--policy",
                str(ROOT / case["policy"]),
                "--policy-result-json",
                str(result_path),
            ]
            returncode, elapsed_ms, max_rss_kib = run_measured(
                command,
                artifact / f"{name}-check-{repeat}.stdout",
                artifact / f"{name}-check-{repeat}.stderr",
            )
            expected_returncode = 0 if case["expected_passed"] else 1
            if returncode != expected_returncode:
                stderr = (artifact / f"{name}-check-{repeat}.stderr").read_text(encoding="utf-8")
                raise AssertionError(
                    f"{name} return code {returncode}, expected {expected_returncode}: {stderr}"
                )
            result = json.loads(result_path.read_text(encoding="utf-8"))
            if result.get("policy_result_schema_version") != 1:
                raise AssertionError(f"{name} policy result schema differs")
            if result.get("policy_schema_version") != 1:
                raise AssertionError(f"{name} input policy schema differs")
            if result.get("building_schema_version") != 4:
                raise AssertionError(f"{name} building schema differs")
            expected_policy = json.loads((ROOT / case["policy"]).read_text(encoding="utf-8"))
            expected_rules = {rule["id"]: rule for rule in expected_policy["rules"]}
            expected_rule_ids = sorted(expected_rules)
            emitted_rule_ids = [rule.get("rule_id") for rule in result["rules"]]
            if emitted_rule_ids != expected_rule_ids:
                raise AssertionError(
                    f"{name} emitted rules {emitted_rule_ids}, expected {expected_rule_ids}"
                )
            required_rule_fields = {
                "rule_id",
                "kind",
                "source_path_prefix",
                "target_path_prefix",
                "edge_kinds",
                "max_matches",
                "matched_edges",
                "passed",
            }
            for rule in result["rules"]:
                if set(rule) != required_rule_fields:
                    raise AssertionError(f"{name} policy result rule contract differs: {set(rule)}")
                expected_rule = expected_rules.get(rule["rule_id"])
                if expected_rule is None:
                    raise AssertionError(f"{name} emitted unknown rule {rule['rule_id']}")
                for field in [
                    "kind",
                    "source_path_prefix",
                    "target_path_prefix",
                    "edge_kinds",
                    "max_matches",
                ]:
                    expected = expected_rule[field]
                    if field == "edge_kinds":
                        expected = sorted(expected)
                    if rule[field] != expected:
                        raise AssertionError(f"{name} rule field {field} differs")
            matches = sum(len(rule["matched_edges"]) for rule in result["rules"])
            if result.get("passed") != case["expected_passed"]:
                raise AssertionError(f"{name} policy pass state differs")
            if matches != case["expected_matches"]:
                raise AssertionError(
                    f"{name} matched {matches} edges, expected {case['expected_matches']}"
                )
            for rule in result["rules"]:
                for match in rule["matched_edges"]:
                    evidence_records = match["edge"].get("evidence")
                    if not evidence_records:
                        raise AssertionError(f"{name} emitted a match without evidence")
                    for evidence in evidence_records:
                        if set(evidence) != {"lane", "backend", "confidence", "provenance"}:
                            raise AssertionError(
                                f"{name} emitted incomplete evidence fields: {set(evidence)}"
                            )
            signature = json.dumps(result, sort_keys=True, separators=(",", ":"))
            if name in signatures and signatures[name] != signature:
                raise AssertionError(f"unstable policy result for {name} repeat {repeat}")
            signatures[name] = signature
            samples.append(
                {
                    "repeat": repeat,
                    "case": name,
                    "command": command,
                    "returncode": returncode,
                    "elapsed_ms": elapsed_ms,
                    "max_rss_kib": max_rss_kib,
                    "passed": result["passed"],
                    "matched_edges": matches,
                }
            )

    cases = {}
    for case in plan["cases"]:
        case_samples = [sample for sample in samples if sample["case"] == case["name"]]
        elapsed = [sample["elapsed_ms"] for sample in case_samples]
        rss = [sample["max_rss_kib"] for sample in case_samples]
        cases[case["name"]] = {
            "identity_stable": True,
            "passed": case["expected_passed"],
            "matched_edges": case["expected_matches"],
            "elapsed_ms_median": round(statistics.median(elapsed), 3),
            "elapsed_ms_min": min(elapsed),
            "elapsed_ms_max": max(elapsed),
            "max_rss_kib_median": statistics.median(rss),
            "max_rss_kib_min": min(rss),
            "max_rss_kib_max": max(rss),
        }
    summary = {
        "schema_version": 1,
        "plan": "policy",
        "repeats": plan["repeats"],
        "cases": cases,
        "samples": samples,
    }
    (artifact / "summary.json").write_text(
        json.dumps(summary, indent=2) + "\n", encoding="utf-8"
    )
    report = ["# Brick benchmark: policy", ""]
    for name, result in cases.items():
        report.extend(
            [
                f"## {name}",
                "",
                "- Stable result: yes",
                f"- Passed: {result['passed']}",
                f"- Matched edges: {result['matched_edges']}",
                f"- Median: {result['elapsed_ms_median']} ms",
                f"- Range: {result['elapsed_ms_min']}–{result['elapsed_ms_max']} ms",
                f"- Max RSS median: {result['max_rss_kib_median']} KiB",
                "",
            ]
        )
    (artifact / "REPORT.md").write_text("\n".join(report), encoding="utf-8")
    return artifact


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("command", choices=["self-test", "run"])
    parser.add_argument("--out-dir", type=pathlib.Path)
    args = parser.parse_args()
    plan = load_plan()
    if args.command == "self-test":
        print(f"policy benchmark self-test passed ({len(plan['cases'])} cases)")
        return 0
    artifact = run(args.out_dir)
    print(f"ARTIFACT={artifact}")
    return 0


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except (OSError, ValueError, RuntimeError, AssertionError) as error:
        print(f"policy benchmark: {error}", file=sys.stderr)
        raise SystemExit(1) from error