"""Tests for the data models."""

import pytest
from pydantic import ValidationError

from lumo_sdk import Message, MessageRole, RunTaskRequest, RunTaskResponse, Step, ToolCall


class TestMessage:
    """Test cases for Message model."""

    def test_message_creation(self):
        """Test creating a message with required fields."""
        message = Message(role=MessageRole.USER, content="Hello")
        assert message.role == MessageRole.USER
        assert message.content == "Hello"
        assert message.tool_call_id is None
        assert message.tool_calls is None

    def test_message_with_optional_fields(self):
        """Test creating a message with optional fields."""
        tool_call = ToolCall(id="123", type="function", function={"name": "test", "arguments": {}})
        message = Message(
            role=MessageRole.ASSISTANT,
            content="Response",
            tool_call_id="123",
            tool_calls=[tool_call],
        )
        assert message.tool_call_id == "123"
        assert len(message.tool_calls) == 1

    def test_message_missing_required_fields(self):
        """Test that missing required fields raise validation error."""
        with pytest.raises(ValidationError):
            Message(role=MessageRole.USER)  # Missing content


class TestRunTaskRequest:
    """Test cases for RunTaskRequest model."""

    def test_request_creation(self):
        """Test creating a request with required fields."""
        request = RunTaskRequest(
            task="Test task",
            model="gpt-4.1-mini",
            base_url="https://api.openai.com/v1/chat/completions",
        )
        assert request.task == "Test task"
        assert request.model == "gpt-4.1-mini"
        assert request.tools is None

    def test_request_with_optional_fields(self):
        """Test creating a request with optional fields."""
        request = RunTaskRequest(
            task="Test task",
            model="gpt-4.1-mini",
            base_url="https://api.openai.com/v1/chat/completions",
            tools=["Tool1", "Tool2"],
            max_steps=5,
            agent_type="function-calling",
        )
        assert request.tools == ["Tool1", "Tool2"]
        assert request.max_steps == 5
        assert request.agent_type == "function-calling"

    def test_request_missing_required_fields(self):
        """Test that missing required fields raise validation error."""
        with pytest.raises(ValidationError):
            RunTaskRequest(task="Test")  # Missing model and base_url


class TestRunTaskResponse:
    """Test cases for RunTaskResponse model."""

    def test_response_creation(self):
        """Test creating a response."""
        step = Step(tool_calls=[], llm_output="Output")
        response = RunTaskResponse(final_answer="Answer", steps=[step])
        assert response.final_answer == "Answer"
        assert len(response.steps) == 1

    def test_response_missing_required_fields(self):
        """Test that missing required fields raise validation error."""
        with pytest.raises(ValidationError):
            RunTaskResponse(final_answer="Answer")  # Missing steps


class TestStep:
    """Test cases for Step model."""

    def test_step_creation(self):
        """Test creating a step."""
        step = Step(tool_calls=[], llm_output="Output")
        assert step.llm_output == "Output"
        assert step.tool_calls == []

    def test_step_with_tool_calls(self):
        """Test creating a step with tool calls."""
        tool_call = ToolCall(id="123", type="function", function={"name": "test", "arguments": {}})
        step = Step(tool_calls=[tool_call], llm_output="Output")
        assert len(step.tool_calls) == 1
        assert step.tool_calls[0].id == "123"

