"""Deterministic runtime identities for frozen representation modules."""

from __future__ import annotations

import hashlib
import json

import torch
from torch import nn


def module_state_sha256(module: nn.Module) -> str:
    """Hash sorted tensor keys, metadata, and loaded tensor bytes."""
    digest = hashlib.sha256()
    state = module.state_dict()
    for key in sorted(state):
        value = state[key]
        if not isinstance(value, torch.Tensor):
            raise TypeError(
                f"state_dict entry {key!r} is not a tensor: "
                f"{type(value).__name__}")
        metadata = json.dumps(
            {
                "key": str(key),
                "dtype": str(value.dtype),
                "shape": [int(size) for size in value.shape],
            },
            sort_keys=True,
            separators=(",", ":"),
        ).encode("utf-8")
        digest.update(len(metadata).to_bytes(8, "little"))
        digest.update(metadata)
        cpu_bytes = (
            value.detach()
            .to(device="cpu")
            .contiguous()
            .reshape(-1)
            .view(torch.uint8)
            .numpy()
        )
        digest.update(int(cpu_bytes.nbytes).to_bytes(8, "little"))
        digest.update(memoryview(cpu_bytes))
    return digest.hexdigest()


def representation_runtime_identity(module: nn.Module) -> dict:
    """Return the strict identity used by frozen-encoder metadata checks."""
    return {
        "wrapper_class": (
            f"{module.__class__.__module__}.{module.__class__.__qualname__}"),
        "torch_version": str(torch.__version__),
        "state_sha256": module_state_sha256(module),
        "backend": "unknown",
        "backend_version": "unknown",
    }