import pytest
import json
from unittest.mock import AsyncMock, MagicMock, patch
from yier_agents.src.llm import LLM
from yier_agents.src.message import Message, ToolCall, LLMResponse


@pytest.mark.asyncio
async def test_llm_complete_stop():
    """测试正常结束的调用"""
    mock_response = MagicMock()
    mock_response.json.return_value = {
        "choices": [{
            "message": {
                "role": "assistant",
                "content": "Hello, how can I help?"
            },
            "finish_reason": "stop"
        }],
        "usage": {
            "prompt_tokens": 10,
            "completion_tokens": 5,
            "total_tokens": 15
        }
    }

    with patch("httpx.AsyncClient") as MockClient:
        mock_client = AsyncMock()
        mock_client.post = AsyncMock(return_value=mock_response)
        MockClient.return_value.__aenter__.return_value = mock_client

        llm = LLM(
            base_url="https://api.example.com",
            api_key="test-key",
            model="gpt-4",
            timeout=30.0
        )

        messages = [Message(role="user", content="Hello")]
        response = await llm.complete(messages)

        assert isinstance(response, LLMResponse)
        assert response.message.content == "Hello, how can I help?"
        assert response.finish_reason == "stop"
        assert response.message.tool_calls is None
        assert response.usage == 15  # Total tokens


@pytest.mark.asyncio
async def test_llm_complete_tool_calls():
    """测试带工具调用的响应"""
    mock_response = MagicMock()
    mock_response.json.return_value = {
        "choices": [{
            "message": {
                "role": "assistant",
                "content": None,
                "tool_calls": [
                    {
                        "id": "call_123",
                        "type": "function",
                        "function": {
                            "name": "get_weather",
                            "arguments": '{"city": "Beijing"}'
                        }
                    }
                ]
            },
            "finish_reason": "tool_calls"
        }]
    }

    with patch("httpx.AsyncClient") as MockClient:
        mock_client = AsyncMock()
        mock_client.post = AsyncMock(return_value=mock_response)
        MockClient.return_value.__aenter__.return_value = mock_client

        llm = LLM(
            base_url="https://api.example.com",
            api_key="test-key",
            model="gpt-4",
            timeout=30.0
        )

        messages = [Message(role="user", content="What's the weather in Beijing?")]
        tools = [
            {
                "type": "function",
                "function": {
                    "name": "get_weather",
                    "description": "Get weather",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "city": {"type": "string"}
                        }
                    }
                }
            }
        ]

        response = await llm.complete(messages, tools=tools)

        assert isinstance(response, LLMResponse)
        assert response.message.content is None
        assert response.finish_reason == "tool_calls"
        assert len(response.message.tool_calls) == 1
        assert response.message.tool_calls[0].id == "call_123"
        assert response.message.tool_calls[0].name == "get_weather"
        assert response.message.tool_calls[0].arguments == {"city": "Beijing"}


def test_llm_message_to_dict_user():
    """测试用户消息转换"""
    llm = LLM(
        base_url="https://api.example.com",
        api_key="test-key",
        model="gpt-4",
        timeout=30.0
    )

    msg = Message(role="user", content="Hello")
    result = llm._message_to_dict(msg)

    assert result == {
        "role": "user",
        "content": "Hello"
    }


def test_llm_message_to_dict_assistant_with_tools():
    """测试助手消息带工具调用的转换"""
    llm = LLM(
        base_url="https://api.example.com",
        api_key="test-key",
        model="gpt-4",
        timeout=30.0
    )

    msg = Message(
        role="assistant",
        content=None,
        tool_calls=[
            ToolCall(
                id="call_123",
                name="get_weather",
                arguments={"city": "Beijing"}
            )
        ]
    )
    result = llm._message_to_dict(msg)

    assert result == {
        "role": "assistant",
        "content": None,
        "tool_calls": [
            {
                "id": "call_123",
                "type": "function",
                "function": {
                    "name": "get_weather",
                    "arguments": '{"city": "Beijing"}'
                }
            }
        ]
    }


def test_normalize_finish_reason():
    """测试 finish_reason 标准化"""
    llm = LLM(
        base_url="https://api.example.com",
        api_key="test-key",
        model="gpt-4",
        timeout=30.0
    )

    # 有 tool_calls 返回 "tool_calls"
    assert llm._normalize_finish_reason("stop", has_tool_calls=True) == "tool_calls"

    # raw_reason == "length" 返回 "length"
    assert llm._normalize_finish_reason("length", has_tool_calls=False) == "length"

    # 其他返回 "stop"
    assert llm._normalize_finish_reason("stop", has_tool_calls=False) == "stop"
    assert llm._normalize_finish_reason(None, has_tool_calls=False) == "stop"
    assert llm._normalize_finish_reason("end_turn", has_tool_calls=False) == "stop"
