"""Resumable, serial labelled-function traces; host fetches pinned sources, Docker stays offline.

Pair blob identities resolve through catalogs to real source commits. The analyzer
is git archive HEAD src plus a hashed snapshot of this uncommitted tool.
"""

from __future__ import annotations

import argparse
import ast
import fcntl
import hashlib
import json
import os
import re
import shutil
import subprocess
import tempfile
import textwrap
import time
import tomllib
from collections import Counter, defaultdict
from pathlib import Path

from engine_trace_worker import EXCLUDED_DIRS, safe_path, unit_record, write_json

from openultrasast.learn.examples import read_jsonl
from openultrasast.learn.labels import load_sources, repo_name
from openultrasast.model.specs import taint_specs
from openultrasast.model.taint import request_params
from openultrasast.pairs import _materialize_side, load_pair_catalog
from openultrasast.plane.engine_alerts import export_pin
from openultrasast.plane.memory import FileStore
from openultrasast.plane.tasks.alerts import engine_languages
from openultrasast.preprocess import detect_language

ROOT = Path(__file__).resolve().parents[2]
DEFAULT_OUT = Path.home() / "ousast-results/plane/engine-trace"
SNAPSHOT_FILES = (
    "benchmarks/learn/engine_trace.py",
    "benchmarks/learn/engine_trace_worker.py",
    "benchmarks/learn/engine_trace_parse.py",
    "src/openultrasast/model/taint.py",
    "src/openultrasast/cpg/queries/taint.sc",
)
MATERIALIZATION_VERSION = 3
EXPORT_CAP = 20 * 1024 * 1024

STATUSES = ("path", "asked-nothing", "failed", "timeout", "unsupported")


def digest(value):
    return hashlib.sha256(json.dumps(value, sort_keys=True).encode()).hexdigest()


def pin_name(pin):
    return digest([pin["repo"], pin["pin"]]) + ".json"


def join_units(units, examples, include_typescript=False):
    index = {row["id"]: row for row in examples}
    missing = [u["unit"] for u in units if u["unit"] not in index]
    if missing:
        raise ValueError(f"example snapshot missing {len(missing)}/{len(units)} unit IDs; supply --examples JSONL; first: {missing[0]}")
    covered = engine_languages() - {"typescript"}
    if include_typescript:
        covered |= {"typescript"}
    result = []
    for unit in units:
        example = index[unit["unit"]]
        if any(example[key] != unit[key] for key in ("family", "label", "group")):
            raise ValueError(f"unit/example identity mismatch: {unit['unit']}")
        file, sep, function = example["candidate"].partition("::")
        if not sep or not function:
            raise ValueError(f"unit has no labelled function: {unit['unit']}")
        language = detect_language(Path(file)) or example.get("language", "other")
        model_language = "javascript" if language == "typescript" else language
        spec = taint_specs(language=model_language).get(unit["family"]) if language in covered else None
        result.append(
            {
                **unit,
                "repo": example["repo"],
                "pin": example["pin"],
                "file": file,
                "function": function,
                "language": language,
                "source": example["source"],
                "excerpt": example.get("excerpt", example["source"] == "pairs" or example["source"].startswith("advisory-fixes")),
                "pin_role": example.get("pin_role", ""),
                "supported": spec is not None,
                "reason": "" if spec else f"no admitted taint model: {language}/{unit['family']}",
            }
        )
    return result


