"""Tests for ieum.server.websocket.protocol -- WebSocket message protocol.

Covers:
- WSMessage dataclass construction and serialization
- MessageType enum values and string representations
- Convenience constructor functions produce correct message formats
- from_dict round-trips correctly
- Invalid message types fall back to ERROR
"""

import pytest

from ieum.server.websocket.protocol import (
    MessageType,
    WSMessage,
    approval_request,
    chat_message_end,
    chat_message_start,
    chat_token,
    connected,
    error_message,
    file_changed,
    interview_complete,
    interview_questions,
    session_policy_updated,
    stream_debug,
    stream_node_end,
    stream_node_start,
    stream_state,
    system_message,
    thinking_end,
    thinking_start,
    thinking_token,
    tool_error,
    tool_rejected,
    tool_start,
    tool_success,
)


# ---------------------------------------------------------------------------
# MessageType enum
# ---------------------------------------------------------------------------


class TestMessageType:
    """Test MessageType enum completeness and values."""

    def test_all_expected_types_exist(self):
        expected = [
            "CHAT_SEND",
            "CHAT_CANCEL",
            "APPROVAL_RESPOND",
            "SESSION_SET_APPROVAL_POLICY",
            "AUDIT_REPLAY",
            "CHAT_TOKEN",
            "CHAT_MESSAGE_START",
            "CHAT_MESSAGE_END",
            "TOOL_START",
            "TOOL_SUCCESS",
            "TOOL_ERROR",
            "TOOL_REJECTED",
            "APPROVAL_REQUEST",
            "FILE_CHANGED",
            "SYSTEM_MESSAGE",
            "SESSION_POLICY_UPDATED",
            "AUDIT_REPLAY_RESULT",
            "ERROR",
            "CONNECTED",
            "DISCONNECTED",
            "PING",
            "PONG",
            "INTERVIEW_QUESTIONS",
            "INTERVIEW_ANSWER",
            "INTERVIEW_CANCEL",
            "INTERVIEW_COMPLETE",
            "THINKING_START",
            "THINKING_TOKEN",
            "THINKING_END",
        ]
        for name in expected:
            assert hasattr(MessageType, name), f"Missing MessageType.{name}"

    def test_message_type_is_string_enum(self):
        assert isinstance(MessageType.CHAT_TOKEN.value, str)
        assert MessageType.CHAT_TOKEN.value == "chat.token"

    def test_message_type_dot_notation_values(self):
        assert MessageType.TOOL_START.value == "tool.start"
        assert MessageType.TOOL_SUCCESS.value == "tool.success"
        assert MessageType.TOOL_ERROR.value == "tool.error"
        assert MessageType.APPROVAL_REQUEST.value == "approval.request"


# ---------------------------------------------------------------------------
# WSMessage dataclass
# ---------------------------------------------------------------------------


class TestWSMessage:
    """Test WSMessage construction and serialization."""

    def test_basic_construction(self):
        msg = WSMessage(type=MessageType.CHAT_TOKEN, data={"content": "hello"})
        assert msg.type == MessageType.CHAT_TOKEN
        assert msg.data == {"content": "hello"}
        assert msg.timestamp is not None

    def test_to_dict(self):
        msg = WSMessage(
            type=MessageType.CHAT_TOKEN,
            data={"content": "hello"},
            id="msg-1",
        )
        d = msg.to_dict()
        assert d["type"] == "chat.token"
        assert d["data"]["content"] == "hello"
        assert d["id"] == "msg-1"
        assert "timestamp" in d

    def test_from_dict_roundtrip(self):
        original = WSMessage(
            type=MessageType.TOOL_START,
            data={"tool_id": "t1", "tool_name": "shell", "args": {}},
            id="t1",
        )
        d = original.to_dict()
        restored = WSMessage.from_dict(d)
        assert restored.type == MessageType.TOOL_START
        assert restored.data == original.data
        assert restored.id == "t1"

    def test_from_dict_unknown_type_defaults_to_error(self):
        d = {"type": "unknown.type", "data": {}}
        msg = WSMessage.from_dict(d)
        assert msg.type == MessageType.ERROR

    def test_from_dict_interrupt_alias_maps_to_cancel(self):
        d = {"type": "chat.interrupt", "data": {}}
        msg = WSMessage.from_dict(d)
        assert msg.type == MessageType.CHAT_CANCEL

    def test_from_dict_missing_fields(self):
        d = {"type": "chat.token"}
        msg = WSMessage.from_dict(d)
        assert msg.type == MessageType.CHAT_TOKEN
        assert msg.data == {}

    def test_default_data_is_empty_dict(self):
        msg = WSMessage(type=MessageType.PING)
        assert msg.data == {}


# ---------------------------------------------------------------------------
# Convenience constructors
# ---------------------------------------------------------------------------


class TestChatTokenConstructor:
    def test_creates_correct_type(self):
        msg = chat_token("hello")
        assert msg.type == MessageType.CHAT_TOKEN
        assert msg.data["content"] == "hello"

    def test_with_message_id(self):
        msg = chat_token("text", message_id="m1")
        assert msg.id == "m1"


