"""Audit separation between durable state, memory, and model context.

The records and thresholds in this module are illustrative local fixtures.  The
audit does not persist a workflow, authenticate a tenant, run retrieval, count
real tokenizer output, measure memory quality, or certify crash recovery.  A
production system must supply those controls and evidence independently.
"""

from __future__ import annotations

from dataclasses import asdict, dataclass
from datetime import datetime, timezone
from hashlib import sha256
import json
from math import isfinite
import re
from typing import Literal


RunStatus = Literal[
    "ready",
    "awaiting_approval",
    "awaiting_execution",
    "succeeded",
    "failed",
    "cancelled",
]
MemoryKind = Literal["episodic", "semantic", "preference"]
SourceKind = Literal["state", "memory", "instruction"]
ContextAuthority = Literal["control", "evidence", "instruction"]

ID_PATTERN = re.compile(r"^[a-z][a-z0-9._:-]{2,95}$")
KEY_PATTERN = re.compile(r"^[a-z][a-z0-9_]{1,63}$")
UTC_PATTERN = re.compile(r"^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}Z$")


def _exact_text(name: str, value: object, *, maximum: int = 1024) -> str:
    if type(value) is not str or not value.strip():
        raise ValueError(f"{name} must be non-empty exact str")
    if value != value.strip():
        raise ValueError(f"{name} must not have surrounding whitespace")
    if len(value) > maximum or any(ord(character) < 32 for character in value):
        raise ValueError(f"{name} is outside the bounded text contract")
    return value


def _canonical_id(name: str, value: object) -> str:
    text = _exact_text(name, value, maximum=96)
    if not ID_PATTERN.fullmatch(text):
        raise ValueError(f"{name} must be a canonical identifier")
    return text


def _exact_int(name: str, value: object, *, minimum: int, maximum: int) -> int:
    if type(value) is not int:
        raise ValueError(f"{name} must be exact int")
    if not minimum <= value <= maximum:
        raise ValueError(f"{name} must be between {minimum} and {maximum}")
    return value


def _finite_probability(name: str, value: object) -> float:
    if type(value) not in {int, float}:
        raise ValueError(f"{name} must be a non-boolean number")
    numeric = float(value)
    if not isfinite(numeric):
        raise ValueError(f"{name} must be finite")
    if not 0 <= numeric <= 1:
        raise ValueError(f"{name} must be between zero and one")
    return numeric


def _utc_timestamp(name: str, value: object) -> str:
    text = _exact_text(name, value, maximum=20)
    if not UTC_PATTERN.fullmatch(text):
        raise ValueError(f"{name} must be canonical UTC seconds")
    try:
        parsed = datetime.fromisoformat(text.replace("Z", "+00:00"))
    except ValueError as error:
        raise ValueError(f"{name} must be a real timestamp") from error
    if parsed.tzinfo != timezone.utc:
        raise ValueError(f"{name} must be UTC")
    return text


def _stable_digest(payload: object) -> str:
    encoded = json.dumps(
        payload,
        allow_nan=False,
        ensure_ascii=True,
        separators=(",", ":"),
        sort_keys=True,
    ).encode("utf-8")
    return sha256(encoded).hexdigest()


@dataclass(frozen=True)
class StateField:
    key: str
    value: str

    def __post_init__(self) -> None:
        _exact_text("state_field.key", self.key, maximum=64)
        if not KEY_PATTERN.fullmatch(self.key):
            raise ValueError("state_field.key must be canonical snake_case")
        _exact_text(f"state_field.{self.key}", self.value, maximum=256)


@dataclass(frozen=True)
class WorkflowState:
    tenant_id: str
    run_id: str
    revision: int
    status: RunStatus
    next_action: str
    remaining_steps: int
    last_event_id: str
    fields: tuple[StateField, ...]

    def __post_init__(self) -> None:
        _canonical_id("state.tenant_id", self.tenant_id)
        _canonical_id("state.run_id", self.run_id)
        _exact_int("state.revision", self.revision, minimum=1, maximum=1_000_000_000)
        _exact_text("state.status", self.status, maximum=32)
        if self.status not in {
            "ready",
            "awaiting_approval",
            "awaiting_execution",
            "succeeded",
            "failed",
            "cancelled",
        }:
            raise ValueError("state.status is unsupported")
        _canonical_id("state.next_action", self.next_action)
        _exact_int("state.remaining_steps", self.remaining_steps, minimum=0, maximum=10_000)
        _canonical_id("state.last_event_id", self.last_event_id)
        try:
            fields = tuple(self.fields)
        except TypeError as error:
            raise ValueError("state.fields must be iterable") from error
        object.__setattr__(self, "fields", fields)
        if any(type(field) is not StateField for field in fields):
            raise ValueError("state.fields must contain exact StateField records")
        keys = [field.key for field in fields]
        if len(keys) != len(set(keys)):
            raise ValueError("state field keys must be unique")
        if tuple(sorted(keys)) != tuple(keys):
            raise ValueError("state fields must be sorted by key")
        if self.status in {"succeeded", "failed", "cancelled"} and self.remaining_steps != 0:
            raise ValueError("terminal state must have zero remaining_steps")

    @property
    def content_id(self) -> str:
        return "workflow-state@sha256:" + _stable_digest(asdict(self))


