"""Fail-closed audit for a small byte-level BPE compression contract.

The illustrative policy is intentionally tiny.  It demonstrates the identities and
evidence a production tokenizer release must bind; it is not a useful vocabulary.
"""

from __future__ import annotations

from dataclasses import asdict, dataclass
from datetime import datetime
import hashlib
import json
import re
from typing import Any


MAX_TEXT_BYTES = 4096
MAX_VOCABULARY = 512
MAX_MERGES = 512
IDENTIFIER = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._:@/-]{0,127}$")
SHA256_ID = re.compile(r"^[a-z][a-z0-9-]*@sha256:[0-9a-f]{64}$")


def _identifier(name: str, value: object) -> str:
    if type(value) is not str or not IDENTIFIER.fullmatch(value):
        raise ValueError(f"{name} must be a bounded identifier")
    return value


def _timestamp(name: str, value: object) -> str:
    if type(value) is not str or len(value) > 64:
        raise ValueError(f"{name} must be a bounded ISO-8601 timestamp")
    try:
        parsed = datetime.fromisoformat(value.replace("Z", "+00:00"))
    except ValueError as exc:
        raise ValueError(f"{name} must be an ISO-8601 timestamp") from exc
    if parsed.tzinfo is None:
        raise ValueError(f"{name} must include an offset")
    return value


def _hex_piece(name: str, value: object, *, maximum_bytes: int = 64) -> str:
    if type(value) is not str or not value or len(value) > maximum_bytes * 2:
        raise ValueError(f"{name} must be non-empty bounded lowercase hex")
    if value != value.lower() or len(value) % 2 or re.fullmatch(r"[0-9a-f]+", value) is None:
        raise ValueError(f"{name} must be non-empty bounded lowercase hex")
    return value


def _canonical(value: Any) -> Any:
    if isinstance(value, dict):
        return {key: _canonical(item) for key, item in sorted(value.items())}
    if isinstance(value, tuple):
        return [_canonical(item) for item in value]
    return value


def _content_id(prefix: str, payload: dict[str, Any]) -> str:
    encoded = json.dumps(
        _canonical(payload), sort_keys=True, separators=(",", ":"), allow_nan=False
    ).encode("utf-8")
    return f"{prefix}@sha256:{hashlib.sha256(encoded).hexdigest()}"


@dataclass(frozen=True)
class TokenDefinition:
    token_id: str
    piece_hex: str

    def __post_init__(self) -> None:
        _identifier("token_id", self.token_id)
        _hex_piece("piece_hex", self.piece_hex)


@dataclass(frozen=True)
class MergeRule:
    rank: int
    left_hex: str
    right_hex: str
    merged_hex: str

    def __post_init__(self) -> None:
        if type(self.rank) is not int or not 0 <= self.rank < MAX_MERGES:
            raise ValueError("rank must be a bounded integer")
        left = _hex_piece("left_hex", self.left_hex)
        right = _hex_piece("right_hex", self.right_hex)
        merged = _hex_piece("merged_hex", self.merged_hex)
        if merged != left + right:
            raise ValueError("merged_hex must concatenate left_hex and right_hex")


def _token_definition(value: object) -> TokenDefinition:
    if type(value) is not TokenDefinition:
        raise TypeError("vocabulary entries must be concrete TokenDefinition records")
    value.__post_init__()
    return TokenDefinition(value.token_id, value.piece_hex)


def _merge_rule(value: object) -> MergeRule:
    if type(value) is not MergeRule:
        raise TypeError("merge entries must be concrete MergeRule records")
    value.__post_init__()
    return MergeRule(value.rank, value.left_hex, value.right_hex, value.merged_hex)