class TestChatMessageStartEnd:
    def test_start(self):
        msg = chat_message_start("m1")
        assert msg.type == MessageType.CHAT_MESSAGE_START
        assert msg.data["message_id"] == "m1"
        assert msg.id == "m1"

    def test_end(self):
        msg = chat_message_end("m1", tokens=150)
        assert msg.type == MessageType.CHAT_MESSAGE_END
        assert msg.data["message_id"] == "m1"
        assert msg.data["tokens"] == 150

    def test_end_default_tokens(self):
        msg = chat_message_end("m1")
        assert msg.data["tokens"] == 0


class TestToolConstructors:
    def test_tool_start(self):
        msg = tool_start("t1", "shell", {"command": "ls"})
        assert msg.type == MessageType.TOOL_START
        assert msg.data["tool_id"] == "t1"
        assert msg.data["tool_name"] == "shell"
        assert msg.data["args"]["command"] == "ls"
        assert msg.id == "t1"

    def test_tool_success(self):
        msg = tool_success("t1", "output text")
        assert msg.type == MessageType.TOOL_SUCCESS
        assert msg.data["tool_id"] == "t1"
        assert msg.data["output"] == "output text"

    def test_tool_error(self):
        msg = tool_error("t1", "something went wrong")
        assert msg.type == MessageType.TOOL_ERROR
        assert msg.data["error"] == "something went wrong"

    def test_tool_rejected(self):
        msg = tool_rejected("t1")
        assert msg.type == MessageType.TOOL_REJECTED
        assert msg.data["tool_id"] == "t1"


class TestApprovalRequest:
    def test_creates_correct_format(self):
        actions = [{"tool": "shell", "args": {"command": "rm -rf build"}}]
        msg = approval_request("int-1", actions, assistant_id="agent-1")
        assert msg.type == MessageType.APPROVAL_REQUEST
        assert msg.data["interrupt_id"] == "int-1"
        assert msg.data["action_requests"] == actions
        assert msg.data["assistant_id"] == "agent-1"
        assert msg.id == "int-1"


class TestInterviewConstructors:
    def test_interview_questions(self):
        questions = [{"id": "q1", "text": "What is your name?"}]
        msg = interview_questions("iv-1", questions, skippable=True)
        assert msg.type == MessageType.INTERVIEW_QUESTIONS
        assert msg.data["interview_id"] == "iv-1"
        assert msg.data["questions"] == questions
        assert msg.data["skippable"] is True

    def test_interview_complete(self):
        msg = interview_complete("iv-1")
        assert msg.type == MessageType.INTERVIEW_COMPLETE
        assert msg.data["interview_id"] == "iv-1"


class TestFileChanged:
    def test_with_diff(self):
        msg = file_changed("/src/main.py", "modified", diff="--- a\n+++ b")
        assert msg.type == MessageType.FILE_CHANGED
        assert msg.data["path"] == "/src/main.py"
        assert msg.data["action"] == "modified"
        assert msg.data["diff"] == "--- a\n+++ b"

    def test_without_diff(self):
        msg = file_changed("/src/new.py", "created")
        assert msg.data["diff"] is None


class TestMiscConstructors:
    def test_system_message(self):
        msg = system_message("Server restarting")
        assert msg.type == MessageType.SYSTEM_MESSAGE
        assert msg.data["content"] == "Server restarting"

    def test_error_message(self):
        msg = error_message("Something failed", code="INTERNAL")
        assert msg.type == MessageType.ERROR
        assert msg.data["error"] == "Something failed"
        assert msg.data["message"] == "Something failed"
        assert msg.data["code"] == "INTERNAL"

    def test_session_policy_updated(self):
        msg = session_policy_updated("safe_write")
        assert msg.type == MessageType.SESSION_POLICY_UPDATED
        assert msg.data["approval_policy"] == "safe_write"

    def test_connected(self):
        msg = connected("sess-1")
        assert msg.type == MessageType.CONNECTED
        assert msg.data["session_id"] == "sess-1"


class TestStreamConstructors:
    def test_stream_state(self):
        msg = stream_state("agent", {"messages": []})
        assert msg.type == MessageType.STREAM_STATE
        assert msg.data["node"] == "agent"

    def test_stream_debug(self):
        msg = stream_debug("task_start", node="agent")
        assert msg.type == MessageType.STREAM_DEBUG
        assert msg.data["event_type"] == "task_start"

    def test_stream_node_start(self):
        msg = stream_node_start("agent", {"query": "hello"})
        assert msg.type == MessageType.STREAM_NODE_START
        assert msg.data["node"] == "agent"

    def test_stream_node_end(self):
        msg = stream_node_end("agent", {"result": "done"})
        assert msg.type == MessageType.STREAM_NODE_END
        assert msg.data["output"]["result"] == "done"


class TestThinkingConstructors:
    def test_thinking_start(self):
        msg = thinking_start("m1")
        assert msg.type == MessageType.THINKING_START

    def test_thinking_token(self):
        msg = thinking_token("reasoning step 1")
        assert msg.type == MessageType.THINKING_TOKEN
        assert msg.data["content"] == "reasoning step 1"

    def test_thinking_end(self):
        msg = thinking_end("m1")
        assert msg.type == MessageType.THINKING_END