def make_plan(
    units,
    out,
    only_family=None,
    limit=0,
    *,
    include_typescript=False,
    only_source=(),
    rerun_status=("failed",),
    allow_excerpt_fallback=False,
):
    grouped = defaultdict(list)
    for unit in units:
        if (not only_family or unit["family"] == only_family) and (not only_source or unit["source"] in only_source):
            # Include admission options even when a unit's support happens to stay
            # unchanged. Old results cannot mark a different selection as complete.
            identity = {
                "unit": unit,
                "include_typescript": include_typescript,
                "materialization_version": MATERIALIZATION_VERSION,
                "allow_excerpt_fallback": allow_excerpt_fallback,
            }
            grouped[unit["repo"], unit["pin"]].append({**unit, "input_digest": digest(identity)})
    plan = []
    for (repo, pin), rows in sorted(grouped.items()):
        item = {
            "repo": repo,
            "pin": pin,
            "materialization_version": MATERIALIZATION_VERSION,
            "units": sorted(rows, key=lambda r: r["unit"]),
        }
        item["selection"] = digest(item["units"])
        done = out / pin_name(item)
        if done.exists():
            old = json.loads(done.read_text())
            recorded = {r["unit"]: r.get("input_digest") for r in old.get("units", [])}
            if (
                old.get("done")
                and old.get("materialization_version") == MATERIALIZATION_VERSION
                and not any(r.get("status") in rerun_status for r in old.get("units", []) if r["unit"] in {u["unit"] for u in rows})
                and all(recorded.get(r["unit"]) == r["input_digest"] for r in rows)
            ):
                continue
        sources = {r["source"] for r in rows}
        # Inventory medians: harvest sides 31s, advisory sides 47s; v1 pins 182s, v2 pins 498s.
        item["estimated_seconds"] = (
            (47 if any(s.startswith("advisory") for s in sources) else 31)
            if any(r.get("excerpt") for r in rows) or any(s.startswith("advisory") for s in sources)
            else (182 if sources == {"population-v1"} else 498)
        )
        if not any(r["supported"] for r in rows):
            item["estimated_seconds"] = 0
        plan.append(item)
    # limit counts runnable pins, not unsupported records.
    if limit:
        chosen = []
        count = 0
        for pin in plan:
            runnable = any(r["supported"] for r in pin["units"])
            if runnable and count >= limit:
                continue
            chosen.append(pin)
            count += runnable
        return chosen
    return plan


