import json
import math
import struct
import sys
from collections import Counter
from pathlib import Path

MAX_HEADER_BYTES = 100_000_000

# Matches safetensors commit 6eb4dc9a as checked on 2026-08-03.
DTYPE_BITS = {
    "BOOL": 8,
    "F4": 4,
    "F6_E2M3": 6,
    "F6_E3M2": 6,
    "U8": 8,
    "I8": 8,
    "F8_E4M3": 8,
    "F8_E5M2": 8,
    "F8_E8M0": 8,
    "F8_E4M3FNUZ": 8,
    "F8_E5M2FNUZ": 8,
    "U16": 16,
    "I16": 16,
    "F16": 16,
    "BF16": 16,
    "U32": 32,
    "I32": 32,
    "F32": 32,
    "C64": 64,
    "U64": 64,
    "I64": 64,
    "F64": 64,
}


def reject_duplicates(pairs):
    result = {}
    for key, value in pairs:
        if key in result:
            raise ValueError(f"duplicate JSON key: {key}")
        result[key] = value
    return result


def inspect_file(path):
    file_size = path.stat().st_size
    if file_size < 8:
        raise ValueError(f"{path}: shorter than the length prefix")

    with path.open("rb") as handle:
        prefix = handle.read(8)
        header_length = struct.unpack("<Q", prefix)[0]
        if not 2 <= header_length <= MAX_HEADER_BYTES:
            raise ValueError(f"{path}: invalid header length {header_length}")
        if header_length > file_size - 8:
            raise ValueError(f"{path}: header extends past end of file")

        raw_header = handle.read(header_length)

    if not raw_header.startswith(b"{"):
        raise ValueError(f"{path}: header does not start with '{{'")

    header = json.loads(
        raw_header.decode("utf-8"),
        object_pairs_hook=reject_duplicates,
    )
    if not isinstance(header, dict):
        raise ValueError(f"{path}: header must be a JSON object")

    data_bytes = file_size - 8 - header_length
    ranges = []
    tensors = {}
    elements = Counter()
    payload = Counter()

    for name, info in header.items():
        if name == "__metadata__":
            if not isinstance(info, dict) or any(
                not isinstance(key, str) or not isinstance(value, str)
                for key, value in info.items()
            ):
                raise ValueError(
                    f"{path}: __metadata__ must be a string-to-string map"
                )
            continue
        if not isinstance(info, dict):
            raise ValueError(f"{path}:{name}: tensor entry is not an object")

        dtype = info.get("dtype")
        shape = info.get("shape")
        offsets = info.get("data_offsets")

        if dtype not in DTYPE_BITS:
            raise ValueError(f"{path}:{name}: unknown dtype {dtype!r}")
        if not isinstance(shape, list) or any(
            type(dimension) is not int or dimension < 0
            for dimension in shape
        ):
            raise ValueError(f"{path}:{name}: invalid shape")
        if (
            not isinstance(offsets, list)
            or len(offsets) != 2
            or any(type(offset) is not int for offset in offsets)
        ):
            raise ValueError(f"{path}:{name}: invalid data_offsets")

        begin, end = offsets
        count = math.prod(shape)
        expected_bits = count * DTYPE_BITS[dtype]
        if expected_bits % 8:
            raise ValueError(
                f"{path}:{name}: sub-byte tensor is not byte-aligned"
            )
        expected_bytes = expected_bits // 8

        if begin < 0 or end < begin or end > data_bytes:
            raise ValueError(f"{path}:{name}: offsets leave the data buffer")
        if end - begin != expected_bytes:
            raise ValueError(
                f"{path}:{name}: shape and dtype require {expected_bytes} "
                f"bytes, offsets describe {end - begin}"
            )

        ranges.append((begin, end, name))

        tensors[name] = path.name
        elements[dtype] += count
        payload[dtype] += expected_bytes

    cursor = 0
    for begin, end, name in sorted(ranges):
        if begin != cursor:
            kind = "overlap" if begin < cursor else "hole"
            raise ValueError(f"{path}:{name}: {kind} before byte {begin}")
        cursor = end

    if cursor != data_bytes:
        raise ValueError(
            f"{path}: tensors cover {cursor} of {data_bytes} data bytes"
        )

    return {
        "tensors": tensors,
        "elements": elements,
        "payload": payload,
        "payload_bytes": data_bytes,
    }


root = Path(sys.argv[1] if len(sys.argv) > 1 else "model-audit")
files = sorted(root.glob("*.safetensors"))
if not files:
    raise SystemExit(f"no .safetensors files found under {root}")

actual = {}
duplicates = []
elements = Counter()
payload = Counter()
file_payload = {}

for path in files:
    result = inspect_file(path)
    file_payload[path.name] = result["payload_bytes"]
    elements.update(result["elements"])
    payload.update(result["payload"])

    for name, shard in result["tensors"].items():
        if name in actual:
            duplicates.append((name, actual[name], shard))
        actual[name] = shard

print("files:", len(files))
print("tensors:", len(actual))
print("duplicate tensors:", duplicates)
print("elements by dtype:", dict(sorted(elements.items())))
print("payload bytes by dtype:", dict(sorted(payload.items())))
print("total elements:", sum(elements.values()))
print("total payload bytes:", sum(payload.values()))

if duplicates:
    raise ValueError(f"duplicate tensors across shards: {duplicates}")

index_path = root / "model.safetensors.index.json"
if index_path.exists():
    raw_index = index_path.read_text(encoding="utf-8")
    if raw_index.startswith("version https://git-lfs.github.com/spec/"):
        raise SystemExit("shard index is still a Git LFS pointer")

    index = json.loads(raw_index, object_pairs_hook=reject_duplicates)
    weight_map = index.get("weight_map")
    if not isinstance(weight_map, dict):
        raise ValueError("index weight_map is missing or invalid")

    expected_names = set(weight_map)
    actual_names = set(actual)
    mapped_shards = set(weight_map.values())
    present_shards = set(file_payload)

    misplaced = sorted(
        (name, weight_map[name], actual[name])
        for name in expected_names & actual_names
        if weight_map[name] != actual[name]
    )

    missing_shards = sorted(mapped_shards - present_shards)
    unmapped_files = sorted(present_shards - mapped_shards)
    missing_tensors = sorted(expected_names - actual_names)
    unindexed_tensors = sorted(actual_names - expected_names)

    print("missing shards:", missing_shards)
    print("unmapped tensor files:", unmapped_files)
    print("missing tensors:", missing_tensors)
    print("unindexed tensors:", unindexed_tensors)
    print("misplaced tensors:", misplaced)

    discrepancies = {
        "missing_shards": missing_shards,
        "unmapped_files": unmapped_files,
        "missing_tensors": missing_tensors,
        "unindexed_tensors": unindexed_tensors,
        "misplaced_tensors": misplaced,
    }
    if any(discrepancies.values()):
        raise ValueError(
            "checkpoint reconciliation failed: "
            + json.dumps(discrepancies, sort_keys=True)
        )

    declared_total = index.get("metadata", {}).get("total_size")
    mapped_payload = sum(
        file_payload[shard]
        for shard in mapped_shards
        if shard in file_payload
    )
    print("declared total size:", declared_total)
    print("mapped payload bytes:", mapped_payload)

    if (
        declared_total is not None
        and mapped_shards <= present_shards
        and declared_total != mapped_payload
    ):
        raise ValueError(
            f"declared total {declared_total} != "
            f"mapped payload {mapped_payload}"
        )