#!/usr/bin/env python3
"""Verify and summarize repeated deterministic decoder-training runs."""

from __future__ import annotations

import argparse
import hashlib
import json
from datetime import UTC, datetime
from pathlib import Path


MATCH_FIELDS = (
    "dataset_manifest_sha256",
    "requested_images",
    "training_images",
    "model_sha256",
    "initialization",
    "initialization_sha256",
    "seed",
    "image_size",
    "encoder_layers",
    "batch_size",
    "steps",
    "data_order",
    "objective",
    "optimizer",
    "initial_learning_rate",
    "final_learning_rate",
    "decoder_parameters",
    "final_l1",
)


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


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--run", action="append", required=True, type=Path)
    parser.add_argument("--output", required=True, type=Path)
    args = parser.parse_args()
    if len(args.run) < 2:
        raise SystemExit("at least two --run directories are required")

    records: list[tuple[Path, dict[str, object], str, str]] = []
    for directory in args.run:
        training_path = directory / "training.json"
        decoder_path = directory / "decoder.bin"
        training = json.loads(training_path.read_text(encoding="utf-8-sig"))
        decoder_hash = sha256(decoder_path)
        if decoder_hash != training["decoder_sha256"]:
            raise SystemExit(f"{decoder_path} disagrees with training.json")
        records.append((directory, training, sha256(training_path), decoder_hash))

    baseline = records[0][1]
    for directory, training, _, _ in records[1:]:
        differences = [field for field in MATCH_FIELDS if training[field] != baseline[field]]
        if differences:
            raise SystemExit(f"{directory} differs in deterministic fields: {differences}")
    decoder_hashes = {record[3] for record in records}
    if len(decoder_hashes) != 1:
        raise SystemExit(f"decoder hashes differ: {sorted(decoder_hashes)}")

    output = {
        "schema_version": 1,
        "created_utc": datetime.now(UTC).isoformat(),
        "purpose": "short deterministic replay; not a quality or throughput result",
        "configuration": {field: baseline[field] for field in MATCH_FIELDS},
        "runs": [
            {
                "name": directory.name,
                "training_record_sha256": training_hash,
                "decoder_sha256": decoder_hash,
            }
            for directory, _, training_hash, decoder_hash in records
        ],
        "all_decoder_bytes_identical": True,
        "decoder_sha256": records[0][3],
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(output, indent=2) + "\n", encoding="utf-8")
    print(
        f"verified {len(records)} deterministic runs: "
        f"decoder SHA-256 {records[0][3]}"
    )


if __name__ == "__main__":
    main()