#!/usr/bin/env python3
"""Small artifact-first benchmark runner for Brick's labeled seed corpus."""

from __future__ import annotations

import argparse
import datetime as dt
import hashlib
import json
import os
import pathlib
import platform
import re
import statistics
import subprocess
import sys
import time
from collections import Counter

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


def git_provenance(path: pathlib.Path) -> dict:
    """Describe a Git worktree and fingerprint all tracked and untracked changes."""

    discovered = subprocess.run(
        ["git", "rev-parse", "--show-toplevel"],
        cwd=path,
        text=True,
        capture_output=True,
        check=False,
    )
    if discovered.returncode != 0:
        return {
            "root": str(path.resolve()),
            "revision": None,
            "dirty": None,
            "working_tree_fingerprint": None,
        }

    repository = pathlib.Path(discovered.stdout.strip()).resolve()
    revision = subprocess.run(
        ["git", "rev-parse", "HEAD"],
        cwd=repository,
        text=True,
        capture_output=True,
        check=False,
    )
    status = subprocess.run(
        ["git", "status", "--porcelain=v1", "--untracked-files=all", "-z"],
        cwd=repository,
        capture_output=True,
        check=False,
    )
    if revision.returncode != 0 or status.returncode != 0:
        return {
            "root": str(repository),
            "revision": None,
            "dirty": None,
            "working_tree_fingerprint": None,
        }

    dirty = bool(status.stdout)
    fingerprint = None
    if dirty:
        diff = subprocess.run(
            ["git", "diff", "--binary", "HEAD"],
            cwd=repository,
            capture_output=True,
            check=False,
        )
        untracked = subprocess.run(
            ["git", "ls-files", "--others", "--exclude-standard", "-z"],
            cwd=repository,
            capture_output=True,
            check=False,
        )
        if diff.returncode == 0 and untracked.returncode == 0:
            digest = hashlib.sha256()
            digest.update(b"status\0")
            digest.update(status.stdout)
            digest.update(b"diff\0")
            digest.update(diff.stdout)
            for encoded_path in filter(None, untracked.stdout.split(b"\0")):
                entry = repository / os.fsdecode(encoded_path)
                if entry.is_symlink():
                    contents = os.fsencode(os.readlink(entry))
                elif entry.is_file():
                    contents = entry.read_bytes()
                else:
                    contents = b""
                digest.update(b"untracked\0")
                digest.update(encoded_path)
                digest.update(b"\0")
                digest.update(contents)
            fingerprint = digest.hexdigest()

    return {
        "root": str(repository),
        "revision": revision.stdout.strip(),
        "dirty": dirty,
        "working_tree_fingerprint": fingerprint,
    }


def load_plan(name: str) -> dict:
    if re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_-]*", name) is None:
        raise ValueError(f"invalid benchmark plan name: {name!r}")
    path = PLANS / f"{name}.json"
    with path.open(encoding="utf-8") as handle:
        plan = json.load(handle)
    if plan.get("schema_version") != 1 or plan.get("name") != name:
        raise ValueError(f"invalid plan identity: {path}")
    if int(plan.get("repeats", 0)) < 1:
        raise ValueError("repeats must be positive")
    return plan


def command_for(
    plan: dict, corpus: pathlib.Path, graph: pathlib.Path, html: pathlib.Path
) -> list[str]:
    features = plan.get("cargo_features", [])
    feature_args = ["--features", ",".join(features)] if features else []
    prefix = (
        [str(ROOT / "target" / "debug" / "brick")]
        if plan.get("direct_binary")
        else ["cargo", "run", "--quiet", *feature_args, "--"]
    )
    return [
        *prefix,
        "build",
        str(corpus),
        "--json",
        str(graph),
        "--html",
        str(html),
        *plan.get("build_args", []),
    ]


def fixture_assertions(metrics: dict, expected: dict) -> dict[str, bool]:
    assertions = {
        "symbols": metrics["symbol_count"] >= expected["min_symbols"],
        "dependencies": metrics["dependency_edge_count"]
        >= expected["min_dependency_edges"],
    }
    if "cycles" in expected:
        assertions["cycles"] = metrics["cycle_count"] == expected["cycles"]
    if "max_cycles" in expected:
        assertions["max_cycles"] = metrics["cycle_count"] <= expected["max_cycles"]
    for field, value in expected.get("p4", {}).items():
        assertions[f"p4.{field}"] = metrics.get("p4", {}).get(field) == value
    return assertions