class Inputs:
    """Resolve spent corpora; fetch pair source commits into the host cache."""

    def __init__(self, root, cache, allow_excerpt_fallback=False):
        self.cache = cache
        self.allow_excerpt_fallback = allow_excerpt_fallback
        self.catalog_rows = defaultdict(list)
        for catalog in (root / "benchmarks/pairs").rglob("catalog.toml"):
            for item in tomllib.loads(catalog.read_text()).get("pair", []):
                self.catalog_rows[item["name"]].append((catalog.parent.name, item))
        self.pairs = []
        self.repos = defaultdict(list)
        sources = load_sources()
        for source in sources.sources:
            if source.kind == "pairs":
                self.pairs.extend(load_pair_catalog(root / source.files[0]))
            elif source.kind == "population":
                if source.id not in {"population-v1", "population-v2"}:
                    continue
                data = tomllib.loads((root / source.files[0]).read_text())
                for case in data.get("case", []):
                    self.repos[repo_name(case["repo"])].append((cache / "independent" / case["id"], case))
            elif source.kind == "recipes":
                for file in source.files:
                    data = tomllib.loads((root / file).read_text())
                    recipe = data.get("repo", data)
                    # Recipes carry a top-level URL and a pinned checkout cache.
                    if isinstance(recipe, dict):
                        self.repos[repo_name(str(recipe.get("url", "")))].append((cache / "repos" / str(data.get("name", "")), data))
        self.hashes = {}

    def pair_for(self, row):
        if row["pin_role"] not in {"fixed", "vulnerable", "vuln"}:
            raise ValueError(f"ambiguous pair side: {row['pin_role']}")
        side = "fixed" if row["pin_role"] == "fixed" else "vuln"
        candidates = [
            c for c in self.pairs if c.relpath == row["file"] and repo_name(c.repo or f"{c.slice}:{c.name}") == repo_name(row["repo"])
        ]
        matches = []
        for case in candidates:
            path = case.fixed_file if side == "fixed" else case.vuln_file
            if path.is_file():
                if path not in self.hashes:
                    data = path.read_bytes()
                    # Harvest hashes decoded text, so universal newline conversion is part
                    # of its identity (one recorded ThreatByte excerpt uses CRLF).
                    normalized = path.read_text().encode("utf-8")
                    self.hashes[path] = {
                        hashlib.sha1(b"blob " + str(len(body)).encode() + b"\0" + body).hexdigest() for body in (data, normalized)
                    }
                if row["pin"] in self.hashes[path]:
                    matches.append(case)
        if len(matches) == 1:
            return matches[0], side
        if len(candidates) == 1:
            # Earlier stores used label_pin rather than the excerpt digest.
            from openultrasast.learn.examples import label_pin

            case = candidates[0]
            if row["pin"] == label_pin({"source": "pairs", "source_ref": case.name, "pin_role": row["pin_role"]}):
                return case, side
        raise ValueError(f"cannot uniquely resolve excerpt {row['repo']}@{row['pin']}:{row['file']}")

    def catalog_for(self, case):
        matches = [
            (source, item)
            for source, item in self.catalog_rows[case.name]
            if item.get("relpath") == case.relpath and repo_name(item.get("repo", "")) == repo_name(case.repo)
        ]
        if len(matches) != 1:
            raise ValueError(f"ambiguous catalog source for {case.name}")
        return matches[0]

    def source_identity(self, case, side):
        _, item = self.catalog_for(case)
        parent, commit = item.get("parent", ""), item.get("commit", "")
        if side not in {"vuln", "fixed"} or not all(re.fullmatch(r"[0-9a-fA-F]{40}", sha) for sha in (parent, commit)) or parent == commit:
            raise ValueError(f"ambiguous catalog commits for {case.name}: need distinct parent and commit SHAs")
        repo = item.get("repo", "")
        if not repo:
            raise ValueError(f"missing catalog repository for {case.name}")
        return repo, parent if side == "vuln" else commit

    def fetch_source(self, repo, commit, roots):
        clone = self.cache / "trace-sources" / digest([repo, commit, roots])
        clone.mkdir(parents=True, exist_ok=True)
        url = repo if repo.startswith("https://") else f"https://github.com/{repo.removesuffix('.git')}.git"
        env = {**os.environ, "GIT_TERMINAL_PROMPT": "0"}
        env.pop("GIT_NO_LAZY_FETCH", None)  # sparse checkout may fetch the selected blobs

        def git(*args):
            return subprocess.run(["git", "-C", str(clone), *args], env=env, capture_output=True, text=True, check=True, timeout=300)

        git("init")
        git("config", "remote.origin.url", url)
        git("config", "remote.origin.promisor", "true")
        git("config", "remote.origin.partialclonefilter", "blob:none")
        git("fetch", "--depth", "1", "--filter=blob:none", url, commit)
        git("sparse-checkout", "init", "--cone")
        git("sparse-checkout", "set", "--", *roots)
        git("checkout", "--detach", "--force", commit)
        actual = git("rev-parse", "HEAD").stdout.strip()
        if actual != commit:
            raise ValueError(f"fetched commit mismatch: {actual} != {commit}")
        return clone

    def export_local(self, clone, pin, target):
        with tempfile.TemporaryDirectory(prefix="trace-export-", dir=target.parent) as scratch:
            export_pin(clone, pin["pin"], Path(scratch))
            roots = sorted({source_root(r["file"], r["language"]) for r in pin["units"]})
            proof = export_sources(Path(scratch), target, pin["units"], roots)
        return {**proof, "mode": "source", "commit": pin["pin"]}

    def materialize(self, pin, target):
        rows = pin["units"]
        for row in rows:
            safe_path(target, row["file"])
            if row.get("mapping_error"):
                raise ValueError(row["mapping_error"])
        if any(r.get("excerpt") or r["source"].startswith("advisory-fixes") for r in rows):
            resolved = [self.pair_for(row) for row in rows]
            identities = set()
            catalog_files = []
            for case, side in resolved:
                _, item = self.catalog_for(case)
                parent, commit = item.get("parent", ""), item.get("commit", "")
                if parent != commit and all(re.fullmatch(r"[0-9a-fA-F]{40}", sha) for sha in (parent, commit)):
                    identities.add(self.source_identity(case, side))
                else:
                    file = case.fixed_file if side == "fixed" else case.vuln_file
                    data = file.read_text()
                    if is_catalog_excerpt(data):
                        data = parseable_excerpt(file, data)
                        mode = "excerpt-parseable"
                    else:
                        # Validate BOTH twins before treating this as a complete-file pair.
                        for twin in (case.vuln_file, case.fixed_file):
                            verify_catalog_file(twin)
                        mode = "catalog-file"
                    catalog_files.append((case, side, mode, data))
            if catalog_files:
                if identities:
                    raise ValueError("mixed source and catalog-file materialization at one pair pin")
                for case, side, mode, data in catalog_files:
                    _materialize_side(target, case, side=side)
                    if mode == "excerpt-parseable":
                        safe_path(target, case.relpath).write_text(data)
                verify_functions(target, rows)
                modes = {case.relpath: mode for case, _, mode, _ in catalog_files}
                for row in rows:
                    if modes[row["file"]] == "excerpt-parseable":
                        row["evidence_scope"] = "in-function-only"
                return {
                    "mode": "excerpt-parseable" if "excerpt-parseable" in modes.values() else "catalog-file",
                    "files": sum(p.is_file() for p in target.rglob("*")),
                    "bytes": sum(p.stat().st_size for p in target.rglob("*") if p.is_file()),
                    "roots": sorted({source_root(r["file"], r["language"]) for r in rows}),
                }
            if len(identities) != 1:
                raise ValueError("ambiguous source commits at one pair pin")
            repo, commit = identities.pop()
            roots = sorted({source_root(row["file"], row["language"]) for row in rows})
            try:
                clone = self.fetch_source(repo, commit, roots)
                proof = export_sources(clone, target, rows, roots)
            except (OSError, ValueError, subprocess.SubprocessError) as exc:
                if not self.allow_excerpt_fallback:
                    raise
                shutil.rmtree(target, ignore_errors=True)
                for case, side in resolved:
                    _materialize_side(target, case, side=side)
                return {"mode": "excerpt-fallback", "reason": str(exc), "repo": repo, "commit": commit}
            verify_functions(target, rows)
            return {**proof, "mode": "source", "repo": repo, "commit": commit}
        if all(r["source"] == "dev-php" for r in rows):
            # Export the exact requested commit from the local recipe checkout.
            for recipe in (ROOT / "benchmarks/repos").glob("*.toml"):
                data = tomllib.loads(recipe.read_text())
                if repo_name(str(data.get("url", ""))) == repo_name(pin["repo"]):
                    checkout = self.cache / "repos" / data["name"] / pin["pin"][:12]
                    return self.export_local(checkout, pin, target)
        for clone, _ in self.repos[repo_name(pin["repo"])]:
            check = subprocess.run(["git", "-C", str(clone), "cat-file", "-e", pin["pin"] + "^{commit}"], capture_output=True)
            if check.returncode == 0:
                return self.export_local(clone, pin, target)
        raise ValueError(f"no local clone contains {pin['repo']}@{pin['pin']}")


