import pytest
from pydantic import ValidationError

from lambda_agent_api.server import *
from lambda_agent_api.client import *


@pytest.mark.parametrize(
    "data, expected_type",
    [
        ({"type": "USER_MESSAGE", "text": "hello"}, MsgUserMessage),
    ],
)
def test_client_message_valid(data, expected_type):
    msg = ClientMessage.validate_python(data)
    assert isinstance(msg, expected_type)


@pytest.mark.parametrize(
    "data",
    [
        {"type": "UNKNOWN", "text": "hello"},
        {"type": "USER_MESSAGE"},  # нет text
    ],
)
def test_client_message_invalid(data):
    with pytest.raises(ValidationError):
        ClientMessage.validate_python(data)


@pytest.mark.parametrize(
    "data, expected_type",
    [
        ({"type": "STATUS"}, MsgStatus),
        ({"type": "AGENT_EVENT_TEXT_CHUNK", "text": "hi"}, MsgEventTextChunk),
        ({"type": "AGENT_EVENT_TOOL_CALL_CHUNK", "tool_name": "search", "args_chunk": "{\"q\": \"hello\"}"}, MsgEventToolCallChunk),
        ({"type": "AGENT_EVENT_TOOL_RESULT", "tool_name": "search", "result": {"items": [1, 2, 3]}}, MsgEventToolResult),
        ({"type": "AGENT_EVENT_CUSTOM_UPDATE", "payload": {"status": "in_progress", "progress": 50}}, MsgEventCustomUpdate),
        ({"type": "AGENT_EVENT_END", "tokens_used": 10}, MsgEventEnd),
        ({"type": "ERROR", "code": "E1", "details": "fail"}, MsgError),
        ({"type": "GRACEFUL_DISCONNECT"}, MsgGracefulDisconnect),
    ],
)
def test_server_message_valid(data, expected_type):
    msg = ServerMessage.validate_python(data)
    assert isinstance(msg, expected_type)


@pytest.mark.parametrize(
    "data",
    [
        {"type": "AGENT_EVENT_TEXT_CHUNK"},  # нет text
        {"type": "AGENT_EVENT_TOOL_RESULT", "tool_name": "search"},  # нет result
        {"type": "AGENT_EVENT_CUSTOM_UPDATE"},  # нет payload
        {"type": "AGENT_EVENT_END"},  # нет tokens_used
        {"type": "ERROR", "code": "E1"},  # нет details
        {"type": "UNKNOWN"},
    ],
)
def test_server_message_invalid(data):
    with pytest.raises(ValidationError):
        ServerMessage.validate_python(data)


def test_validate_json():
    json_data = '{"type": "AGENT_EVENT_TEXT_CHUNK", "text": "hello"}'
    msg = ServerMessage.validate_json(json_data)

    assert isinstance(msg, MsgEventTextChunk)
    assert msg.text == "hello"


def test_model_dump_roundtrip():
    data = {"type": "ERROR", "code": "E1", "details": "fail"}

    msg = ServerMessage.validate_python(data)
    dumped = msg.model_dump()

    assert dumped == data