"""Canonical JSON serialization for hash-chain integrity.

THIS MODULE IS FROZEN. Any change to the serialization format
invalidates all existing audit chains. Do not modify without
an explicit migration strategy and version bump.

Canonical form rules:
  1. Keys sorted alphabetically (recursive)
  2. No whitespace (separators=(',', ':'))
  3. Timestamps as ISO 8601 UTC strings with microsecond precision
     e.g., "2026-03-21T14:32:01.847392+00:00"
  4. Unicode NFC normalization on all string values
  5. The `chain_hash` field is EXCLUDED from canonical form
     (it's computed from the canonical form of everything else)
  6. None values are serialized as JSON null
  7. Enums are serialized as their string value
  8. Nested objects follow the same rules recursively

Format version: 1
"""

from __future__ import annotations

import hashlib
import json
import unicodedata
from datetime import datetime
from enum import StrEnum
from typing import Any

FORMAT_VERSION = 1

# The genesis hash used as the "previous hash" for the first record in a chain.
# This is a fixed, well-known constant. Changing it would break all chains.
GENESIS_HASH = "0" * 64  # 64 hex zeros = SHA-256 of nothing, by convention


def canonical_json(obj: dict[str, Any], *, exclude_fields: set[str] | None = None) -> str:
    """Serialize a dictionary to canonical JSON form.

    Args:
        obj: The dictionary to serialize (typically an AuditRecord.model_dump()).
        exclude_fields: Field names to exclude from serialization.
            For AuditRecord hashing, this should be {"chain_hash"}.

    Returns:
        Deterministic JSON string suitable for hashing.
    """
    filtered = {k: v for k, v in obj.items() if k not in (exclude_fields or set())}
    return json.dumps(
        filtered,
        sort_keys=True,
        separators=(",", ":"),
        default=_canonical_default,
        ensure_ascii=False,
    )


def _canonical_default(obj: Any) -> Any:
    """Custom JSON serializer for types not natively supported."""
    if isinstance(obj, datetime):
        return obj.isoformat()
    if isinstance(obj, StrEnum):
        return str(obj)
    msg = f"Object of type {type(obj).__name__} is not JSON serializable"
    raise TypeError(msg)


def normalize_string(value: str) -> str:
    """Apply Unicode NFC normalization to a string.

    This ensures that equivalent Unicode representations
    (e.g., é as a single codepoint vs e + combining accent)
    produce identical bytes when hashed.
    """
    return unicodedata.normalize("NFC", value)


def compute_hash(data: str) -> str:
    """Compute SHA-256 hash of a string, returning hex digest.

    The input string is encoded as UTF-8 before hashing.
    NFC normalization should be applied before calling this function
    if the data contains user-provided strings.
    """
    return hashlib.sha256(data.encode("utf-8")).hexdigest()


def compute_params_hash(params: dict[str, Any] | None) -> str:
    """Hash the parameters of an MCP request.

    Args:
        params: The JSON-RPC params object. None for methods without params.

    Returns:
        SHA-256 hex digest of the canonical JSON of params,
        or hash of empty string for None params.
    """
    if params is None:
        return compute_hash("")
    return compute_hash(canonical_json(params))


def compute_result_hash(result_body: bytes | str | None) -> str | None:
    """Hash the upstream response body.

    Args:
        result_body: Raw response bytes or string. None if the request was denied.

    Returns:
        SHA-256 hex digest, or None if result_body is None.
    """
    if result_body is None:
        return None
    if isinstance(result_body, str):
        return compute_hash(result_body)
    return hashlib.sha256(result_body).hexdigest()


def compute_chain_hash(previous_chain_hash: str, record_canonical_json: str) -> str:
    """Compute the chain hash for a new audit record.

    chain_hash = SHA-256(previous_chain_hash + canonical_json_of_record)

    This creates a tamper-evident chain: modifying any record
    invalidates all subsequent chain_hash values.

    Args:
        previous_chain_hash: The chain_hash of the preceding record,
            or GENESIS_HASH for the first record.
        record_canonical_json: Canonical JSON of the current record
            (with chain_hash excluded).

    Returns:
        SHA-256 hex digest of the concatenation.
    """
    combined = previous_chain_hash + record_canonical_json
    return compute_hash(combined)
