"""Core data models for AttestMCP audit records.

These schemas are the foundation of the evidence chain. Once records exist
in production, changing these models requires a migration strategy.
Design carefully — this is the hardest thing to change later.
"""

from __future__ import annotations

import uuid
from datetime import UTC, datetime
from enum import StrEnum

from pydantic import BaseModel, Field


class IdentitySource(StrEnum):
    """How the agent identity was obtained."""

    OAUTH_JWT = "oauth_jwt"
    API_KEY = "api_key"
    MTLS = "mtls"
    MCP_SESSION = "mcp_session"
    ANONYMOUS = "anonymous"


class TrustLevel(StrEnum):
    """How much we trust the identity claim.

    - verified: Cryptographically validated (JWT signature checked against JWKS)
    - claimed: Token present but signature not validated
    - inferred: Identity derived from session or context, not directly asserted
    - anonymous: No identity information available
    """

    VERIFIED = "verified"
    CLAIMED = "claimed"
    INFERRED = "inferred"
    ANONYMOUS = "anonymous"


class PolicyDecision(StrEnum):
    """Result of policy evaluation for a memory operation."""

    ALLOW = "allow"
    DENY = "deny"
    ESCALATE = "escalate"
    REDACT = "redact"


class AgentIdentity(BaseModel):
    """Identity of the agent or user performing a memory operation.

    This block appears on every audit record. Even anonymous requests
    get an identity block with source="anonymous" and trust_level="anonymous".

    The schema is designed to be extensible: future fields (like delegated_by
    for agent-to-agent delegation chains) can be added without breaking
    existing records because Pydantic ignores unknown fields on deserialization.
    """

    source: IdentitySource = Field(
        description="How the identity was obtained"
    )
    user: str | None = Field(
        default=None,
        description="Human user email or ID, if present in the token",
    )
    agent_name: str | None = Field(
        default=None,
        description="Agent identifier, if present",
    )
    org: str | None = Field(
        default=None,
        description="Organization or tenant identifier",
    )
    role: str | None = Field(
        default=None,
        description="Role claim from token or policy lookup",
    )
    session_id: str | None = Field(
        default=None,
        description="MCP session ID (Mcp-Session-Id header)",
    )
    verified: bool = Field(
        default=False,
        description="Whether the identity was cryptographically verified",
    )
    trust_level: TrustLevel = Field(
        default=TrustLevel.ANONYMOUS,
        description="How much we trust this identity assertion",
    )

    @classmethod
    def anonymous(cls, session_id: str | None = None) -> AgentIdentity:
        """Create an anonymous identity for requests without credentials."""
        return cls(
            source=IdentitySource.ANONYMOUS,
            trust_level=TrustLevel.ANONYMOUS,
            session_id=session_id,
        )


def _generate_record_id() -> str:
    """Generate a unique record ID.

    Uses UUID v4 for broad compatibility. Records are ordered by
    sequence_num in the database, not by UUID sorting.
    UUID v7 (time-sortable) can be adopted when targeting Python 3.14+.
    """
    return str(uuid.uuid4())


def _now_utc() -> datetime:
    """Current UTC time with microsecond precision."""
    return datetime.now(UTC)


class AuditRecord(BaseModel):
    """A single tamper-evident audit record for one MCP operation.

    Every intercepted JSON-RPC call produces exactly one AuditRecord.
    Records are append-only and hash-chained: each record's chain_hash
    is SHA-256(previous_chain_hash + canonical_json(this_record_without_chain_hash)).

    The chain_hash field is computed after all other fields are set,
    so it is excluded from the canonical serialization used for hashing.
    """

    record_id: str = Field(
        default_factory=_generate_record_id,
        description="Time-sortable unique ID (UUID v7)",
    )
    timestamp: datetime = Field(
        default_factory=_now_utc,
        description="UTC timestamp with microsecond precision",
    )
    session_id: str | None = Field(
        default=None,
        description="MCP session ID from Mcp-Session-Id header",
    )
    agent_identity: AgentIdentity = Field(
        description="Verified or claimed identity of the requester",
    )
    method: str = Field(
        description="JSON-RPC method (e.g., tools/call, tools/list, resources/read)",
    )
    tool_name: str | None = Field(
        default=None,
        description="Tool name for tools/call requests (e.g., store_memory, recall)",
    )
    params_hash: str = Field(
        description="SHA-256 hash of the canonical JSON of request parameters",
    )
    policy_decision: PolicyDecision = Field(
        default=PolicyDecision.ALLOW,
        description="Result of policy evaluation",
    )
    policy_rule: str | None = Field(
        default=None,
        description="OPA rule or policy ID that triggered the decision",
    )
    result_hash: str | None = Field(
        default=None,
        description="SHA-256 hash of the upstream response body (None if denied)",
    )
    duration_ms: int = Field(
        default=0,
        description="Round-trip time through the proxy in milliseconds",
    )
    error: str | None = Field(
        default=None,
        description="Error message if the operation failed (upstream error or policy denial)",
    )
    chain_hash: str = Field(
        default="",
        description="SHA-256(previous_chain_hash + canonical_json(this_record)). "
        "Computed after all other fields are set. Excluded from canonical form.",
    )


class ChainVerification(BaseModel):
    """Result of verifying the integrity of an audit chain."""

    valid: bool = Field(description="Whether the entire chain is intact")
    records_checked: int = Field(description="Number of records verified")
    first_record_id: str | None = Field(
        default=None, description="ID of the first record in the range"
    )
    last_record_id: str | None = Field(
        default=None, description="ID of the last record in the range"
    )
    break_at_record_id: str | None = Field(
        default=None,
        description="ID of the first record where the chain breaks (None if valid)",
    )
    expected_hash: str | None = Field(
        default=None,
        description="What the chain_hash should have been at the break point",
    )
    actual_hash: str | None = Field(
        default=None,
        description="What the chain_hash actually was at the break point",
    )
    verified_at: datetime = Field(
        default_factory=_now_utc,
        description="When this verification was performed",
    )

    @classmethod
    def empty(cls) -> ChainVerification:
        """Verification result for an empty chain (vacuously valid)."""
        return cls(valid=True, records_checked=0)
