#!/usr/bin/env python3
"""Validate the non-destructive /data/wxy/sskj-h3 experiment archive."""

from __future__ import annotations

import hashlib
import os
from pathlib import Path


ROOT = Path(os.environ.get("SSKJ_H3_ROOT", "/data/wxy/sskj-h3"))

MAPPINGS = {
    "sglang-base": [
        ("/data/wxy/run_sglang_h3_mixed_matrix_6000d.sh", "scripts/run_sglang_h3_mixed_matrix_6000d.sh"),
        ("/data/wxy/minimax_h3_mixed_bench.py", "scripts/minimax_h3_mixed_bench.py"),
        ("/data/wxy/results/minimax_h3_mixed_matrix/mixed64-20steps-5s-20260822-100844", "results/mixed64-20steps-5s-20260822-100844"),
        ("/data/wxy/results/minimax_h3_mixed_matrix/balanced-tp4-tp2-20steps-5s-20260822-175030", "results/balanced-tp4-tp2-20steps-5s-20260822-175030"),
    ],
    "vllm-omni-base": [
        ("/data/wxy/run_vllm_omni_h3_matrix_6000d.sh", "scripts/run_vllm_omni_h3_matrix_6000d.sh"),
        ("/data/wxy/minimax_h3_vllm_bench.py", "scripts/minimax_h3_vllm_bench.py"),
        ("/data/wxy/results/minimax_h3_vllm_matrix/vllm-balanced64-20steps-5s-r2-20260822-225317", "results/vllm-balanced64-20steps-5s-r2-20260822-225317"),
    ],
    "sglang-profile": [
        ("/data/wxy/h3_profile", "scripts/h3_profile"),
        ("/data/wxy/profile_results/h3-quick-input-matrix-20260824-run1", "results/h3-quick-input-matrix-20260824-run1"),
        ("/data/wxy/profile_results/h3-targeted-profile-20260824-run1", "results/h3-targeted-profile-20260824-run1"),
        ("/data/wxy/profile_results/h3-sdpa-kernel-matrix-20260825-run1", "results/h3-sdpa-kernel-matrix-20260825-run1"),
    ],
    "vbench-base": [
        ("/data/wxy/results/h3_vbench_base/h3-vbench-base-dense-tp2x4-20260826-run1", "results/h3-vbench-base-dense-tp2x4-20260826-run1"),
        ("/data/wxy/vbench_score_setup_logs", "environment/vbench_score_setup_logs"),
    ],
}

VBENCH_SCRIPTS = [
    "check_vbench_score_env.py",
    "h3_vbench_generate.py",
    "prefetch_vbench_weights.sh",
    "rebuild_vbench_score_env.sh",
    "run_h3_vbench_base_tp2x4.sh",
    "setup_vbench_6000d.sh",
]


def tree_stats(path: Path) -> tuple[int, int]:
    if path.is_file():
        return 1, path.stat().st_size
    files = [
        p
        for p in path.rglob("*")
        if p.is_file() and "__pycache__" not in p.parts and p.suffix != ".pyc"
    ]
    return len(files), sum(p.stat().st_size for p in files)


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def main() -> None:
    errors: list[str] = []
    if not ROOT.is_dir():
        raise SystemExit(f"archive root missing: {ROOT}")

    for section, mappings in MAPPINGS.items():
        section_root = ROOT / section
        for required in ("README.md", "scripts", "results"):
            if not (section_root / required).exists():
                errors.append(f"missing {section}/{required}")
        for source_text, relative_dest in mappings:
            source = Path(source_text)
            dest = section_root / relative_dest
            if not dest.exists():
                errors.append(f"missing copied target: {dest}")
                continue
            if tree_stats(source) != tree_stats(dest):
                errors.append(
                    f"tree mismatch: {source} {tree_stats(source)} != {dest} {tree_stats(dest)}"
                )

    for name in VBENCH_SCRIPTS:
        source = Path("/data/wxy") / name
        dest = ROOT / "vbench-base/scripts" / name
        if not dest.is_file():
            errors.append(f"missing VBench script: {dest}")
        elif sha256(source) != sha256(dest):
            errors.append(f"VBench script checksum mismatch: {name}")

    required_docs = {
        "README.md": "sskj-h3",
        "vbench-base/VBENCH_SCORING_ADAPTATIONS.md": "timm==0.9.12",
        "vbench-base/README.md": "0.8159256025548192",
        "SOURCE_MAP.tsv": "h3-targeted-profile-20260824-run1",
    }
    for relative, marker in required_docs.items():
        path = ROOT / relative
        if not path.is_file() or marker not in path.read_text(encoding="utf-8"):
            errors.append(f"missing documentation marker {marker!r}: {path}")

    for section in MAPPINGS:
        sums = ROOT / section / "scripts/SHA256SUMS"
        if not sums.is_file() or not sums.read_text(encoding="utf-8").strip():
            errors.append(f"missing script checksum manifest: {sums}")

    if errors:
        raise SystemExit("archive validation failed:\n- " + "\n- ".join(errors))
    print(f"ARCHIVE_GREEN root={ROOT} sections={len(MAPPINGS)}")


if __name__ == "__main__":
    main()