def is_catalog_excerpt(data):
    # Older Python exports have this header but no upstream_start field.
    return re.match(r"\s*(?://|#)\s*Provenance:", data) is not None


def parseable_excerpt(path, data):
    """Dedent the code independently of its unindented provenance comments."""
    language = detect_language(path)
    lines = data.splitlines(keepends=True)
    start = 0
    while start < len(lines) and (not lines[start].strip() or lines[start].lstrip().startswith(("#", "//"))):
        start += 1
    header = "".join(lines[:start])
    if language == "python":
        header = re.sub(r"(?m)^(\s*)//", r"\1#", header)
    data = header + textwrap.dedent("".join(lines[start:]))
    if not "".join(lines[start:]).strip():
        raise ValueError(f"excerpt does not parse: empty code ({path.name})")
    if language == "python":
        try:
            ast.parse(data, filename=str(path))
        except SyntaxError as exc:
            raise ValueError(f"excerpt does not parse: {exc}") from exc
    else:
        from openultrasast.semantic.cst import parse_with_cst

        parsed = parse_with_cst(str(path), data, language)
        if parsed is None:
            raise ValueError(f"excerpt does not parse: no cheap parser for {language} ({path.name})")
        if not parsed.parse_ok:
            raise ValueError(f"excerpt does not parse: {language} {parsed.reason} ({path.name})")
    return data


def verify_catalog_file(path):
    """Admit complete source only; syntactically valid provenance excerpts still fail."""
    data = path.read_text()
    if is_catalog_excerpt(data):
        raise ValueError(f"catalog side is an upstream excerpt without distinct commits: {path.name}")
    language = detect_language(path)
    if not data.strip():
        raise ValueError(f"catalog side is empty: {path.name}")
    if language == "python":
        try:
            ast.parse(data, filename=str(path))
        except SyntaxError as exc:
            raise ValueError(f"catalog side is not complete Python source: {path.name}: {exc.msg}") from exc
    else:
        from openultrasast.semantic.extra import grammar_for

        grammar = grammar_for(language)
        if grammar is None:
            raise ValueError(f"cannot verify complete catalog source for {language}: {path.name}")
        from tree_sitter import Parser

        if Parser(grammar).parse(data.encode()).root_node.has_error:
            raise ValueError(f"catalog side is not complete {language} source: {path.name}")


def source_root(file, language):
    parts = Path(file).parts
    if language == "java":
        for i in range(len(parts) - 2):
            if parts[i : i + 3] == ("src", "main", "java"):
                return "/".join(parts[: i + 3])
    # Monorepo and source-layout containers are not themselves a package.
    if len(parts) > 2 and parts[0] in {"packages", "apps", "src", "web"}:
        return "/".join(parts[:2])
    return parts[0] if len(parts) > 1 else "."