def p4_relationship_quality(model: dict, labels: list[dict]) -> dict:
    nodes = {node["id"]: node for node in model["nodes"]}

    def node_identity(node_id: str) -> tuple[str | None, str | None, str | None]:
        node = nodes.get(node_id, {})
        owner = nodes.get(node.get("parent"), {})
        return node.get("path"), owner.get("name"), node.get("name")

    expected = {
        (
            label["source_path"],
            label["source_owner"],
            label["source"],
            label["target_path"],
            label["target_owner"],
            label["target"],
            label["relation"],
        )
        for label in labels
    }
    relation_kinds = {label["relation"] for label in labels}
    observed = set()
    for edge in model["edges"]:
        raw_relations = edge.get("attributes", {}).get("p4.relation")
        edge_relations = (
            raw_relations if isinstance(raw_relations, list) else [raw_relations]
        )
        for relation in edge_relations:
            if relation not in relation_kinds:
                continue
            source = node_identity(edge["source"])
            target = node_identity(edge["target"])
            if source[2] is not None and target[2] is not None:
                observed.add((*source, *target, relation))
    true_positives = len(expected & observed)
    false_positives = len(observed - expected)
    false_negatives = len(expected - observed)
    precision = true_positives / len(observed) if observed else float(not expected)
    recall = true_positives / len(expected) if expected else float(not observed)
    return {
        "expected": [list(item) for item in sorted(expected)],
        "observed": [list(item) for item in sorted(observed)],
        "true_positives": true_positives,
        "false_positives": false_positives,
        "false_negatives": false_negatives,
        "precision": round(precision, 3),
        "recall": round(recall, 3),
    }


