"""AICC protocol message types and enums."""

from __future__ import annotations

import uuid
from enum import Enum
from typing import Any, Literal

from pydantic import BaseModel, ConfigDict, Field

PROTOCOL_VERSION = "aicc/0.1"


class MessageType(str, Enum):
    SESSION_INIT = "session_init"
    SESSION_RESUME = "session_resume"
    SESSION_CLOSE = "session_close"
    MANIFEST_REQUEST = "manifest_request"
    TOOL_CALL = "tool_call"
    TOOL_RESULT = "tool_result"
    EVENT = "event"
    ERROR = "error"
    HEARTBEAT = "heartbeat"


class TickMode(str, Enum):
    FIXED = "fixed"
    EVENT = "event"
    HYBRID = "hybrid"


class WorldKind(str, Enum):
    TWO_D = "2d"
    THREE_D = "3d"
    TEXT = "text"
    ABSTRACT = "abstract"


class ModelClass(str, Enum):
    EDGE_SMALL = "edge_small"
    EDGE_MEDIUM = "edge_medium"
    CLOUD_MEDIUM = "cloud_medium"
    CLOUD_LARGE = "cloud_large"


class ToolClass(str, Enum):
    SENSOR = "sensor"
    ACTUATOR = "actuator"
    GENERATOR = "generator"


class ErrorCode(str, Enum):
    PROTOCOL_MISMATCH = "protocol_mismatch"
    SESSION_EXPIRED = "session_expired"
    TOOL_UNKNOWN = "tool_unknown"
    TOOL_UNAVAILABLE = "tool_unavailable"
    INVALID_INPUT = "invalid_input"
    EXECUTION_FAILED = "execution_failed"
    TIMEOUT = "timeout"
    INTERNAL_ERROR = "internal_error"


class EventTopic(str, Enum):
    TICK = "tick"
    COLLISION = "collision"
    AUDIO = "audio"
    STATE_CHANGE = "state_change"
    AGENT_MESSAGE = "agent_message"


def _uuid() -> str:
    return str(uuid.uuid4())


class Envelope(BaseModel):
    """Common envelope for all AICC messages."""

    model_config = ConfigDict(extra="forbid")

    protocol: Literal["aicc/0.1"] = PROTOCOL_VERSION
    session_id: str = Field(default_factory=_uuid)
    message_id: str = Field(default_factory=_uuid)


class WorldSpec(BaseModel):
    model_config = ConfigDict(extra="forbid")

    name: str
    kind: WorldKind


class ModelDecl(BaseModel):
    model_config = ConfigDict(extra="forbid", populate_by_name=True)

    class_: ModelClass = Field(alias="class")
    expected_first_token_ms: int | None = None
    expected_full_response_ms: int | None = None


class ToolLimits(BaseModel):
    model_config = ConfigDict(extra="forbid")

    calls_per_minute: int | None = None
    calls_per_session: int | None = None


class ToolSpec(BaseModel):
    """Tool definition advertised in the session manifest."""

    model_config = ConfigDict(extra="forbid", populate_by_name=True)

    id: str = Field(pattern=r"^[a-z][a-z0-9_]*$")
    class_: ToolClass = Field(alias="class")
    description: str
    input_schema: dict[str, Any]
    output_schema: dict[str, Any]
    requires_capability: str | None = None
    limits: ToolLimits | None = None
    strict_input: bool = False


class Capabilities(BaseModel):
    model_config = ConfigDict(extra="forbid")

    sensors: list[str] = Field(default_factory=list)
    actuators: list[str] = Field(default_factory=list)
    generators: list[str] = Field(default_factory=list)


class SessionInit(Envelope):
    type: Literal[MessageType.SESSION_INIT] = MessageType.SESSION_INIT
    tick_rate_hz: float = Field(gt=0)
    tick_mode: TickMode
    world: WorldSpec
    agent_model: ModelDecl
    capabilities: Capabilities
    tools: list[ToolSpec]
    event_subscriptions: list[str] = Field(default_factory=list)


class SessionClose(Envelope):
    type: Literal[MessageType.SESSION_CLOSE] = MessageType.SESSION_CLOSE
    reason: str | None = None


class ToolCall(Envelope):
    type: Literal[MessageType.TOOL_CALL] = MessageType.TOOL_CALL
    call_id: str
    tool: str = Field(pattern=r"^[a-z][a-z0-9_]*$")
    input: dict[str, Any] = Field(default_factory=dict)


class Meta(BaseModel):
    model_config = ConfigDict(extra="allow")

    tick: int | None = None
    latency_ms: int | None = None
    source: str | None = None


class ErrorPayload(BaseModel):
    model_config = ConfigDict(extra="forbid")

    code: ErrorCode
    message: str
    retryable: bool


class ToolResult(Envelope):
    type: Literal[MessageType.TOOL_RESULT] = MessageType.TOOL_RESULT
    call_id: str
    ok: bool
    output: dict[str, Any] | None = None
    error: ErrorPayload | None = None
    meta: Meta | None = None


class EventMessage(Envelope):
    type: Literal[MessageType.EVENT] = MessageType.EVENT
    topic: EventTopic | str
    payload: dict[str, Any]
    meta: Meta | None = None


class ErrorMessage(Envelope):
    type: Literal[MessageType.ERROR] = MessageType.ERROR
    error: ErrorPayload


class Heartbeat(Envelope):
    type: Literal[MessageType.HEARTBEAT] = MessageType.HEARTBEAT
    manifest_update: dict[str, Any] | None = None