"""WebSocket message protocol definitions."""

from dataclasses import dataclass, field
from datetime import UTC, datetime
from enum import Enum
from typing import Any


class MessageType(str, Enum):
    """WebSocket message types."""

    # Client → Server
    CHAT_SEND = "chat.send"
    CHAT_CANCEL = "chat.cancel"
    APPROVAL_RESPOND = "approval.respond"
    SESSION_SET_APPROVAL_POLICY = "session.set_approval_policy"
    AUDIT_REPLAY = "audit.replay"

    # Server → Client
    CHAT_TOKEN = "chat.token"
    CHAT_MESSAGE_START = "chat.message_start"
    CHAT_MESSAGE_END = "chat.message_end"
    TOOL_START = "tool.start"
    TOOL_SUCCESS = "tool.success"
    TOOL_ERROR = "tool.error"
    TOOL_REJECTED = "tool.rejected"
    APPROVAL_REQUEST = "approval.request"
    FILE_CHANGED = "file.changed"
    SYSTEM_MESSAGE = "system.message"
    SESSION_POLICY_UPDATED = "session.policy_updated"
    AUDIT_REPLAY_RESULT = "audit.replay_result"
    ERROR = "error"

    # Stream mode events
    STREAM_STATE = "stream.state"  # values mode - full state snapshot
    STREAM_DEBUG = "stream.debug"  # debug mode - detailed execution info
    STREAM_NODE_START = "stream.node_start"  # node execution start
    STREAM_NODE_END = "stream.node_end"  # node execution end

    # Thinking events (for reasoning models like GPT-OSS, Qwen)
    THINKING_START = "thinking.start"
    THINKING_TOKEN = "thinking.token"
    THINKING_END = "thinking.end"

    # Interview (ask_user_question)
    INTERVIEW_QUESTIONS = "interview.questions"  # Server → Client
    INTERVIEW_ANSWER = "interview.answer"  # Client → Server
    INTERVIEW_CANCEL = "interview.cancel"  # Client → Server
    INTERVIEW_COMPLETE = "interview.complete"  # Server → Client

    # Plan mode
    PLAN_SET_MODE = "plan.set_mode"  # Client → Server
    PLAN_APPROVE = "plan.approve"  # Client → Server
    PLAN_MODE_CHANGED = "plan.mode_changed"  # Server → Client

    # Background tasks
    TASK_SUBMIT = "task.submit"  # Client → Server
    TASK_CANCEL = "task.cancel"  # Client → Server
    TASK_CREATED = "task.created"  # Server → Client
    TASK_PROGRESS = "task.progress"  # Server → Client
    TASK_COMPLETED = "task.completed"  # Server → Client
    TASK_FAILED = "task.failed"  # Server → Client
    TASK_CANCELLED = "task.cancelled"  # Server → Client

    # Connection management
    CONNECTED = "connected"
    DISCONNECTED = "disconnected"
    PING = "ping"
    PONG = "pong"


MESSAGE_TYPE_ALIASES: dict[str, MessageType] = {
    # Backward compatibility for older web clients
    "chat.interrupt": MessageType.CHAT_CANCEL,
}


