"""Message types and payload models."""

from __future__ import annotations

from enum import Enum

from pydantic import BaseModel, Field


class MessageType(str, Enum):
    # Core
    TEXT = "text"
    STRUCTURED = "structured"

    # Collaboration
    TASK_REQUEST = "task_request"
    TASK_RESPONSE = "task_response"
    QUERY = "query"
    QUERY_RESPONSE = "query_response"

    # File transfer
    FILE_OFFER = "file_offer"
    FILE_CHUNK = "file_chunk"
    FILE_ACK = "file_ack"

    # Control
    PING = "ping"
    PONG = "pong"
    ACK = "ack"
    STATUS = "status"

    # Handshake
    HELLO = "hello"
    CHALLENGE = "challenge"
    CHALLENGE_RESPONSE = "challenge_response"
    APPROVAL_REQUEST = "approval_request"
    APPROVAL_GRANT = "approval_grant"
    APPROVAL_DENY = "approval_deny"


# ── Payloads ────────────────────────────────────────────────────────


class TextPayload(BaseModel):
    text: str
    context: str = ""


class StructuredPayload(BaseModel):
    content_type: str = "json"
    data: dict = Field(default_factory=dict)


class TaskRequestPayload(BaseModel):
    task: str
    context: str = ""
    constraints: list[str] = Field(default_factory=list)
    timeout_seconds: float = 300.0


class TaskResponsePayload(BaseModel):
    status: str  # completed, failed, partial, rejected
    result: str = ""
    error: str = ""


class QueryPayload(BaseModel):
    text: str
    context: str = ""


class QueryResponsePayload(BaseModel):
    text: str
    confidence: float = 0.0


class FileOfferPayload(BaseModel):
    filename: str
    size_bytes: int
    mime_type: str = "application/octet-stream"
    sha256: str = ""
    chunk_size: int = 65536


class FileChunkPayload(BaseModel):
    offer_id: str
    index: int
    total_chunks: int
    data_b64: str


class FileAckPayload(BaseModel):
    offer_id: str
    status: str = "complete"
    error: str = ""


class StatusPayload(BaseModel):
    status: str  # online, busy, offline_soon
    message: str = ""


class HelloPayload(BaseModel):
    agent_name: str
    owner_display_name: str = ""
    capabilities: list[str] = Field(default_factory=list)
    protocol_version: int = 1
    introduction: str = ""
    ed25519_pubkey_hex: str = ""  # Ed25519 public key for identity verification


class ChallengePayload(BaseModel):
    nonce_hex: str
    timestamp: float


class ChallengeResponsePayload(BaseModel):
    nonce_hex: str
    signature_hex: str


class ApprovalGrantPayload(BaseModel):
    trust_level: int = 1
    scope_instructions: str = ""


class ApprovalDenyPayload(BaseModel):
    reason: str = ""


PAYLOAD_MODELS: dict[MessageType, type[BaseModel] | None] = {
    MessageType.TEXT: TextPayload,
    MessageType.STRUCTURED: StructuredPayload,
    MessageType.TASK_REQUEST: TaskRequestPayload,
    MessageType.TASK_RESPONSE: TaskResponsePayload,
    MessageType.QUERY: QueryPayload,
    MessageType.QUERY_RESPONSE: QueryResponsePayload,
    MessageType.FILE_OFFER: FileOfferPayload,
    MessageType.FILE_CHUNK: FileChunkPayload,
    MessageType.FILE_ACK: FileAckPayload,
    MessageType.PING: None,
    MessageType.PONG: None,
    MessageType.ACK: None,
    MessageType.STATUS: StatusPayload,
    MessageType.HELLO: HelloPayload,
    MessageType.CHALLENGE: ChallengePayload,
    MessageType.CHALLENGE_RESPONSE: ChallengeResponsePayload,
    MessageType.APPROVAL_REQUEST: None,
    MessageType.APPROVAL_GRANT: ApprovalGrantPayload,
    MessageType.APPROVAL_DENY: ApprovalDenyPayload,
}
