"""Offline integrity and model-structure checks for the release directory."""

from __future__ import annotations

import hashlib
import json
from pathlib import Path

import torch
from safetensors import safe_open
from safetensors.torch import load_file

from modeling_anka import RMW3, RMW3Config, count_parameters


ROOT = Path(__file__).resolve().parent
EXPECTED_PARAMETERS = 48_944_657


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 main() -> None:
    required = (
        "model.safetensors",
        "tokenizer.json",
        "config.json",
        "modeling_anka.py",
        "inference.py",
    )
    for name in required:
        if not (ROOT / name).is_file():
            raise FileNotFoundError(name)

    config_data = json.loads((ROOT / "config.json").read_text())
    config = RMW3Config(
        **{field: config_data[field] for field in RMW3Config.__dataclass_fields__}
    )
    model = RMW3(config).to(dtype=torch.bfloat16)
    state = load_file(str(ROOT / "model.safetensors"), device="cpu")
    model.load_state_dict(state, strict=True)

    if count_parameters(model) != EXPECTED_PARAMETERS:
        raise RuntimeError("Parameter-count check failed")
    if model.lm_head.weight.data_ptr() != model.token_embeddings.weight.data_ptr():
        raise RuntimeError("Weight-tying check failed")
    if not torch.equal(state["lm_head.weight"], state["token_embeddings.weight"]):
        raise RuntimeError("Serialized tied weights differ")

    with safe_open(ROOT / "model.safetensors", framework="pt", device="cpu") as file:
        if len(file.keys()) != 199:
            raise RuntimeError("Unexpected tensor count")

    print("ANKA-50M RMW3 release: PASS")
    print(f"Parameters : {EXPECTED_PARAMETERS:,}")
    print(f"Tensors    : {len(state)}")
    print(f"Weights SHA: {sha256(ROOT / 'model.safetensors')}")


if __name__ == "__main__":
    main()