# Copyright 2026 Spawn Lau
# SPDX-License-Identifier: Apache-2.0
"""Rebuild the published compact bundle from the pinned upstream file."""

import argparse
import hashlib
import json
from pathlib import Path
import subprocess
import sys

SOURCE_BYTES = 387710976
SOURCE_SHA256 = "92dcbea108899e5d6e30d919b0744f90d9967e80c67a4ab5503ac16d54f62eb0"
OUTPUT_BYTES = 309559296
OUTPUT_SHA256 = "a08a6bcbcba893d8e52dcf2ae05ce06847f942e3b0e5ea8b8df93f32aa23b027"


def sha256(path):
    with path.open("rb") as stream:
        return hashlib.file_digest(stream, "sha256").hexdigest()


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model", type=Path, required=True)
    parser.add_argument("--work-dir", type=Path, required=True,
                        help="A new directory; existing directories are never overwritten.")
    args = parser.parse_args()
    model = args.model.resolve(strict=True)
    if model.stat().st_size != SOURCE_BYTES or sha256(model) != SOURCE_SHA256:
        raise ValueError("Input is not the pinned generic CPU/GPU upstream bundle.")
    args.work_dir.mkdir(parents=True, exist_ok=False)
    scripts = Path(__file__).resolve().parent
    candidate_dir = args.work_dir.resolve() / "candidates"
    commands = [
        ["build_size_candidates.py", "--model", str(model), "--output-dir", str(candidate_dir)],
        ["refine_position_candidate.py", "--candidate-dir", str(candidate_dir)],
        ["compact_candidate_metadata.py", "--candidate-dir", str(candidate_dir), "--reproducible"],
    ]
    for script, *arguments in commands:
        subprocess.run([sys.executable, str(scripts / script), *arguments], check=True)
    output = candidate_dir / "eg2-gallery-compact.litertlm"
    result = {"path": str(output), "bytes": output.stat().st_size, "sha256": sha256(output)}
    if result["bytes"] != OUTPUT_BYTES or result["sha256"] != OUTPUT_SHA256:
        raise RuntimeError(f"Rebuild differs from the published bundle: {result}")
    print(json.dumps({"status": "verified", **result}, indent=2))


if __name__ == "__main__":
    main()