mech_workbench/bundle.py

"""Bundle verification for completed run folders."""
from __future__ import annotations

import csv
import hashlib
import json
import zipfile
from pathlib import Path

REQUIRED_FILES = [
    "RUN_CARD.json",
    "config/resolved_config.json",
    "logs/events.jsonl",
]

DATA_CSVS = [
    "data/head_recovery.csv",
    "data/logit_lens.csv",
    "data/gradient_attribution.csv",
]

OPTIONAL_CSVS = [
    "data/ablation_results.csv",
    "data/source_results.csv",
    "summaries/candidate_heads.csv",
    "summaries/ablation_summary.csv",
    "summaries/source_summary.csv",
]

SUMMARY_FILES = [
    "summaries/run_summary.json",
]

def verify_bundle(run_dir: Path) -> bool:
    """Verify a completed run bundle. Prints pass/fail for each check. Returns True if all pass."""
    run_dir = Path(run_dir)
    print(f"\n=== Verifying: {run_dir} ===")
    ok = True

    # 1. Directory exists
    if not run_dir.is_dir():
        print(f"  [FAIL] Directory not found: {run_dir}")
        return False

    # 2. Required files
    for rel in REQUIRED_FILES:
        path = run_dir / rel
        if path.exists():
            print(f"  [OK]   {rel}")
        else:
            print(f"  [FAIL] Missing: {rel}")
            ok = False

    # 3. At least one summary JSON
    for rel in SUMMARY_FILES:
        path = run_dir / rel
        if path.exists():
            try:
                json.loads(path.read_text(encoding="utf-8"))
                print(f"  [OK]   {rel} (valid JSON)")
            except Exception as exc:
                print(f"  [FAIL] {rel} - JSON parse error: {exc}")
                ok = False
        else:
            print(f"  [WARN] Missing (optional): {rel}")

    # 4. Raw JSON files
    raw_dir = run_dir / "raw"
    if raw_dir.is_dir():
        raw_jsons = list(raw_dir.glob("*.json"))
        if raw_jsons:
            for f in raw_jsons:
                try:
                    json.loads(f.read_text(encoding="utf-8"))
                    print(f"  [OK]   raw/{f.name} (valid JSON)")
                except Exception as exc:
                    print(f"  [FAIL] raw/{f.name} - JSON parse error: {exc}")
                    ok = False
        else:
            print("  [WARN] raw/ directory is empty")
    else:
        print("  [WARN] raw/ directory not found")

    # 5. Core CSV files
    for rel in DATA_CSVS:
        path = run_dir / rel
        if path.exists():
            row_count = _count_csv_rows(path)
            if row_count > 0:
                print(f"  [OK]   {rel} ({row_count} rows)")
            else:
                print(f"  [FAIL] {rel} is empty (0 rows)")
                ok = False
        else:
            print(f"  [WARN] Missing (may be absent for ablation/source-only runs): {rel}")

    # 6. Optional CSVs
    for rel in OPTIONAL_CSVS:
        path = run_dir / rel
        if path.exists():
            row_count = _count_csv_rows(path)
            print(f"  [OK]   {rel} ({row_count} rows)")

    # 7. Integrity file
    integrity_path = run_dir / "integrity.json"
    if integrity_path.exists():
        try:
            integrity = json.loads(integrity_path.read_text(encoding="utf-8"))
            files = integrity.get("files", {})
            mismatches = []
            for rel_path, file_info in files.items():
                actual = run_dir / rel_path
                if actual.exists():
                    expected_hash = file_info["sha256"] if isinstance(file_info, dict) else file_info
                    actual_hash = hashlib.sha256(actual.read_bytes()).hexdigest()
                    if actual_hash != expected_hash:
                        mismatches.append(rel_path)
            if mismatches:
                print(f"  [FAIL] Checksum mismatches: {mismatches}")
                ok = False
            else:
                print(f"  [OK]   integrity.json ({len(files)} checksums verified)")
        except Exception as exc:
            print(f"  [WARN] integrity.json - parse error: {exc}")
    else:
        print("  [WARN] integrity.json not found")

    # 8. Zip file
    zip_search_roots = [run_dir.parent, run_dir.parent.parent, Path.cwd()]
    zips = []
    seen_roots = set()
    for root in zip_search_roots:
        resolved = root.resolve()
        if resolved in seen_roots:
            continue
        seen_roots.add(resolved)
        zips.extend(root.glob(f"{run_dir.name}*.zip"))
    if zips:
        zip_path = zips[0]
        try:
            with zipfile.ZipFile(zip_path) as zf:
                names = zf.namelist()
            print(f"  [OK]   Archive: {zip_path.name} ({len(names)} files)")
        except Exception as exc:
            print(f"  [FAIL] Archive {zip_path.name} - {exc}")
            ok = False
    else:
        print("  [INFO] No run archive found")

    status = "PASS" if ok else "FAIL"
    print(f"\n  Result: {status}\n")
    return ok

def _count_csv_rows(path: Path) -> int:
    try:
        with path.open(newline="", encoding="utf-8") as f:
            reader = csv.reader(f)
            rows = sum(1 for _ in reader)
        return max(0, rows - 1)  # subtract header
    except Exception:
        return -1