@dataclass(frozen=True)
class TokenizerContract:
    contract_version: str
    tokenizer_id: str
    tokenizer_version: str
    normalization: str
    byte_encoding: str
    merge_semantics: str
    unknown_policy: str
    fallback_policy: str
    fallback_token_prefix: str
    maximum_text_bytes: int
    vocabulary: tuple[TokenDefinition, ...]
    merges: tuple[MergeRule, ...]
    scope: str
    vocabulary_owner: str
    audit_owner: str

    def __post_init__(self) -> None:
        for field in (
            "contract_version",
            "tokenizer_id",
            "tokenizer_version",
            "scope",
            "vocabulary_owner",
            "audit_owner",
        ):
            _identifier(field, getattr(self, field))
        if self.normalization != "none-codepoints-preserved-v1":
            raise ValueError("unsupported normalization")
        if self.byte_encoding != "utf-8-strict-v1":
            raise ValueError("unsupported byte_encoding")
        if self.merge_semantics != "lowest-rank-leftmost-repeat-v1":
            raise ValueError("unsupported merge_semantics")
        if self.unknown_policy != "no-unknown-token-v1":
            raise ValueError("unsupported unknown_policy")
        if self.fallback_policy != "all-raw-bytes-v1":
            raise ValueError("unsupported fallback_policy")
        _identifier("fallback_token_prefix", self.fallback_token_prefix)
        if type(self.maximum_text_bytes) is not int or not 1 <= self.maximum_text_bytes <= MAX_TEXT_BYTES:
            raise ValueError("maximum_text_bytes is outside the supported bound")
        if type(self.vocabulary) is not tuple or not 1 <= len(self.vocabulary) <= MAX_VOCABULARY:
            raise TypeError("vocabulary must be a non-empty bounded immutable tuple")
        if type(self.merges) is not tuple or len(self.merges) > MAX_MERGES:
            raise TypeError("merges must be a bounded immutable tuple")
        vocabulary = tuple(_token_definition(item) for item in self.vocabulary)
        merges = tuple(_merge_rule(item) for item in self.merges)
        token_ids = [item.token_id for item in vocabulary]
        pieces = [item.piece_hex for item in vocabulary]
        if len(token_ids) != len(set(token_ids)) or len(pieces) != len(set(pieces)):
            raise ValueError("vocabulary token IDs and pieces must be unique")
        if any(item.token_id.startswith(self.fallback_token_prefix) for item in vocabulary):
            raise ValueError("vocabulary token ID collides with byte fallback namespace")
        ranks = [item.rank for item in merges]
        pairs = [(item.left_hex, item.right_hex) for item in merges]
        if len(ranks) != len(set(ranks)) or sorted(ranks) != list(range(len(ranks))):
            raise ValueError("merge ranks must be unique and contiguous from zero")
        if len(pairs) != len(set(pairs)):
            raise ValueError("merge pairs must be unique")
        vocabulary_pieces = set(pieces)
        available = {f"{byte:02x}" for byte in range(256)}
        for rule in sorted(merges, key=lambda item: item.rank):
            if rule.left_hex not in available or rule.right_hex not in available:
                raise ValueError("merge inputs must be available at their declared rank")
            if rule.merged_hex not in vocabulary_pieces:
                raise ValueError("every merge output must have a vocabulary token")
            available.add(rule.merged_hex)
        object.__setattr__(self, "vocabulary", vocabulary)
        object.__setattr__(self, "merges", merges)

    @property
    def content_id(self) -> str:
        return _content_id("tokenizer-contract", asdict(self))


