#!/usr/bin/env python3
"""Emit a compact, reproducible tensor and parameter audit for GGUF files."""

from __future__ import annotations

import argparse
import hashlib
import json
from collections import Counter
from pathlib import Path

from gguf import GGUFReader


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 audit(path: Path, with_hash: bool) -> dict[str, object]:
    reader = GGUFReader(path, "r")
    tensors: Counter[str] = Counter()
    parameters: Counter[str] = Counter()
    protected: dict[str, str] = {}

    for tensor in reader.tensors:
        tensor_type = tensor.tensor_type.name.lower()
        tensors[tensor_type] += 1
        parameters[tensor_type] += int(tensor.n_elements)
        if tensor.name.endswith(
            (
                "audio_patch_proj.weight",
                "video_patch_proj.weight",
                "condition_proj.weight",
            )
        ):
            protected[tensor.name] = tensor_type

    result: dict[str, object] = {
        "file": path.name,
        "bytes": path.stat().st_size,
        "tensor_count": len(reader.tensors),
        "tensor_histogram": dict(sorted(tensors.items())),
        "parameter_histogram": dict(sorted(parameters.items())),
        "protected_tensors": dict(sorted(protected.items())),
    }
    if with_hash:
        result["sha256"] = sha256(path)
    return result


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("files", nargs="+", type=Path)
    parser.add_argument("--sha256", action="store_true")
    parser.add_argument("--output", type=Path)
    args = parser.parse_args()

    report = [audit(path.resolve(), args.sha256) for path in args.files]
    rendered = json.dumps(report, indent=2, sort_keys=True) + "\n"
    if args.output:
        args.output.write_text(rendered, encoding="utf-8")
    else:
        print(rendered, end="")


if __name__ == "__main__":
    main()