def local_imports(source, rows):
    """Cheap one-level Python imports from labelled files, without importing code."""
    found = set()
    for row in rows:
        if row["language"] != "python":
            continue
        file = safe_path(source, row["file"])
        try:
            tree = ast.parse(file.read_bytes())
        except (SyntaxError, ValueError):
            continue
        for node in ast.walk(tree):
            if isinstance(node, ast.Import):
                modules, bases = [alias.name for alias in node.names], [source, source / "src"]
            elif isinstance(node, ast.ImportFrom):
                modules = [node.module or ""]
                modules += [".".join(filter(None, (node.module, alias.name))) for alias in node.names if alias.name != "*"]
                bases = [file.parents[node.level - 1]] if node.level and node.level <= len(file.parents) else [source, source / "src"]
            else:
                continue
            for base in bases:
                for module in modules:
                    path = base.joinpath(*module.split("."))
                    for candidate in (path.with_suffix(".py"), path / "__init__.py"):
                        if candidate.is_relative_to(source) and candidate.is_file() and not candidate.is_symlink():
                            rel = candidate.relative_to(source).as_posix()
                            if not any(part in EXCLUDED_DIRS for part in Path(rel).parts):
                                found.add(rel)
    return found


def export_sources(source, target, rows, roots=(".",), cap=EXPORT_CAP):
    """Copy labels intact, narrowing oversized roots before bounding context bytes."""
    labelled = {r["file"] for r in rows}
    languages = {r["language"] for r in rows}
    files = set()
    protected = {parent.as_posix() for rel in labelled for parent in Path(rel).parents}
    for directory, dirs, names in os.walk(source):
        dirs[:] = sorted(
            d
            for d in dirs
            if d != ".git"
            and not (Path(directory) / d).is_symlink()
            and (d not in EXCLUDED_DIRS or (Path(directory) / d).relative_to(source).as_posix() in protected)
        )
        for name in names:
            file = Path(directory) / name
            rel = file.relative_to(source).as_posix()
            excluded = any(part in EXCLUDED_DIRS for part in Path(rel).parts[:-1])
            if not file.is_symlink() and (
                rel in labelled
                or (
                    not excluded
                    and detect_language(file) in languages
                    and any(root == "." or rel == root or rel.startswith(root + "/") for root in roots)
                )
            ):
                files.add(rel)
    missing = labelled - files
    if missing:
        raise ValueError(f"labelled files missing or excluded from export: {sorted(missing)}")
    imports = local_imports(source, rows)
    files.update(imports)
    sizes = {rel: safe_path(source, rel).stat().st_size for rel in files}
    original_roots = sorted(roots)
    roots = original_roots
    # Avoid lexicographically truncating a whole large package such as vllm.
    if sum(sizes.values()) > cap:
        roots = sorted({str(Path(rel).parent) for rel in labelled})
        files = {
            rel for rel in files if rel in labelled or rel in imports or any(root == "." or rel.startswith(root + "/") for root in roots)
        }
    proof = {
        "bytes": 0,
        "files": 0,
        "cap_bytes": cap,
        "cap_hit": False,
        "excludes": sorted(EXCLUDED_DIRS),
        "roots": roots,
        "requested_roots": original_roots,
        "scope_narrowed": roots != original_roots,
        "import_files": sorted(imports),
    }
    for rel in sorted(files, key=lambda rel: (rel not in labelled, rel)):
        file = safe_path(source, rel)
        size = sizes[rel]
        # The cap bounds context only. No labelled source is ever cut or omitted.
        if proof["bytes"] + size > cap:
            proof["cap_hit"] = True
            if rel not in labelled:
                continue
        dest = safe_path(target, rel)
        dest.parent.mkdir(parents=True, exist_ok=True)
        dest.write_bytes(file.read_bytes())
        proof["bytes"] += size
        proof["files"] += 1
    return proof


def verify_functions(root, rows):
    for row in rows:
        data = safe_path(root, row["file"]).read_text(errors="replace")
        if not re.search(r"(?<![\w$])" + re.escape(row["function"]) + r"(?![\w$])", data):
            raise ValueError(f"labelled function not found in fetched source: {row['file']}::{row['function']}")