@dataclass(frozen=True)
class TextEvidence:
    evidence_id: str
    evidence_content_id: str
    contract_content_id: str
    source_id: str
    source_version: str
    observed_at: str
    scope: str
    text: str
    input_hex: str

    def __post_init__(self) -> None:
        for field in ("evidence_id", "source_id", "source_version", "scope"):
            _identifier(field, getattr(self, field))
        for field in ("evidence_content_id", "contract_content_id"):
            if type(getattr(self, field)) is not str or SHA256_ID.fullmatch(getattr(self, field)) is None:
                raise ValueError(f"{field} must be a content identity")
        _timestamp("observed_at", self.observed_at)
        if type(self.text) is not str:
            raise TypeError("text must be a concrete string")
        _hex_piece("input_hex", self.input_hex, maximum_bytes=MAX_TEXT_BYTES)

    @classmethod
    def capture(
        cls,
        *,
        contract: TokenizerContract,
        evidence_id: str,
        source_id: str,
        source_version: str,
        observed_at: str,
        scope: str,
        text: str,
    ) -> "TextEvidence":
        if type(contract) is not TokenizerContract:
            raise TypeError("contract must be a concrete TokenizerContract")
        contract.__post_init__()
        if type(text) is not str or not text:
            raise ValueError("text must be a non-empty concrete string")
        encoded = text.encode("utf-8", errors="strict")
        if len(encoded) > contract.maximum_text_bytes:
            raise ValueError("text exceeds maximum_text_bytes")
        checked_scope = _identifier("scope", scope)
        if checked_scope != contract.scope:
            raise ValueError("evidence scope does not match contract scope")
        payload = {
            "contract_content_id": contract.content_id,
            "evidence_id": _identifier("evidence_id", evidence_id),
            "source_id": _identifier("source_id", source_id),
            "source_version": _identifier("source_version", source_version),
            "observed_at": _timestamp("observed_at", observed_at),
            "scope": checked_scope,
            "text": text,
            "input_hex": encoded.hex(),
        }
        return cls(
            evidence_content_id=_content_id("tokenization-evidence", payload),
            **payload,
        )


@dataclass(frozen=True)
class TokenizationAudit:
    decision: str
    contract_content_id: str
    evidence_content_id: str
    input_bytes: int
    token_ids: tuple[str, ...]
    token_pieces_hex: tuple[str, ...]
    fallback_tokens: int
    bytes_per_token: float
    round_trip: str
    compression_claim: str
    material_evidence_id: str


def _revalidate(contract: TokenizerContract, evidence: TextEvidence) -> None:
    if type(contract) is not TokenizerContract or type(evidence) is not TextEvidence:
        raise TypeError("audit requires concrete contract and evidence records")
    contract.__post_init__()
    evidence.__post_init__()
    if evidence.contract_content_id != contract.content_id:
        raise ValueError("evidence belongs to a different tokenizer contract")
    if evidence.scope != contract.scope:
        raise ValueError("evidence scope does not match contract scope")
    encoded = evidence.text.encode("utf-8", errors="strict")
    if not encoded or len(encoded) > contract.maximum_text_bytes:
        raise ValueError("evidence text is empty or outside the byte bound")
    if encoded.hex() != evidence.input_hex:
        raise ValueError("text and captured input bytes differ")
    payload = {
        "contract_content_id": evidence.contract_content_id,
        "evidence_id": evidence.evidence_id,
        "source_id": evidence.source_id,
        "source_version": evidence.source_version,
        "observed_at": evidence.observed_at,
        "scope": evidence.scope,
        "text": evidence.text,
        "input_hex": evidence.input_hex,
    }
    if evidence.evidence_content_id != _content_id("tokenization-evidence", payload):
        raise ValueError("evidence content identity does not match its contents")