@dataclass
class WSMessage:
    """WebSocket message structure."""

    type: MessageType
    data: dict[str, Any] = field(default_factory=dict)
    timestamp: str = field(default_factory=lambda: datetime.now(UTC).isoformat())
    id: str | None = None

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary for JSON serialization."""
        return {
            "type": self.type.value,
            "data": self.data,
            "timestamp": self.timestamp,
            "id": self.id,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "WSMessage":
        """Create from dictionary."""
        msg_type = data.get("type", "")
        type_enum = MESSAGE_TYPE_ALIASES.get(msg_type)
        if type_enum is None:
            try:
                type_enum = MessageType(msg_type)
            except ValueError:
                type_enum = MessageType.ERROR

        return cls(
            type=type_enum,
            data=data.get("data", {}),
            timestamp=data.get("timestamp", datetime.now(UTC).isoformat()),
            id=data.get("id"),
        )


# Convenience constructors
def chat_token(content: str, message_id: str | None = None) -> WSMessage:
    """Create a chat token message."""
    return WSMessage(
        type=MessageType.CHAT_TOKEN,
        data={"content": content},
        id=message_id,
    )


def chat_message_start(message_id: str) -> WSMessage:
    """Create a chat message start message."""
    return WSMessage(
        type=MessageType.CHAT_MESSAGE_START,
        data={"message_id": message_id},
        id=message_id,
    )


def chat_message_end(
    message_id: str,
    tokens: int = 0,
    cost: float = 0.0,
    session_cost: float = 0.0,
) -> WSMessage:
    """Create a chat message end message."""
    data: dict[str, Any] = {"message_id": message_id, "tokens": tokens}
    if cost > 0:
        data["cost"] = round(cost, 6)
    if session_cost > 0:
        data["session_cost"] = round(session_cost, 6)
    return WSMessage(
        type=MessageType.CHAT_MESSAGE_END,
        data=data,
        id=message_id,
    )


def tool_start(tool_id: str, tool_name: str, args: dict[str, Any]) -> WSMessage:
    """Create a tool start message."""
    return WSMessage(
        type=MessageType.TOOL_START,
        data={
            "tool_id": tool_id,
            "tool_name": tool_name,
            "args": args,
        },
        id=tool_id,
    )


def tool_success(tool_id: str, output: str) -> WSMessage:
    """Create a tool success message."""
    return WSMessage(
        type=MessageType.TOOL_SUCCESS,
        data={
            "tool_id": tool_id,
            "output": output,
        },
        id=tool_id,
    )


def tool_error(tool_id: str, error: str) -> WSMessage:
    """Create a tool error message."""
    return WSMessage(
        type=MessageType.TOOL_ERROR,
        data={
            "tool_id": tool_id,
            "error": error,
        },
        id=tool_id,
    )


def tool_rejected(tool_id: str) -> WSMessage:
    """Create a tool rejected message."""
    return WSMessage(
        type=MessageType.TOOL_REJECTED,
        data={"tool_id": tool_id},
        id=tool_id,
    )


def approval_request(
    interrupt_id: str,
    action_requests: list[dict[str, Any]],
    assistant_id: str | None = None,
) -> WSMessage:
    """Create an approval request message."""
    return WSMessage(
        type=MessageType.APPROVAL_REQUEST,
        data={
            "interrupt_id": interrupt_id,
            "action_requests": action_requests,
            "assistant_id": assistant_id,
        },
        id=interrupt_id,
    )


def file_changed(
    path: str,
    action: str,
    diff: str | None = None,
) -> WSMessage:
    """Create a file changed message."""
    return WSMessage(
        type=MessageType.FILE_CHANGED,
        data={
            "path": path,
            "action": action,
            "diff": diff,
        },
    )


def system_message(content: str) -> WSMessage:
    """Create a system message."""
    return WSMessage(
        type=MessageType.SYSTEM_MESSAGE,
        data={"content": content},
    )


def error_message(error: str, code: str | None = None) -> WSMessage:
    """Create an error message."""
    return WSMessage(
        type=MessageType.ERROR,
        data={
            "error": error,
            "message": error,
            "code": code,
        },
    )


def connected(session_id: str) -> WSMessage:
    """Create a connected message."""
    return WSMessage(
        type=MessageType.CONNECTED,
        data={"session_id": session_id},
    )


def session_policy_updated(
    approval_policy: str,
) -> WSMessage:
    """Create a session policy update event."""
    return WSMessage(
        type=MessageType.SESSION_POLICY_UPDATED,
        data={"approval_policy": approval_policy},
    )


def stream_state(node: str, state: dict[str, Any]) -> WSMessage:
    """Create a stream state message (values mode)."""
    return WSMessage(
        type=MessageType.STREAM_STATE,
        data={
            "node": node,
            "state": state,
        },
    )


def stream_debug(
    event_type: str,
    node: str | None = None,
    data: dict[str, Any] | None = None,
) -> WSMessage:
    """Create a stream debug message."""
    return WSMessage(
        type=MessageType.STREAM_DEBUG,
        data={
            "event_type": event_type,
            "node": node,
            "data": data or {},
        },
    )


def stream_node_start(node: str, input_data: dict[str, Any] | None = None) -> WSMessage:
    """Create a node start message."""
    return WSMessage(
        type=MessageType.STREAM_NODE_START,
        data={
            "node": node,
            "input": input_data or {},
        },
    )


def stream_node_end(node: str, output_data: dict[str, Any] | None = None) -> WSMessage:
    """Create a node end message."""
    return WSMessage(
        type=MessageType.STREAM_NODE_END,
        data={
            "node": node,
            "output": output_data or {},
        },
    )


def interview_questions(
    interview_id: str,
    questions: list[dict[str, Any]],
    skippable: bool = True,
) -> WSMessage:
    """Create an interview questions message."""
    return WSMessage(
        type=MessageType.INTERVIEW_QUESTIONS,
        data={
            "interview_id": interview_id,
            "questions": questions,
            "skippable": skippable,
        },
        id=interview_id,
    )


def interview_complete(interview_id: str) -> WSMessage:
    """Create an interview complete message."""
    return WSMessage(
        type=MessageType.INTERVIEW_COMPLETE,
        data={"interview_id": interview_id},
        id=interview_id,
    )


def thinking_start(message_id: str | None = None) -> WSMessage:
    """Create a thinking start message (for reasoning models)."""
    return WSMessage(
        type=MessageType.THINKING_START,
        data={"message_id": message_id},
        id=message_id,
    )


def thinking_token(content: str, message_id: str | None = None) -> WSMessage:
    """Create a thinking token message."""
    return WSMessage(
        type=MessageType.THINKING_TOKEN,
        data={"content": content},
        id=message_id,
    )


def thinking_end(message_id: str | None = None) -> WSMessage:
    """Create a thinking end message."""
    return WSMessage(
        type=MessageType.THINKING_END,
        data={"message_id": message_id},
        id=message_id,
    )


def plan_mode_changed(
    plan_mode: bool,
    approval_policy: str,
    current_plan: str | None = None,
) -> WSMessage:
    """Create a plan mode changed message."""
    return WSMessage(
        type=MessageType.PLAN_MODE_CHANGED,
        data={
            "plan_mode": plan_mode,
            "approval_policy": approval_policy,
            "current_plan": current_plan,
        },
    )
