"""Protocol types matching the agents-rendezvous WebSocket protocol."""

from __future__ import annotations

import enum
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, List, Literal, Optional, TypedDict


# ---------------------------------------------------------------------------
# Connection state
# ---------------------------------------------------------------------------

class ConnectionState(enum.Enum):
    DISCONNECTED = "DISCONNECTED"
    CONNECTING = "CONNECTING"
    CONNECTED = "CONNECTED"
    RECONNECTING = "RECONNECTING"


# ---------------------------------------------------------------------------
# Options
# ---------------------------------------------------------------------------

@dataclass
class AgentOptions:
    """Configuration options for AnycastAgent."""

    rendezvous_url: str
    """WebSocket URL of the rendezvous server."""

    token: str
    """Agent auth token from the Anycast portal."""

    machine_id: Optional[str] = None
    """Unique machine identifier (auto-generated if not provided)."""

    public_key: Optional[str] = None
    """Public key for end-to-end encryption (auto-generated if not provided)."""

    version: Optional[str] = None
    """Agent version string."""

    reconnect: bool = True
    """Auto-reconnect on disconnect."""

    reconnect_interval: float = 3.0
    """Base reconnect interval in seconds."""

    timeout: float = 30.0
    """Connection timeout in seconds."""

    heartbeat_interval: float = 25.0
    """Heartbeat interval in seconds."""


# ---------------------------------------------------------------------------
# Server → Client message types (as TypedDicts for JSON parsing)
# ---------------------------------------------------------------------------

class RegisteredMsg(TypedDict, total=False):
    type: Literal["registered"]
    agentId: str
    tenantId: str
    name: str
    popIata: str
    heartbeatIntervalMs: int


class HeartbeatAckMsg(TypedDict):
    type: Literal["heartbeat_ack"]
    timestamp: int


class ConnectIncomingMsg(TypedDict):
    type: Literal["connect_incoming"]
    sessionId: str
    fromAgentId: str
    fromAgentName: str


class ConnectAcceptedMsg(TypedDict):
    type: Literal["connect_accepted"]
    sessionId: str
    agentId: str
    agentName: str
    publicKey: str


class ConnectRejectedMsg(TypedDict):
    type: Literal["connect_rejected"]
    sessionId: str
    reason: str


class SignalData(TypedDict):
    type: Literal["offer", "answer", "ice"]
    data: str


class SignalForwardMsg(TypedDict):
    type: Literal["signal_forward"]
    sessionId: str
    fromAgentId: str
    signal: SignalData


class RelayAssignedMsg(TypedDict):
    type: Literal["relay_assigned"]
    sessionId: str
    relayHost: str
    relayPort: int
    relayToken: str


class AgentStatusMsg(TypedDict):
    type: Literal["agent_status"]
    agentId: str
    status: Literal["online", "offline"]


class ErrorMsg(TypedDict, total=False):
    type: Literal["error"]
    code: str
    message: str
    fatal: bool


# ---------------------------------------------------------------------------
# Agent info returned on connect
# ---------------------------------------------------------------------------

@dataclass
class AgentInfo:
    agent_id: str
    tenant_id: str
    name: str
    pop_iata: Optional[str] = None