@dataclass(frozen=True)
class MemoryRecord:
    tenant_id: str
    memory_id: str
    kind: MemoryKind
    text: str
    source_event_id: str
    observed_at: str
    confidence: float
    authority: Literal["evidence"] = "evidence"

    def __post_init__(self) -> None:
        _canonical_id("memory.tenant_id", self.tenant_id)
        _canonical_id("memory.memory_id", self.memory_id)
        _exact_text("memory.kind", self.kind, maximum=16)
        if self.kind not in {"episodic", "semantic", "preference"}:
            raise ValueError("memory.kind is unsupported")
        _exact_text("memory.text", self.text, maximum=1024)
        _canonical_id("memory.source_event_id", self.source_event_id)
        _utc_timestamp("memory.observed_at", self.observed_at)
        confidence = _finite_probability("memory.confidence", self.confidence)
        object.__setattr__(self, "confidence", confidence)
        if self.authority != "evidence" or type(self.authority) is not str:
            raise ValueError("memory.authority must be exact evidence")


@dataclass(frozen=True)
class ContextItem:
    source_kind: SourceKind
    source_id: str
    text: str
    declared_tokens: int
    authority: ContextAuthority

    def __post_init__(self) -> None:
        if self.source_kind not in {"state", "memory", "instruction"}:
            raise ValueError("context.source_kind is unsupported")
        _exact_text("context.source_kind", self.source_kind, maximum=16)
        _exact_text("context.source_id", self.source_id, maximum=96)
        _exact_text("context.text", self.text, maximum=2048)
        _exact_int("context.declared_tokens", self.declared_tokens, minimum=1, maximum=100_000)
        if self.authority not in {"control", "evidence", "instruction"}:
            raise ValueError("context.authority is unsupported")
        _exact_text("context.authority", self.authority, maximum=16)


@dataclass(frozen=True)
class AgentSnapshot:
    tenant_id: str
    state: WorkflowState
    state_content_id: str
    memories: tuple[MemoryRecord, ...]
    context: tuple[ContextItem, ...]
    max_context_tokens: int

    def __post_init__(self) -> None:
        _canonical_id("snapshot.tenant_id", self.tenant_id)
        if type(self.state) is not WorkflowState:
            raise ValueError("snapshot.state must be exact WorkflowState")
        _exact_text("snapshot.state_content_id", self.state_content_id, maximum=96)
        if not re.fullmatch(r"workflow-state@sha256:[0-9a-f]{64}", self.state_content_id):
            raise ValueError("snapshot.state_content_id must be a state digest")
        try:
            memories = tuple(self.memories)
            context = tuple(self.context)
        except TypeError as error:
            raise ValueError("snapshot collections must be iterable") from error
        object.__setattr__(self, "memories", memories)
        object.__setattr__(self, "context", context)
        if any(type(memory) is not MemoryRecord for memory in memories):
            raise ValueError("snapshot.memories must contain exact MemoryRecord records")
        if any(type(item) is not ContextItem for item in context):
            raise ValueError("snapshot.context must contain exact ContextItem records")
        memory_ids = [memory.memory_id for memory in memories]
        if len(memory_ids) != len(set(memory_ids)):
            raise ValueError("memory IDs must be unique")
        context_sources = [(item.source_kind, item.source_id) for item in context]
        if len(context_sources) != len(set(context_sources)):
            raise ValueError("context sources must be unique")
        _exact_int(
            "snapshot.max_context_tokens",
            self.max_context_tokens,
            minimum=1,
            maximum=1_000_000,
        )


@dataclass(frozen=True)
class AuditResult:
    snapshot: AgentSnapshot
    declared_context_tokens: int
    selected_memory_count: int
    decision: Literal["PASS_SEPARATION_INVARIANTS"]


def _rebuild_state_field(value: object) -> StateField:
    if type(value) is not StateField:
        raise ValueError("state.fields must contain exact StateField records")
    return StateField(value.key, value.value)


def _rebuild_state(value: object) -> WorkflowState:
    if type(value) is not WorkflowState:
        raise ValueError("snapshot.state must be exact WorkflowState")
    try:
        fields = tuple(_rebuild_state_field(field) for field in value.fields)
        return WorkflowState(
            value.tenant_id,
            value.run_id,
            value.revision,
            value.status,
            value.next_action,
            value.remaining_steps,
            value.last_event_id,
            fields,
        )
    except (AttributeError, TypeError) as error:
        raise ValueError("state failed reconstruction") from error


def _rebuild_memory(value: object) -> MemoryRecord:
    if type(value) is not MemoryRecord:
        raise ValueError("snapshot.memories must contain exact MemoryRecord records")
    try:
        return MemoryRecord(
            value.tenant_id,
            value.memory_id,
            value.kind,
            value.text,
            value.source_event_id,
            value.observed_at,
            value.confidence,
            value.authority,
        )
    except (AttributeError, TypeError) as error:
        raise ValueError("memory failed reconstruction") from error


