#!/usr/bin/env python3
"""Verify the release checksum and the required GGUF schema without dependencies."""

from __future__ import annotations

import argparse
import hashlib
import json
import struct
from collections import Counter
from pathlib import Path
from typing import BinaryIO


SCALAR_FORMATS = {
    0: "<B",
    1: "<b",
    2: "<H",
    3: "<h",
    4: "<I",
    5: "<i",
    6: "<f",
    7: "<?",
    10: "<Q",
    11: "<q",
    12: "<d",
}


def read_exact(stream: BinaryIO, size: int) -> bytes:
    value = stream.read(size)
    if len(value) != size:
        raise ValueError("unexpected end of GGUF file")
    return value


def read_scalar(stream: BinaryIO, value_type: int):
    fmt = SCALAR_FORMATS[value_type]
    return struct.unpack(fmt, read_exact(stream, struct.calcsize(fmt)))[0]


def read_string(stream: BinaryIO, keep: bool = True):
    size = read_scalar(stream, 10)
    value = read_exact(stream, size)
    return value.decode("utf-8") if keep else None


def read_value(stream: BinaryIO, value_type: int, keep: bool = True):
    if value_type in SCALAR_FORMATS:
        value = read_scalar(stream, value_type)
        return value if keep else None
    if value_type == 8:
        return read_string(stream, keep)
    if value_type == 9:
        element_type = read_scalar(stream, 4)
        count = read_scalar(stream, 10)
        values = [read_value(stream, element_type, keep) for _ in range(count)]
        return values if keep else None
    raise ValueError(f"unsupported GGUF metadata type {value_type}")


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


def inspect(path: Path, expected_keys: set[str]) -> dict:
    with path.open("rb") as stream:
        if read_exact(stream, 4) != b"GGUF":
            raise ValueError("invalid GGUF magic")
        version = read_scalar(stream, 4)
        tensor_count = read_scalar(stream, 10)
        metadata_count = read_scalar(stream, 10)
        metadata = {}
        for _ in range(metadata_count):
            key = read_string(stream)
            value_type = read_scalar(stream, 4)
            value = read_value(stream, value_type, key in expected_keys)
            if key in expected_keys:
                metadata[key] = value

        tensor_types = Counter()
        offsets = []
        for _ in range(tensor_count):
            read_string(stream, keep=False)
            dimensions = read_scalar(stream, 4)
            for _ in range(dimensions):
                read_scalar(stream, 10)
            tensor_types[str(read_scalar(stream, 4))] += 1
            offsets.append(read_scalar(stream, 10))

    return {
        "version": version,
        "tensor_count": tensor_count,
        "metadata": metadata,
        "tensor_type_counts": dict(tensor_types),
        "aligned": all(offset % 32 == 0 for offset in offsets),
    }


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("artifact", nargs="?", type=Path)
    parser.add_argument("--manifest", type=Path, default=Path("manifest/release.json"))
    args = parser.parse_args()
    expected = json.loads(args.manifest.read_text())
    artifact = args.artifact or Path(expected["artifact"]["file"])

    failures = []
    actual_size = artifact.stat().st_size
    actual_hash = sha256(artifact)
    if actual_size != expected["artifact"]["size"]:
        failures.append(f"size: expected {expected['artifact']['size']}, got {actual_size}")
    if actual_hash != expected["artifact"]["sha256"]:
        failures.append(f"sha256: expected {expected['artifact']['sha256']}, got {actual_hash}")

    gguf_expected = expected["gguf"]
    actual = inspect(artifact, set(gguf_expected["metadata"]))
    for key in ("version", "tensor_count", "tensor_type_counts", "aligned"):
        if actual[key] != gguf_expected[key]:
            failures.append(f"{key}: expected {gguf_expected[key]!r}, got {actual[key]!r}")
    for key, value in gguf_expected["metadata"].items():
        if actual["metadata"].get(key) != value:
            failures.append(
                f"metadata {key}: expected {value!r}, got {actual['metadata'].get(key)!r}"
            )

    if failures:
        raise SystemExit("verification failed:\n- " + "\n- ".join(failures))
    print(f"verified {artifact}: {actual_size} bytes {actual_hash}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())