def freeze_source(target):
    target.mkdir(parents=True)
    archive = target.parent / "source.tar"
    with archive.open("wb") as stream:
        subprocess.run(["git", "-C", str(ROOT), "archive", "HEAD", "src"], stdout=stream, check=True)
    subprocess.run(["tar", "-xf", str(archive), "-C", str(target)], check=True)
    archive.unlink()
    overlay = {}
    for name in SNAPSHOT_FILES:
        data = (ROOT / name).read_bytes()
        dest = target / name
        dest.parent.mkdir(parents=True, exist_ok=True)
        dest.write_bytes(data)
        overlay[name] = hashlib.sha256(data).hexdigest()
    head = subprocess.run(["git", "-C", str(ROOT), "rev-parse", "HEAD"], capture_output=True, text=True, check=True).stdout.strip()
    return {"head": head, "overlay_sha256": overlay}


def docker_command(source, checkout, out, name, image, deadline, question_deadline):
    return [
        "docker",
        "run",
        "--rm",
        "--init",
        "--name",
        name,
        "--network",
        "none",
        "--memory",
        "3g",
        "--pull",
        "never",
        "--entrypoint",
        "python",
        "-e",
        "PYTHONPATH=/frozen/src",
        "-v",
        f"{source}:/frozen:ro",
        "-v",
        f"{checkout}:/case:ro",
        "-v",
        f"{out}:/out",
        image,
        "/frozen/benchmarks/learn/engine_trace_worker.py",
        "--worker",
        "/out/pin.json",
        "--deadline",
        str(deadline),
        "--question-deadline",
        str(question_deadline),
    ]


def prepare_questions(pin, root, question_deadline):
    """Prepare metadata on the host, before entering the dependency-minimal image."""
    active = [u for u in pin["units"] if u["supported"]]
    from openultrasast.learn.features import entry_names
    from openultrasast.mapping import php_hook_callbacks
    from openultrasast.plane.harvest import declares
    from openultrasast.preprocess import preprocess_repository

    _, targets = preprocess_repository(root)
    entries, _ = entry_names(root, {u["file"] for u in active})
    hooks = ";".join(f"{h}:{f}" for h, functions in php_hook_callbacks(root, targets).items() for f in functions)
    requests = {}
    for u in active:
        specs = taint_specs(language="javascript" if u["language"] == "typescript" else u["language"])
        if not declares(safe_path(root, u["file"]).read_text(errors="replace").splitlines(), u["file"], u["function"]):
            raise ValueError(f"labelled function not declared: {u['file']}::{u['function']}")
        excerpt = u.get("evidence_scope") == "in-function-only" or pin.get("materialization", {}).get("mode") == "excerpt-fallback"
        params = request_params(
            specs[u["family"]],
            file=u["file"],
            function=u["function"],
            trace=True,
            parameter_sources=excerpt or u["function"] in entries,
            call_depth=0 if excerpt else 3,
            hook_callbacks=hooks,
        )
        params["questionDeadline"] = str(question_deadline)
        requests[u["unit"]] = params
    return requests


def run_pin(pin, source, provenance, inputs, out, args):
    with tempfile.TemporaryDirectory(prefix="trace-pin-", dir=out) as scratch:
        scratch = Path(scratch)
        checkout, output = scratch / "case", scratch / "output"
        output.mkdir()
        try:
            pin = {**pin, "materialization": inputs.materialize(pin, checkout)}
            write_json(output / "pin.json", {**pin, "questions": prepare_questions(pin, checkout, args.question_deadline)})
            name = "ousast-trace-" + pin_name(pin)[:16]
            command = docker_command(source, checkout, output, name, args.image, args.deadline, args.question_deadline)
            started = time.monotonic()
            done = None
            try:
                done = subprocess.run(command, capture_output=True, text=True, timeout=args.deadline)
                status, reason = "failed", f"container exit {done.returncode}: {done.stderr[-1000:]}"
                (output / "container.log").write_text(done.stdout + "\n" + done.stderr)
            except subprocess.TimeoutExpired:
                done = None
                status, reason = "timeout", f"container deadline {args.deadline}s"
            finally:
                # Never start the next pin unless this container is known to be gone.
                try:
                    removed = subprocess.run(["docker", "rm", "-f", name], capture_output=True, timeout=20)
                except subprocess.SubprocessError as exc:
                    raise RuntimeError(f"cannot confirm container {name} stopped; refusing another launch") from exc
                if done is None and removed.returncode != 0:
                    raise RuntimeError(f"cannot remove timed-out container {name}; refusing another launch")
            result_path = output / "result.json"
            if not result_path.exists():
                reason = f"container produced no result: {reason}"
            record = json.loads(result_path.read_text()) if result_path.exists() else {**pin, "units": []}
            if not record.get("done"):
                previous = {u["unit"]: u for u in record["units"]}
                record["units"] = []
                for unit in pin["units"]:
                    row = previous.get(unit["unit"], unit_record(unit, status))
                    if row["status"] not in {"path", "asked-nothing"}:
                        row.update(
                            status=status if unit["supported"] else "unsupported", reason=reason if unit["supported"] else unit["reason"]
                        )
                    record["units"].append(row)
            record["container_exit"] = done.returncode if done else None
            if done is not None and done.returncode != 0:
                for row in record["units"]:
                    if row["supported"]:
                        row.update(status="failed", reason=reason)
            record.update(done=True, seconds=time.monotonic() - started, analyzer=provenance, image=args.image)
            if (output / "container.log").exists():
                shutil.copyfile(output / "container.log", out / pin_name(pin).replace(".json", ".log"))
        except (OSError, ValueError, subprocess.SubprocessError) as exc:
            record = {
                **pin,
                "done": True,
                "analyzer": provenance,
                "units": [unit_record(u, "failed" if u["supported"] else "unsupported", str(exc)) for u in pin["units"]],
            }
        save_pin(out / pin_name(pin), record)
        print(
            json.dumps({"repo": pin["repo"], "pin": pin["pin"], "statuses": dict(Counter(u["status"] for u in record["units"]))}),
            flush=True,
        )
        return any(u["status"] in {"failed", "timeout"} for u in record["units"])


