"""Deterministic identities for frozen files and JSON metadata."""

from __future__ import annotations

import hashlib
import json
import os
from typing import Callable, Optional

import torch.distributed as dist


def canonical_sha256(value) -> str:
    """Hash a JSON value with one canonical serialization."""
    payload = json.dumps(
        value,
        sort_keys=True,
        separators=(",", ":"),
        ensure_ascii=False,
    ).encode("utf-8")
    return hashlib.sha256(payload).hexdigest()


def canonical_identity_value(value):
    """Normalize tuples and lists through strict, finite canonical JSON."""
    try:
        payload = json.dumps(
            value,
            sort_keys=True,
            separators=(",", ":"),
            ensure_ascii=False,
            allow_nan=False,
        )
    except (TypeError, ValueError) as exc:
        raise ValueError(
            "artifact identity must contain only finite JSON-serializable values"
        ) from exc
    return json.loads(payload)


def file_sha256(path: str) -> str:
    """Hash a file without loading it wholly into memory."""
    digest = hashlib.sha256()
    with open(path, "rb") as handle:
        while chunk := handle.read(8 * 1024 * 1024):
            digest.update(chunk)
    return digest.hexdigest()


def distributed_file_sha256(path: str) -> str:
    """Hash a shared file on rank zero and broadcast the result."""
    if (
        not dist.is_available()
        or not dist.is_initialized()
        or dist.get_world_size() == 1
    ):
        return file_sha256(path)
    payload = [None]
    if dist.get_rank() == 0:
        try:
            payload[0] = {"ok": True, "sha256": file_sha256(path)}
        except BaseException as exc:
            payload[0] = {
                "ok": False,
                "error": f"{type(exc).__name__}: {exc}",
            }
    dist.broadcast_object_list(payload, src=0)
    result = payload[0]
    if not isinstance(result, dict) or not result.get("ok"):
        error = (
            result.get("error", "unknown error")
            if isinstance(result, dict)
            else "no rank-zero result"
        )
        raise RuntimeError(f"rank-zero file hashing failed: {error}")
    return str(result["sha256"])


def runtime_file_identity(
    path,
    *,
    file_hasher: Optional[Callable[[str], str]] = None,
) -> dict:
    """Return path, stable stat fields, and SHA256 for one file."""
    resolved = os.path.realpath(os.path.expanduser(os.fspath(path)))
    if not os.path.isfile(resolved):
        raise FileNotFoundError(
            f"runtime identity target is not a file: {resolved}")
    before = os.stat(resolved)
    digest = str((file_hasher or distributed_file_sha256)(resolved))
    after = os.stat(resolved)
    if (
        int(before.st_size) != int(after.st_size)
        or int(before.st_mtime_ns) != int(after.st_mtime_ns)
    ):
        raise RuntimeError(
            f"runtime identity target changed while hashing: {resolved}")
    if (
        len(digest) != 64
        or any(character not in "0123456789abcdef" for character in digest)
    ):
        raise ValueError(
            f"runtime identity returned an invalid SHA256 for {resolved}")
    return {
        "path": resolved,
        "size_bytes": int(after.st_size),
        "mtime_ns": int(after.st_mtime_ns),
        "sha256": digest,
    }


def distributed_runtime_file_identity(path) -> dict:
    """Build the complete file identity on rank zero and broadcast it."""
    if (
        not dist.is_available()
        or not dist.is_initialized()
        or dist.get_world_size() == 1
    ):
        return runtime_file_identity(path, file_hasher=file_sha256)
    payload = [None]
    if dist.get_rank() == 0:
        try:
            payload[0] = {
                "ok": True,
                "identity": runtime_file_identity(
                    path, file_hasher=file_sha256),
            }
        except BaseException as exc:
            payload[0] = {
                "ok": False,
                "error": f"{type(exc).__name__}: {exc}",
            }
    dist.broadcast_object_list(payload, src=0)
    result = payload[0]
    if not isinstance(result, dict) or not result.get("ok"):
        error = (
            result.get("error", "unknown error")
            if isinstance(result, dict)
            else "no rank-zero result"
        )
        raise RuntimeError(
            f"rank-zero runtime file identity failed: {error}")
    identity = result.get("identity")
    if not isinstance(identity, dict):
        raise RuntimeError(
            "rank-zero runtime file identity returned invalid data")
    return dict(identity)