def run_plan(
    name: str,
    out_dir: pathlib.Path | None = None,
    corpus_override: pathlib.Path | None = None,
) -> pathlib.Path:
    plan = load_plan(name)
    corpus = (corpus_override or ROOT / plan["corpus"]).resolve()
    if not corpus.is_dir():
        raise ValueError(f"benchmark corpus is missing: {corpus}")
    stamp = dt.datetime.now(dt.UTC).strftime("%Y%m%dT%H%M%SZ")
    artifact = out_dir or ROOT / "build" / "benchmarks" / name / stamp
    artifact.mkdir(parents=True, exist_ok=False)
    brick_provenance = git_provenance(ROOT)
    corpus_provenance = git_provenance(corpus)
    manifest = {
        "schema_version": 1,
        "started_at": stamp,
        "plan": name,
        "corpus": str(corpus),
        "root": str(ROOT),
        "revision": brick_provenance["revision"],
        "provenance": {
            "brick": brick_provenance,
            "corpus": corpus_provenance,
        },
        "host": {
            "platform": platform.platform(),
            "python": platform.python_version(),
        },
    }
    if plan.get("direct_binary"):
        prepare_command = ["cargo", "build", "--quiet"]
        features = plan.get("cargo_features", [])
        if features:
            prepare_command.extend(["--features", ",".join(features)])
        prepared = subprocess.run(
            prepare_command,
            cwd=ROOT,
            text=True,
            capture_output=True,
            check=False,
        )
        (artifact / "prepare.stdout").write_text(prepared.stdout, encoding="utf-8")
        (artifact / "prepare.stderr").write_text(prepared.stderr, encoding="utf-8")
        if prepared.returncode:
            raise RuntimeError(prepared.stderr or prepared.stdout)
        manifest["prepare_command"] = prepare_command
    (artifact / "run_manifest.json").write_text(
        json.dumps(manifest, indent=2) + "\n", encoding="utf-8"
    )
    samples = []
    reference = None
    reference_model = None
    for repeat in range(plan["repeats"]):
        graph = artifact / f"graph-{repeat}.json"
        html = artifact / f"building-{repeat}.html"
        started = time.perf_counter()
        command = command_for(plan, corpus, graph, html)
        rss_path = artifact / f"rss-{repeat}.json"
        max_rss_kib = None
        if plan.get("measure_rss"):
            if not hasattr(os, "wait4"):
                raise RuntimeError("RSS measurement requires os.wait4")
            stdout_path = artifact / f"command-{repeat}.stdout"
            stderr_path = artifact / f"command-{repeat}.stderr"
            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)
            max_rss_kib = int(usage.ru_maxrss)
            if sys.platform == "darwin":
                max_rss_kib //= 1024
            rss_path.write_text(
                json.dumps({"max_rss_kib": max_rss_kib}, indent=2) + "\n",
                encoding="utf-8",
            )
            completed = subprocess.CompletedProcess(
                command,
                process.returncode,
                stdout_path.read_text(encoding="utf-8"),
                stderr_path.read_text(encoding="utf-8"),
            )
        else:
            completed = subprocess.run(
                command,
                cwd=ROOT,
                text=True,
                capture_output=True,
                check=False,
            )
        elapsed_ms = (time.perf_counter() - started) * 1000
        if completed.returncode:
            raise RuntimeError(completed.stderr or completed.stdout)
        with graph.open(encoding="utf-8") as handle:
            model = json.load(handle)
        signature = json.dumps(model, sort_keys=True, separators=(",", ":"))
        if reference is None:
            reference = signature
            reference_model = model
        if signature != reference:
            raise AssertionError(f"unstable graph identity on repeat {repeat}")
        sample = {
            "repeat": repeat,
            "command": command,
            "elapsed_ms": round(elapsed_ms, 3),
            "metrics": model["metrics"],
        }
        if max_rss_kib is not None:
            sample["max_rss_kib"] = max_rss_kib
        samples.append(sample)

    metrics = samples[0]["metrics"]
    expected = plan["expect"]
    assertions = fixture_assertions(metrics, expected)
    label_quality = None
    source_quality = None
    if plan.get("source_labels"):
        if reference_model is None:
            raise AssertionError("benchmark produced no reference model")
        labels = json.loads((corpus / plan["source_labels"]).read_text())
        source_quality = source_language_quality(reference_model, labels)
        assertions["source_exact_symbols"] = source_quality["symbols"]["exact"]
        assertions["source_exact_relationships"] = source_quality["relationships"]["exact"]
        assertions["source_clean_diagnostics"] = not reference_model.get("diagnostics")
    if plan.get("p4_relationship_labels"):
        if reference_model is None:
            raise AssertionError("benchmark produced no reference model")
        label_quality = p4_relationship_quality(
            reference_model, plan["p4_relationship_labels"]
        )
        assertions["p4_relationship_precision"] = label_quality["precision"] == 1.0
        assertions["p4_relationship_recall"] = label_quality["recall"] == 1.0
    if not all(assertions.values()):
        raise AssertionError(
            f"fixture expectations failed: {assertions}; metrics={metrics}"
        )

    elapsed = [sample["elapsed_ms"] for sample in samples]
    summary = {
        "schema_version": 1,
        "plan": name,
        "repeats": len(samples),
        "elapsed_ms_median": round(statistics.median(elapsed), 3),
        "elapsed_ms_min": min(elapsed),
        "elapsed_ms_max": max(elapsed),
        "identity_stable": True,
        "full_model_stable": True,
        "fixture_assertions": assertions,
        "metrics": metrics,
        "samples": samples,
    }
    if label_quality is not None:
        summary["p4_relationship_quality"] = label_quality
    if source_quality is not None:
        summary["source_language_quality"] = source_quality
    if plan.get("measure_rss"):
        rss = [sample["max_rss_kib"] for sample in samples]
        summary["max_rss_kib_median"] = statistics.median(rss)
        summary["max_rss_kib_min"] = min(rss)
        summary["max_rss_kib_max"] = max(rss)
    (artifact / "plan.resolved.json").write_text(
        json.dumps(plan, indent=2) + "\n", encoding="utf-8"
    )
    (artifact / "summary.json").write_text(
        json.dumps(summary, indent=2) + "\n", encoding="utf-8"
    )
    report = (
        f"# Brick benchmark: {name}\n\n"
        f"- Repeats: {len(samples)}\n"
        f"- Median: {summary['elapsed_ms_median']} ms\n"
        f"- Range: {summary['elapsed_ms_min']}–{summary['elapsed_ms_max']} ms\n"
        f"- Stable full-model JSON: yes\n"
        f"- Fixture assertions: {assertions}\n"
        f"- Structural score: {metrics['structural_health']['score']}\n"
    )
    if label_quality is not None:
        report += (
            f"- P4 relationship precision: {label_quality['precision']}\n"
            f"- P4 relationship recall: {label_quality['recall']}\n"
        )
    if plan.get("measure_rss"):
        report += (
            f"- Max RSS median: {summary['max_rss_kib_median']} KiB\n"
            f"- Max RSS range: {summary['max_rss_kib_min']}–{summary['max_rss_kib_max']} KiB\n"
        )
    (artifact / "REPORT.md").write_text(report, encoding="utf-8")
    return artifact