def _rebuild_context(value: object) -> ContextItem:
    if type(value) is not ContextItem:
        raise ValueError("snapshot.context must contain exact ContextItem records")
    try:
        return ContextItem(
            value.source_kind,
            value.source_id,
            value.text,
            value.declared_tokens,
            value.authority,
        )
    except (AttributeError, TypeError) as error:
        raise ValueError("context item failed reconstruction") from error


def _rebuild_snapshot(value: object) -> AgentSnapshot:
    if type(value) is not AgentSnapshot:
        raise ValueError("snapshot must be exact AgentSnapshot")
    try:
        state = _rebuild_state(value.state)
        memories = tuple(_rebuild_memory(memory) for memory in value.memories)
        context = tuple(_rebuild_context(item) for item in value.context)
        return AgentSnapshot(
            value.tenant_id,
            state,
            value.state_content_id,
            memories,
            context,
            value.max_context_tokens,
        )
    except (AttributeError, TypeError) as error:
        raise ValueError("snapshot failed reconstruction") from error


def audit(snapshot: object) -> AuditResult:
    """Validate separation invariants without changing the supplied snapshot."""

    checked = _rebuild_snapshot(snapshot)
    if checked.state.tenant_id != checked.tenant_id:
        raise ValueError("state tenant does not match snapshot tenant")
    if checked.state_content_id != checked.state.content_id:
        raise ValueError("snapshot is not bound to current state content")
    for memory in checked.memories:
        if memory.tenant_id != checked.tenant_id:
            raise ValueError(f"memory {memory.memory_id} crosses tenant boundary")

    memories_by_id = {memory.memory_id: memory for memory in checked.memories}
    state_items = 0
    selected_memory_ids: set[str] = set()
    for item in checked.context:
        if item.source_kind == "state":
            state_items += 1
            if item.source_id != checked.state.content_id:
                raise ValueError("state context is not bound to current state content")
            if item.authority != "control":
                raise ValueError("state context must have control authority")
        elif item.source_kind == "memory":
            if item.source_id not in memories_by_id:
                raise ValueError("memory context references an unknown memory")
            if item.authority != "evidence":
                raise ValueError("memory context cannot have control authority")
            selected_memory_ids.add(item.source_id)
        else:
            if item.authority != "instruction":
                raise ValueError("instruction context must have instruction authority")
            _canonical_id("instruction source_id", item.source_id)
    if state_items != 1:
        raise ValueError("context must contain exactly one current state item")

    declared_tokens = sum(item.declared_tokens for item in checked.context)
    if declared_tokens > checked.max_context_tokens:
        raise ValueError("declared context token budget exceeded")
    return AuditResult(
        checked,
        declared_tokens,
        len(selected_memory_ids),
        "PASS_SEPARATION_INVARIANTS",
    )


ILLUSTRATIVE_STATE = WorkflowState(
    tenant_id="tenant_demo",
    run_id="run_0042",
    revision=7,
    status="awaiting_execution",
    next_action="create_support_credit",
    remaining_steps=2,
    last_event_id="event_0007",
    fields=(
        StateField("approval_status", "granted"),
        StateField("credit_request_id", "request_demo_42"),
    ),
)

ILLUSTRATIVE_MEMORY = MemoryRecord(
    tenant_id="tenant_demo",
    memory_id="memory_channel_preference",
    kind="preference",
    text="The illustrative customer previously preferred email updates.",
    source_event_id="event_0003",
    observed_at="2026-09-21T09:30:00Z",
    confidence=0.8,
)

ILLUSTRATIVE_SNAPSHOT = AgentSnapshot(
    tenant_id="tenant_demo",
    state=ILLUSTRATIVE_STATE,
    state_content_id=ILLUSTRATIVE_STATE.content_id,
    memories=(ILLUSTRATIVE_MEMORY,),
    context=(
        ContextItem(
            "state",
            ILLUSTRATIVE_STATE.content_id,
            "Run awaits execution; approval is granted in durable state.",
            11,
            "control",
        ),
        ContextItem(
            "memory",
            "memory_channel_preference",
            "Unverified preference evidence: email updates may be preferred.",
            9,
            "evidence",
        ),
        ContextItem(
            "instruction",
            "policy_support_v1",
            "Revalidate authority at the tool boundary before execution.",
            9,
            "instruction",
        ),
    ),
    max_context_tokens=64,
)


def format_example() -> str:
    result = audit(ILLUSTRATIVE_SNAPSHOT)
    return "\n".join(
        (
            "example=illustrative_only",
            f"state_id={result.snapshot.state.content_id}",
            "state=AUTHORITATIVE_DURABLE_RECORD",
            "memory=RETRIEVABLE_NON_AUTHORITY",
            "context=BOUNDED_MODEL_VIEW",
            f"declared_context_tokens={result.declared_context_tokens}/64",
            f"selected_memories={result.selected_memory_count}",
            f"decision={result.decision}",
            "claim=LOCAL_FIXTURE_ONLY",
            "certification=NOT_A_DURABILITY_OR_MEMORY_QUALITY_CERTIFICATION",
        )
    )


def main() -> None:
    print(format_example())


if __name__ == "__main__":
    main()