def save_pin(path, record):
    if path.exists():
        old = json.loads(path.read_text())
        selected = {u["unit"] for u in record["units"]}
        record["units"].extend(u for u in old.get("units", []) if u["unit"] not in selected)
    write_json(path, record)


def summary(out):
    if not out.is_dir():
        raise ValueError(f"results root does not exist: {out}")
    family, source = defaultdict(Counter), defaultdict(Counter)
    pairs = defaultdict(dict)
    units = {}
    result_files = 0
    for path in sorted(out.glob("*.json")):
        record = json.loads(path.read_text())
        if not record.get("done"):
            continue
        result_files += 1
        for unit in record["units"]:
            if unit["unit"] in units:
                raise ValueError(f"duplicate result for {unit['unit']}")
            units[unit["unit"]] = unit
            family[unit["family"]][unit["status"]] += 1
            source[unit["source"]][unit["status"]] += 1
            if unit["pair"]:
                pairs[unit["family"], unit["pair"]][unit["label"]] = unit["status"]
    asymmetry = Counter()
    completed = Counter()
    for sides in pairs.values():
        if set(sides) != {0, 1}:
            asymmetry["incomplete-pair"] += 1
            continue
        key = {(True, False): "vulnerable-only", (False, True): "fixed-only", (True, True): "both", (False, False): "neither"}[
            sides[1] == "path", sides[0] == "path"
        ]
        asymmetry[key] += 1
        if all(s in {"path", "asked-nothing"} for s in sides.values()):
            completed[key] += 1
    return {
        "result_files_read": result_files,
        "units": len(units),
        "by_family": {k: {s: v[s] for s in STATUSES} for k, v in sorted(family.items())},
        "by_source": {k: {s: v[s] for s in STATUSES} for k, v in sorted(source.items())},
        "pair_asymmetry": dict(asymmetry),
        "pair_asymmetry_both_answered": dict(completed),
    }


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--units", type=Path, default=ROOT / "plane/experiments/exp-005-graph-slice.units.jsonl")
    parser.add_argument("--examples", type=Path, help="frozen example JSONL instead of the local memory store")
    parser.add_argument("--memory", type=Path, default=Path.home() / "ousast-results/plane/memory")
    parser.add_argument("--cache", type=Path, default=Path.home() / ".cache/openultrasast")
    parser.add_argument("--out", type=Path, default=DEFAULT_OUT)
    parser.add_argument("--image", default="openultrasast:dev")
    parser.add_argument("--deadline", type=int, default=900)
    parser.add_argument("--question-deadline", type=int, default=120)
    parser.add_argument("--limit", type=int, default=0)
    parser.add_argument("--only-family")
    parser.add_argument("--only-source", default="", help="comma-separated source catalog IDs")
    parser.add_argument("--rerun-status", default="failed,timeout", help="comma-separated statuses to re-run (including timeout)")
    parser.add_argument(
        "--allow-excerpt-fallback", action="store_true", help="allow fragments only if fetching/exporting resolved source fails"
    )
    parser.add_argument("--include-typescript", action="store_true")
    parser.add_argument("--dry-run", action="store_true")
    parser.add_argument("--summary", action="store_true")
    args = parser.parse_args()
    if min(args.deadline, args.question_deadline) <= 0 or args.limit < 0:
        parser.error("deadlines must be positive; limit must be nonnegative")
    rerun_status = set(filter(None, args.rerun_status.split(",")))
    if rerun_status - set(STATUSES):
        parser.error("unknown --rerun-status value")
    if args.summary:
        print(json.dumps(summary(args.out), indent=2))
        return 0
    # Even inherited opt-ins must not make catalog resolution fetch.
    os.environ["GIT_NO_LAZY_FETCH"] = "1"
    os.environ.pop("OPENULTRASAST_PAIRS_NETWORK", None)
    os.environ.pop("OPENULTRASAST_REPOS_NETWORK", None)
    examples = read_jsonl(args.examples) if args.examples else [r.row for r in FileStore(args.memory).rows(kind="example")]
    try:
        units = join_units(read_jsonl(args.units), examples, args.include_typescript)
    except ValueError as exc:
        parser.error(str(exc))
    inputs = Inputs(ROOT, args.cache, args.allow_excerpt_fallback)
    # Resolve pair source names for leak-audit coverage before calculating the resume key.
    for unit in units:
        if unit.get("excerpt") or unit["source"].startswith("advisory-fixes"):
            try:
                case, _ = inputs.pair_for(unit)
                catalog_source, _ = inputs.catalog_for(case)
                unit["source"] = catalog_source if catalog_source.startswith("advisory-fixes") else case.slice
                unit["excerpt"] = True
            except ValueError as exc:
                unit["mapping_error"] = str(exc)
    plan = make_plan(
        units,
        args.out,
        args.only_family,
        args.limit,
        include_typescript=args.include_typescript,
        only_source=set(filter(None, args.only_source.split(","))),
        rerun_status=rerun_status,
        allow_excerpt_fallback=args.allow_excerpt_fallback,
    )
    # Planning is read-only, including unsupported pins. All persistence stays below this return.
    if args.dry_run:
        print(
            json.dumps(
                {
                    "pins": len(plan),
                    "runnable_pins": sum(any(u["supported"] for u in p["units"]) for p in plan),
                    "units": sum(len(p["units"]) for p in plan),
                    "estimated_seconds": sum(p["estimated_seconds"] for p in plan),
                    "estimate_basis": "inventory medians: excerpt 31/47s, repository 182/498s; not a deadline guarantee",
                    "plan": plan,
                },
                indent=2,
            )
        )
        return 0
    args.out.mkdir(parents=True, exist_ok=True)
    lock = (args.out / ".lock").open("w")
    try:
        fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
    except BlockingIOError:
        parser.error("another trace runner holds this results root")
    failures = False
    if any(u["supported"] for pin in plan for u in pin["units"]):
        image = subprocess.run(
            ["docker", "image", "inspect", "--format", "{{.Id}}", args.image], capture_output=True, text=True, check=True
        )
        args.image = image.stdout.strip()
        if not args.image:
            raise ValueError("Docker returned no local image identity")
    with tempfile.TemporaryDirectory(prefix="trace-source-", dir=args.out) as scratch:
        source = Path(scratch) / "frozen"
        provenance = freeze_source(source)
        progress = {"pins_done": 0, "pins_total": len(plan), "last_pin": None, "statuses": {}}
        counts = Counter()
        write_json(args.out / "progress.json", progress)
        for pin in plan:
            if (args.out / "STOP").exists():
                break
            if not any(u["supported"] for u in pin["units"]):
                save_pin(
                    args.out / pin_name(pin),
                    {**pin, "done": True, "units": [unit_record(u, "unsupported", u["reason"]) for u in pin["units"]]},
                )
            else:
                failures |= run_pin(pin, source, provenance, inputs, args.out, args)
            result = json.loads((args.out / pin_name(pin)).read_text())
            selected = {u["unit"] for u in pin["units"]}
            counts.update(u["status"] for u in result["units"] if u["unit"] in selected)
            progress.update(pins_done=progress["pins_done"] + 1, last_pin={"repo": pin["repo"], "pin": pin["pin"]}, statuses=dict(counts))
            write_json(args.out / "progress.json", progress)
    return 2 if failures else 0


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