def source_language_quality(model: dict, labels: dict) -> dict:
    """Compare every declaration and dependency, including clean-negative files."""
    nodes = {node["id"]: node for node in model["nodes"]}

    def identity(node: dict) -> str:
        if node["kind"] == "external":
            return f"external:{node.get('language')}:{node['name']}"
        names = []
        current = node
        visited = set()
        while current["kind"] in ("class", "function"):
            if current["id"] in visited:
                raise AssertionError("cyclic declaration ownership")
            visited.add(current["id"])
            qualifier = current.get("attributes", {}).get("source.qualifier")
            names.append(f"{qualifier}::{current['name']}" if qualifier else current["name"])
            current = nodes[current["parent"]]
        scope = node.get("attributes", {}).get("source.scope")
        return "::".join([node["path"], *([f"scope={scope}"] if scope else []), *reversed(names)])

    observed_symbols = [(identity(n), n["kind"]) for n in nodes.values() if n["kind"] in ("class", "function")]
    observed_relationships = [(identity(nodes[e["source"]]), e["kind"], identity(nodes[e["target"]])) for e in model["edges"] if e["kind"] in ("calls", "imports", "inherits")]

    def compare(expected: list, observed: list) -> dict:
        expected, observed = Counter(map(tuple, expected)), Counter(map(tuple, observed))
        tp = sum((expected & observed).values())
        fp, fn = observed - expected, expected - observed
        return {
            "exact": not fp and not fn,
            "true_positives": tp,
            "false_positives": list(fp.elements()),
            "false_negatives": list(fn.elements()),
            "precision": tp / sum(observed.values()) if observed else float(not expected),
            "recall": tp / sum(expected.values()) if expected else float(not observed),
        }

    result = {"symbols": compare(labels["symbols"], observed_symbols), "relationships": compare(labels["relationships"], observed_relationships)}
    # Corpus directories are the authored language strata (tsx is reported separately).
    prefixes = labels.get("language_prefixes") or {
        name: name + "/" for name in sorted({item[0].split("/")[0] for item in labels["symbols"]})
    }
    result["by_language"] = {
        language: {
            "symbols": compare([s for s in labels["symbols"] if s[0].startswith(prefix)], [s for s in observed_symbols if s[0].startswith(prefix)]),
            "relationships": compare([s for s in labels["relationships"] if s[0].startswith(prefix)], [s for s in observed_relationships if s[0].startswith(prefix)]),
        }
        for language, prefix in prefixes.items()
    }
    return result


def main() -> int:
    parser = argparse.ArgumentParser()
    subcommands = parser.add_subparsers(dest="command", required=True)
    subcommands.add_parser("self-test")
    run_parser = subcommands.add_parser("run")
    run_parser.add_argument("plan")
    run_parser.add_argument("--out-dir", type=pathlib.Path)
    run_parser.add_argument("--corpus", type=pathlib.Path)
    args = parser.parse_args()
    if args.command == "self-test":
        plan = load_plan("lexical_smoke")
        expected = plan["expect"]
        if expected["min_symbols"] < 1 or expected["cycles"] < 0:
            raise AssertionError("invalid benchmark expectations")
        if not (ROOT / plan["corpus"]).is_dir():
            raise AssertionError("benchmark corpus is missing")
        print("benchmark harness self-test passed")
        return 0
    artifact = run_plan(args.plan, args.out_dir, args.corpus)
    print(f"ARTIFACT={artifact}")
    return 0


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