#!/usr/bin/env python3

"""Build a repro bundle for a published drop.

A repro bundle pins everything a third party needs to reproduce a run:

suite revision, scenario/config hashes, fixture seeds + generator versions,

subject binary SHA256s, runner profile, noise floor, and raw-JSONL hashes

(plus optional R2 storage keys). The bundle carries hashes of the raw rows,

not the rows themselves.

Usage:

  python3 scripts/repro_bundle.py results/rows.jsonl ... \

      --out bundle.json --suite-revision <oak-hash> \

      [--noise-floor-pct 4.2] [--r2-key raw/rows.jsonl ...]

Null means unmeasured (ADR-0002): an absent noise floor is recorded as null,

never 0.

"""

from __future__ import annotations

import argparse

import json

import sys

from datetime import datetime, timezone

from pathlib import Path

from typing import Any

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

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

from oakbench.baseline_book import sha256_file  # noqa: E402

from oakbench.fixture_registry import load_fixture_registry  # noqa: E402

from oakbench.runner import machine_profile, runner_class, runner_id  # noqa: E402

SCHEMA_VERSION = 1

CONTENT_INTEGRITY_SOURCE_STRENGTH = {

    "git_head_payload_sha256": 3,

    "oak_independent_reader_payload_sha256": 3,

    "git_lfs_worktree_payload_sha256": 1,

    "worktree_payload_sha256": 1,

    "mixed_payload_sources": 1,

}

def hash_directory(directory: Path, pattern: str, root: Path) -> dict[str, str]:

    """SHA256 every file matching ``pattern`` under ``directory``."""

    hashes: dict[str, str] = {}

    if directory.is_dir():

        for path in sorted(directory.glob(pattern)):

            if path.is_file():

                hashes[path.relative_to(root).as_posix()] = sha256_file(path)

    return hashes

def count_jsonl_rows(path: Path) -> int:

    count = 0

    with path.open() as fh:

        for line in fh:

            if line.strip():

                count += 1

    return count

def iter_jsonl_rows(paths: list[Path]) -> list[dict[str, Any]]:

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

    for path in paths:

        with path.open() as fh:

            for raw_line in fh:

                line = raw_line.strip()

                if line:

                    rows.append(json.loads(line))

    return rows

def subject_binaries_from_rows(rows: list[dict[str, Any]]) -> dict[str, Any]:

    """Map subject -> binary SHA256 collected from rows' binary_sha256 fields.

    A subject that recorded more than one distinct binary hash keeps the full

    sorted list — silently picking one would hide a mixed-binary run.

    """

    found: dict[str, set[str]] = {}

    for row in rows:

        subject = row.get("subject")

        sha = row.get("binary_sha256")

        if isinstance(subject, str) and subject and isinstance(sha, str) and sha:

            found.setdefault(subject, set()).add(sha)

    result: dict[str, Any] = {}

    for subject in sorted(found):

        ordered = sorted(found[subject])

        result[subject] = ordered[0] if len(ordered) == 1 else ordered

    return result

def content_integrity_summary_from_rows(rows: list[dict[str, Any]]) -> dict[str, Any]:

    """Summarize content-attestation strength by comparable row group.

    The boolean content_integrity_check_passed says only whether the row's

    chosen oracle passed. Publication/replay decisions also need the source

    strength: HEAD-payload attestation is stronger than worktree/LFS-worktree

    attestation, and mixed groups must be treated as their weakest source.

    """

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

    for row in rows:

        source = row.get("content_integrity_source")

        sources = row.get("content_integrity_sources")

        if not isinstance(source, str) and not isinstance(sources, list):

            continue

        profile = row.get("profile")

        key = "/".join(

            [

                "-" if profile is None else str(profile),

                str(row.get("scenario") or ""),

                str(row.get("operation") or ""),

                str(row.get("subject") or ""),

            ]

        )

        entry = grouped.setdefault(

            key,

            {

                "sources": set(),

                "rows": 0,

                "passed": 0,

                "failed": 0,

                "weakest_source_strength": None,

            },

        )

        row_sources = [value for value in ([source] + sources if isinstance(sources, list) else [source]) if isinstance(value, str)]

        for value in row_sources:

            entry["sources"].add(value)

            strength = CONTENT_INTEGRITY_SOURCE_STRENGTH.get(value, 0)

            current = entry["weakest_source_strength"]

            entry["weakest_source_strength"] = strength if current is None else min(current, strength)

        entry["rows"] += 1

        if row.get("content_integrity_check_passed") is True:

            entry["passed"] += 1

        elif row.get("content_integrity_check_passed") is False:

            entry["failed"] += 1

    return {

        key: {

            "sources": sorted(value["sources"]),

            "rows": value["rows"],

            "passed": value["passed"],

            "failed": value["failed"],

            "weakest_source_strength": value["weakest_source_strength"],

        }

        for key, value in sorted(grouped.items())

    }

