#!/usr/bin/env python3
"""Run every analysis script in this package and, optionally, check that the
results match the ones shipped with it.

Why this file exists:
    Each script under code/<name>/ runs on its own, but reproducing the paper
    means running all of them. This runner does that in the order listed in
    manifest.py, runs each script in its own process so one failure does not
    stop the rest, and prints one line per script.

    With --check it also compares each fresh output with the copy that shipped
    in the package, ignoring the "generated" timestamp and floating-point
    noise. Running a script overwrites its shipped output, so --check first
    confirms each shipped output is untouched, using the checksums recorded
    in shipped_outputs.sha256 when the package was built. A script whose
    shipped output is missing or already overwritten is reported as not
    verified (re-extract the package) instead of being compared. When a
    script reproduces, its shipped files are restored byte-for-byte, so
    --check can be run again on the same copy.

Usage:
    python code/run_all.py                    # run everything
    python code/run_all.py --check            # run everything and verify
    python code/run_all.py --only main_spec   # run scripts whose path matches

Upstream dependencies:
    manifest.py (script list, outputs, run order).
    shipped_outputs.sha256 (written when the package is built; absent in a
    development checkout, where --check compares against the files on disk).
"""

from __future__ import annotations

import argparse
import hashlib
import json
import math
import subprocess
import sys
import time
from typing import Any

from manifest import CODE_DIR, SCRIPTS, Script

# Keys whose values change on every run by design.
VOLATILE_KEYS = {"generated"}

# SHA-256 of every shipped output, one "<digest>  <path under code/>" per line.
CHECKSUMS = CODE_DIR / "shipped_outputs.sha256"


def shipped_checksums() -> dict[str, str] | None:
    """Return {path under code/: sha256}, or None when no checksum file exists."""
    if not CHECKSUMS.exists():
        return None
    checksums = {}
    for line in CHECKSUMS.read_text(encoding="utf-8").splitlines():
        if line.strip():
            digest, rel = line.split(maxsplit=1)
            checksums[rel] = digest
    return checksums


def tampered_outputs(script: Script, checksums: dict[str, str]) -> list[str]:
    """Describe each shipped output of `script` that is missing or no longer as shipped."""
    problems = []
    for name in script.outputs:
        path = script.output_dir / name
        rel = path.relative_to(CODE_DIR).as_posix()
        if rel not in checksums:
            problems.append(f"{name}: not listed in {CHECKSUMS.name}")
        elif not path.exists():
            problems.append(f"{name}: shipped copy is missing")
        elif hashlib.sha256(path.read_bytes()).hexdigest() != checksums[rel]:
            problems.append(f"{name}: shipped copy was modified (overwritten by an earlier run?)")
    return problems


def differences(expected: Any, actual: Any, path: str = "") -> list[str]:
    """List where two parsed JSON values differ beyond floating-point noise."""
    if isinstance(expected, dict) and isinstance(actual, dict):
        found: list[str] = []
        for key in sorted(set(expected) | set(actual)):
            if key in VOLATILE_KEYS:
                continue
            if key not in expected or key not in actual:
                found.append(f"{path}.{key}: present in only one version")
                continue
            found += differences(expected[key], actual[key], f"{path}.{key}")
        return found
    if isinstance(expected, list) and isinstance(actual, list):
        if len(expected) != len(actual):
            return [f"{path}: expected {len(expected)} items, got {len(actual)}"]
        found = []
        for i, (e, a) in enumerate(zip(expected, actual)):
            found += differences(e, a, f"{path}[{i}]")
        return found
    if (isinstance(expected, (int, float)) and isinstance(actual, (int, float))
            and not isinstance(expected, bool) and not isinstance(actual, bool)):
        if math.isclose(expected, actual, rel_tol=1e-9, abs_tol=1e-12):
            return []
        return [f"{path}: expected {expected}, got {actual}"]
    return [] if expected == actual else [f"{path}: expected {expected!r}, got {actual!r}"]


def read_outputs(script: Script) -> dict[str, Any]:
    """Parse the script's output files that currently exist."""
    return {
        name: json.loads((script.output_dir / name).read_text(encoding="utf-8"))
        for name in script.outputs
        if (script.output_dir / name).exists()
    }


def main() -> int:
    parser = argparse.ArgumentParser(description="Run the replication package.")
    parser.add_argument("--check", action="store_true",
                        help="compare fresh outputs with the shipped ones")
    parser.add_argument("--only", nargs="+", metavar="NAME",
                        help="run only scripts whose path contains NAME")
    args = parser.parse_args()

    scripts = [s for s in SCRIPTS
               if not args.only or any(name in s.entry for name in args.only)]
    if not scripts:
        print(f"No scripts match --only {args.only}", file=sys.stderr)
        return 2

    checksums = shipped_checksums() if args.check else None
    if args.check and checksums is None:
        print(f"Note: no {CHECKSUMS.name}; comparing against the output files on disk, "
              "which an earlier run may have overwritten.\n")

    failed: list[str] = []
    mismatched: list[str] = []
    for script in scripts:

        # Show the shared panel numbers beside the script that reproduces them.
        panel_label = f" [{', '.join(script.panels)}]" if script.panels else ""
        print(f"{script.entry + panel_label:<78} ", end="", flush=True)
        if checksums is not None:
            tampered = tampered_outputs(script, checksums)
            if tampered:
                print("NOT VERIFIED")
                for problem in tampered:
                    print(f"    {problem}")
                print("    Re-extract the package to restore the shipped outputs, then re-run.")
                mismatched.append(script.entry)
                continue
        expected = read_outputs(script) if args.check else {}
        shipped_bytes = {name: (script.output_dir / name).read_bytes()
                         for name in expected} if args.check else {}
        start = time.perf_counter()
        proc = subprocess.run([sys.executable, str(script.path)],
                              cwd=script.path.parent, capture_output=True, text=True)
        seconds = time.perf_counter() - start
        if proc.returncode != 0:
            print(f"FAILED ({seconds:.1f}s)")
            detail = proc.stderr.strip() or proc.stdout.strip() or "(no output)"
            print("    " + detail.replace("\n", "\n    "))
            failed.append(script.entry)
            continue
        if not args.check:
            print(f"ok ({seconds:.1f}s)")
            continue
        actual = read_outputs(script)
        problems = [f"{name}: no shipped copy to compare against"
                    for name in script.outputs if name not in expected]
        problems += [f"{name}{d}" for name in expected
                     for d in differences(expected[name], actual.get(name))]
        if problems:
            print(f"DIFFERS ({seconds:.1f}s)")
            for problem in problems[:5]:
                print(f"    {problem}")
            mismatched.append(script.entry)
        else:
            for name, content in shipped_bytes.items():
                (script.output_dir / name).write_bytes(content)
            print(f"reproduced ({seconds:.1f}s)")

    print()
    if failed:
        print(f"{len(failed)} script(s) failed: {', '.join(failed)}")
    if mismatched:
        print(f"{len(mismatched)} script(s) not verified: {', '.join(mismatched)}")
    if failed or mismatched:
        return 1
    verb = "reproduced the shipped results" if args.check else "ran successfully"
    print(f"All {len(scripts)} script(s) {verb}.")
    return 0


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