def audit_tokenization(
    contract: TokenizerContract, evidence: TextEvidence
) -> TokenizationAudit:
    """Encode, decode, and bind a deterministic compression observation."""

    _revalidate(contract, evidence)
    pieces = [f"{byte:02x}" for byte in bytes.fromhex(evidence.input_hex)]
    rules = {(rule.left_hex, rule.right_hex): rule for rule in contract.merges}
    while True:
        candidates = [
            (rules[(pieces[index], pieces[index + 1])].rank, index)
            for index in range(len(pieces) - 1)
            if (pieces[index], pieces[index + 1]) in rules
        ]
        if not candidates:
            break
        _, index = min(candidates)
        pieces[index : index + 2] = [pieces[index] + pieces[index + 1]]

    vocabulary = {item.piece_hex: item.token_id for item in contract.vocabulary}
    token_ids: list[str] = []
    fallback_tokens = 0
    for piece in pieces:
        if piece in vocabulary:
            token_ids.append(vocabulary[piece])
        elif len(piece) == 2:
            token_ids.append(f"{contract.fallback_token_prefix}{piece}")
            fallback_tokens += 1
        else:
            raise ValueError("merged piece is absent from the bound vocabulary")

    decoded = bytearray()
    reverse_vocabulary = {item.token_id: item.piece_hex for item in contract.vocabulary}
    for token_id in token_ids:
        if token_id in reverse_vocabulary:
            decoded.extend(bytes.fromhex(reverse_vocabulary[token_id]))
        elif token_id.startswith(contract.fallback_token_prefix):
            byte_hex = token_id[len(contract.fallback_token_prefix) :]
            if len(byte_hex) != 2 or re.fullmatch(r"[0-9a-f]{2}", byte_hex) is None:
                raise ValueError("malformed byte fallback token")
            decoded.extend(bytes.fromhex(byte_hex))
        else:
            raise ValueError("unknown token violates no-unknown-token policy")
    if decoded.hex() != evidence.input_hex:
        raise AssertionError("round-trip invariant failed")

    input_bytes = len(decoded)
    bytes_per_token = input_bytes / len(token_ids)
    material = {
        "contract_content_id": contract.content_id,
        "evidence_content_id": evidence.evidence_content_id,
        "input_bytes": input_bytes,
        "token_ids": tuple(token_ids),
        "token_pieces_hex": tuple(pieces),
        "fallback_tokens": fallback_tokens,
        "round_trip": "EXACT_UTF8_BYTES",
    }
    return TokenizationAudit(
        decision="PASS",
        contract_content_id=contract.content_id,
        evidence_content_id=evidence.evidence_content_id,
        input_bytes=input_bytes,
        token_ids=tuple(token_ids),
        token_pieces_hex=tuple(pieces),
        fallback_tokens=fallback_tokens,
        bytes_per_token=bytes_per_token,
        round_trip="EXACT_UTF8_BYTES",
        compression_claim="OBSERVED_FOR_THIS_EVIDENCE_ONLY",
        material_evidence_id=_content_id("tokenization-material", material),
    )


ILLUSTRATIVE_CONTRACT = TokenizerContract(
    contract_version="tokenizer-contract-v1",
    tokenizer_id="illustrative-byte-bpe",
    tokenizer_version="v1",
    normalization="none-codepoints-preserved-v1",
    byte_encoding="utf-8-strict-v1",
    merge_semantics="lowest-rank-leftmost-repeat-v1",
    unknown_policy="no-unknown-token-v1",
    fallback_policy="all-raw-bytes-v1",
    fallback_token_prefix="byte-",
    maximum_text_bytes=64,
    vocabulary=(
        TokenDefinition("piece-ba", "6261"),
        TokenDefinition("piece-na", "6e61"),
        TokenDefinition("piece-nana", "6e616e61"),
    ),
    merges=(
        MergeRule(0, "62", "61", "6261"),
        MergeRule(1, "6e", "61", "6e61"),
        MergeRule(2, "6e61", "6e61", "6e616e61"),
    ),
    scope="academy:foundation-model-internals",
    vocabulary_owner="team:tokenization",
    audit_owner="team:model-platform",
)

ILLUSTRATIVE_EVIDENCE = TextEvidence.capture(
    contract=ILLUSTRATIVE_CONTRACT,
    evidence_id="sample-banana",
    source_id="academy-fixture",
    source_version="fixture-v1",
    observed_at="2026-08-25T00:00:00+00:00",
    scope="academy:foundation-model-internals",
    text="banana!",
)


def format_example() -> str:
    report = audit_tokenization(ILLUSTRATIVE_CONTRACT, ILLUSTRATIVE_EVIDENCE)
    return "\n".join(
        (
            "example=illustrative_only",
            f"contract_id={report.contract_content_id}",
            f"evidence_id={report.evidence_content_id}",
            f"tokens={','.join(report.token_ids)}",
            f"compression={report.input_bytes}_bytes/{len(report.token_ids)}_tokens",
            f"fallback_tokens={report.fallback_tokens}",
            f"round_trip={report.round_trip}",
            f"claim={report.compression_claim}",
            f"decision={report.decision}",
        )
    )


if __name__ == "__main__":
    print(format_example())