#!/usr/bin/env python3
"""Validate the self-contained VoxCPM2 FP8 W8A8 V1 release payload."""

from __future__ import annotations

import hashlib
import json
import struct
from pathlib import Path
from typing import Any


ROOT = Path(__file__).resolve().parent
MANIFEST_PATH = ROOT / "MANIFEST.json"
EXPECTED_FORMAT = "voxcpm2-torchao-fp8-static-v1"
EXPECTED_SOURCE_REVISION = "32279effe8c19989596f05d353d1447f51d9e915"
EXPECTED_MODEL_SHA256 = "fd86c0740d0faa9ba6017b28453851b68f754c3ac9b5d3051f0987e4a77cbd82"
EXPECTED_QUANTIZED_MODULES = 424
IGNORED_TOP_LEVEL = {".cache", ".venv", "__pycache__"}


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


def distributed_files() -> set[str]:
    result: set[str] = set()
    for path in ROOT.rglob("*"):
        relative = path.relative_to(ROOT)
        if relative.parts[0] in IGNORED_TOP_LEVEL or not path.is_file():
            continue
        if relative.as_posix() != "MANIFEST.json":
            result.add(relative.as_posix())
    return result


def safetensors_header(path: Path) -> dict[str, Any]:
    with path.open("rb") as handle:
        raw_length = handle.read(8)
        if len(raw_length) != 8:
            raise ValueError("Truncated safetensors length header")
        header_length = struct.unpack("<Q", raw_length)[0]
        if not 2 <= header_length <= 128 * 1024 * 1024:
            raise ValueError(f"Invalid safetensors header length: {header_length}")
        raw_header = handle.read(header_length)
    if len(raw_header) != header_length:
        raise ValueError("Truncated safetensors JSON header")
    value = json.loads(raw_header)
    if not isinstance(value, dict):
        raise ValueError("Safetensors header is not an object")
    return value


def main() -> None:
    manifest = json.loads(MANIFEST_PATH.read_text(encoding="utf-8"))
    records = manifest.get("files")
    if manifest.get("schema_version") != 1 or not isinstance(records, dict):
        raise ValueError("Invalid release manifest")

    actual_files = distributed_files()
    if actual_files != set(records):
        missing = sorted(set(records) - actual_files)
        unexpected = sorted(actual_files - set(records))
        raise ValueError(f"Release file set mismatch; missing={missing}, unexpected={unexpected}")

    for relative in sorted(records):
        path = ROOT / relative
        record = records[relative]
        actual_size = path.stat().st_size
        if actual_size != record.get("bytes"):
            raise ValueError(f"Size mismatch for {relative}")
        actual_hash = sha256(path)
        if actual_hash != record.get("sha256"):
            raise ValueError(f"SHA-256 mismatch for {relative}")
        print(f"ok {relative} {actual_size} bytes")

    fp8_config = json.loads((ROOT / "fp8_config.json").read_text(encoding="utf-8"))
    if fp8_config.get("format") != EXPECTED_FORMAT:
        raise ValueError("Unexpected FP8 artifact format")
    if fp8_config.get("source", {}).get("revision") != EXPECTED_SOURCE_REVISION:
        raise ValueError("Unexpected source revision")
    quantization = fp8_config.get("quantization", {})
    selected = quantization.get("quantized_modules", [])
    state_names = quantization.get("fp8_state_names", [])
    if len(selected) != EXPECTED_QUANTIZED_MODULES or len(state_names) != EXPECTED_QUANTIZED_MODULES:
        raise ValueError("Unexpected FP8 module/state count")

    model_path = ROOT / fp8_config["artifact"]["model_file"]
    if fp8_config["artifact"].get("model_sha256") != EXPECTED_MODEL_SHA256:
        raise ValueError("FP8 config contains an unexpected model checksum")
    if records[model_path.name]["sha256"] != EXPECTED_MODEL_SHA256:
        raise ValueError("Manifest model checksum does not match the release checksum")

    header = safetensors_header(model_path)
    tensor_keys = set(header) - {"__metadata__"}
    for state_name in state_names:
        required = {
            state_name + ".__fp8_qdata",
            state_name + ".__fp8_scale",
            state_name + ".__fp8_activation_scale",
        }
        missing = required - tensor_keys
        if missing:
            raise ValueError(f"Missing FP8 tensors for {state_name}: {sorted(missing)}")

    config = json.loads((ROOT / "config.json").read_text(encoding="utf-8"))
    if config.get("architecture") != "voxcpm2" or config.get("audio_vae_config", {}).get("out_sample_rate") != 48000:
        raise ValueError("Unexpected VoxCPM2 architecture configuration")

    print(
        json.dumps(
            {
                "status": "valid",
                "files": len(records),
                "bytes": sum(record["bytes"] for record in records.values()),
                "model_sha256": EXPECTED_MODEL_SHA256,
                "quantized_modules": EXPECTED_QUANTIZED_MODULES,
                "source_revision": EXPECTED_SOURCE_REVISION,
            },
            indent=2,
            sort_keys=True,
        )
    )


if __name__ == "__main__":
    main()