def build_bundle(

    row_paths: list[Path],

    *,

    suite_revision: str,

    noise_floor_pct: float | None = None,

    r2_keys: list[str] | None = None,

    root: Path = ROOT,

) -> dict[str, Any]:

    profile = machine_profile()

    registry = load_fixture_registry(root / "config" / "fixtures.json")

    rows = iter_jsonl_rows(row_paths)

    return {

        "schema_version": SCHEMA_VERSION,

        "created_utc": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),

        "suite_revision": suite_revision,

        "runner_profile": profile,

        "runner_id": runner_id(profile),

        "runner_class": runner_class(profile),

        "noise_floor_pct": noise_floor_pct,

        "config_hashes": {

            **hash_directory(root / "config", "**/*.json", root),

            **hash_directory(root / "config", "**/*.toml", root),

        },

        "scenario_hashes": hash_directory(root / "scenarios", "*.yaml", root),

        "fixtures": [

            {

                "id": spec.fixture_id,

                "seed": spec.seed,

                "generator": spec.generator,

                "version": spec.version,

            }

            for _, spec in sorted(registry.items())

        ],

        "subject_binaries": subject_binaries_from_rows(rows),

        "content_integrity": {

            "source_strength": dict(sorted(CONTENT_INTEGRITY_SOURCE_STRENGTH.items())),

            "groups": content_integrity_summary_from_rows(rows),

        },

        "raw_inputs": [

            {

                "path": str(path),

                "sha256": sha256_file(path),

                "bytes": path.stat().st_size,

                "rows": count_jsonl_rows(path),

            }

            for path in row_paths

        ],

        "r2_keys": list(r2_keys or []),

    }

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

    parser = argparse.ArgumentParser(description="Build a repro bundle for a published drop.")

    parser.add_argument("rows", nargs="+", type=Path, help="raw row JSONL paths covered by the drop")

    parser.add_argument("--out", type=Path, required=True, help="output bundle JSON path")

    parser.add_argument(

        "--suite-revision",

        required=True,

        help="caller-supplied suite revision (e.g. the oak hash of the harness checkout)",

    )

    parser.add_argument(

        "--noise-floor-pct",

        type=float,

        default=None,

        help="measured same-binary noise floor percentage; omitted means unmeasured (null)",

    )

    parser.add_argument(

        "--r2-key",

        action="append",

        default=None,

        dest="r2_keys",

        help="R2 storage key for an uploaded raw artifact (repeatable)",

    )

    return parser.parse_args(argv)

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

    args = parse_args(argv)

    missing = [str(path) for path in args.rows if not path.is_file()]

    if missing:

        print(f"error: row file(s) not found: {', '.join(missing)}", file=sys.stderr)

        return 2

    bundle = build_bundle(

        args.rows,

        suite_revision=args.suite_revision,

        noise_floor_pct=args.noise_floor_pct,

        r2_keys=args.r2_keys,

    )

    args.out.parent.mkdir(parents=True, exist_ok=True)

    args.out.write_text(json.dumps(bundle, indent=2, sort_keys=True) + "\n")

    print(

        f"wrote {args.out}: {len(bundle['raw_inputs'])} raw input(s), "

        f"{len(bundle['config_hashes'])} config hash(es), "

        f"{len(bundle['scenario_hashes'])} scenario hash(es), "

        f"{len(bundle['fixtures'])} fixture(s)"

    )

    return 0

if __name__ == "__main__":

    sys.exit(main())