"""Integration tests: LangChain RotatorChatModel."""

from __future__ import annotations

import pytest
from langchain_core.messages import AIMessage, AIMessageChunk, HumanMessage, SystemMessage

from llm_rotator import LLMRotator, RotatorConfig, RoutingContext
from llm_rotator._types import LLMResponse, ToolCall, Usage
from llm_rotator.exceptions import ServerError
from llm_rotator.integrations.langchain import RotatorChatModel
from tests.conftest import FakeLLMClient


def _make_rotator(client: FakeLLMClient) -> LLMRotator:
    config = RotatorConfig(
        providers=[
            {
                "name": "test_provider",
                "client_type": "openai",
                "priority": 1,
                "models": ["gpt-4o"],
                "keys": [{"token": "sk-test", "alias": "key1"}],
            }
        ]
    )
    return LLMRotator(config=config, clients={"openai": client})


def _make_rotator_two_providers(
    client: FakeLLMClient,
) -> LLMRotator:
    config = RotatorConfig(
        providers=[
            {
                "name": "primary",
                "client_type": "openai",
                "priority": 1,
                "models": ["gpt-4o"],
                "keys": [{"token": "sk-1", "alias": "key1"}],
            },
            {
                "name": "fallback",
                "client_type": "openai",
                "priority": 2,
                "models": ["gpt-4o"],
                "keys": [{"token": "sk-2", "alias": "key2"}],
            },
        ]
    )
    return LLMRotator(config=config, clients={"openai": client})


def _response(
    content: str = "Hello!",
    model: str = "gpt-4o",
    tool_calls: list[ToolCall] | None = None,
) -> LLMResponse:
    return LLMResponse(
        content=content,
        usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15),
        model=model,
        provider="test_provider",
        key_alias="key1",
        tool_calls=tool_calls,
    )


# --- Tests ---


@pytest.mark.asyncio
async def test_invoke_returns_ai_message():
    client = FakeLLMClient()
    client.enqueue(_response("Hi there!"))
    rotator = _make_rotator(client)
    model = RotatorChatModel(rotator=rotator)

    result = await model.ainvoke([HumanMessage(content="Hello")])

    assert isinstance(result, AIMessage)
    assert result.content == "Hi there!"


@pytest.mark.asyncio
async def test_invoke_with_system():
    client = FakeLLMClient()
    client.enqueue(_response("I'm an assistant"))
    rotator = _make_rotator(client)
    model = RotatorChatModel(rotator=rotator)

    result = await model.ainvoke(
        [
            SystemMessage(content="You are helpful."),
            HumanMessage(content="Who are you?"),
        ]
    )

    assert isinstance(result, AIMessage)
    assert result.content == "I'm an assistant"
    # Verify messages were converted correctly
    call = client.calls[0]
    assert call["messages"][0] == {"role": "system", "content": "You are helpful."}
    assert call["messages"][1] == {"role": "user", "content": "Who are you?"}


@pytest.mark.asyncio
async def test_stream_returns_chunks():
    client = FakeLLMClient()
    client.enqueue(_response("streamed content"))
    rotator = _make_rotator(client)
    model = RotatorChatModel(rotator=rotator)

    chunks = []
    async for chunk in model.astream([HumanMessage(content="Hi")]):
        chunks.append(chunk)

    assert len(chunks) >= 1
    assert isinstance(chunks[0], AIMessageChunk)
    # The FakeLLMClient yields the full content as a single chunk
    assert chunks[0].content == "streamed content"


@pytest.mark.asyncio
async def test_bind_tools():
    client = FakeLLMClient()
    client.enqueue(_response("I'll search for that"))
    rotator = _make_rotator(client)
    model = RotatorChatModel(rotator=rotator)

    tools = [
        {
            "type": "function",
            "function": {
                "name": "search",
                "description": "Search the web",
                "parameters": {
                    "type": "object",
                    "properties": {"query": {"type": "string"}},
                    "required": ["query"],
                },
            },
        }
    ]
    bound = model.bind_tools(tools)
    result = await bound.ainvoke([HumanMessage(content="Search for Python")])

    assert isinstance(result, AIMessage)
    # Verify tools were forwarded to the rotator
    call = client.calls[0]
    assert "tools" in call


@pytest.mark.asyncio
async def test_routing_context_via_kwargs():
    client = FakeLLMClient()
    client.enqueue(_response())
    config = RotatorConfig(
        providers=[
            {
                "name": "test_provider",
                "client_type": "openai",
                "priority": 1,
                "model_groups": [
                    {"name": "premium", "tier": 1, "models": ["gpt-4o"]},
                    {"name": "cheap", "tier": 3, "models": ["gpt-4o-mini"]},
                ],
                "keys": [{"token": "sk-test", "alias": "key1"}],
            }
        ]
    )
    rotator = LLMRotator(config=config, clients={"openai": client})
    model = RotatorChatModel(rotator=rotator)

    result = await model.ainvoke(
        [HumanMessage(content="Hi")],
        routing=RoutingContext(tier=3),
    )

    assert isinstance(result, AIMessage)
    # With tier=3, should use the cheap model
    call = client.calls[0]
    assert call["model"] == "gpt-4o-mini"


@pytest.mark.asyncio
async def test_fallback_transparent():
    client = FakeLLMClient()
    client.enqueue(ServerError("Provider down", status_code=500))
    client.enqueue(_response("Fallback succeeded"))
    rotator = _make_rotator_two_providers(client)
    model = RotatorChatModel(rotator=rotator)

    result = await model.ainvoke([HumanMessage(content="Hi")])

    assert isinstance(result, AIMessage)
    assert result.content == "Fallback succeeded"


@pytest.mark.asyncio
async def test_usage_in_response_metadata():
    client = FakeLLMClient()
    client.enqueue(_response("Hi"))
    rotator = _make_rotator(client)
    model = RotatorChatModel(rotator=rotator)

    result = await model.ainvoke([HumanMessage(content="Hello")])

    assert result.usage_metadata is not None
    assert result.usage_metadata["input_tokens"] == 10
    assert result.usage_metadata["output_tokens"] == 5
    assert result.usage_metadata["total_tokens"] == 15


@pytest.mark.asyncio
async def test_tool_call_response():
    client = FakeLLMClient()
    client.enqueue(
        _response(
            content="",
            tool_calls=[
                ToolCall(
                    id="call_123",
                    name="search",
                    arguments='{"query": "Python"}',
                )
            ],
        )
    )
    rotator = _make_rotator(client)
    model = RotatorChatModel(rotator=rotator)

    result = await model.ainvoke([HumanMessage(content="Search for Python")])

    assert isinstance(result, AIMessage)
    assert len(result.tool_calls) == 1
    assert result.tool_calls[0]["name"] == "search"
    assert result.tool_calls[0]["args"] == {"query": "Python"}
    assert result.tool_calls[0]["id"] == "call_123"


@pytest.mark.asyncio
async def test_in_chain():
    """Test that RotatorChatModel works in a LangChain pipe."""
    from langchain_core.output_parsers import StrOutputParser
    from langchain_core.prompts import ChatPromptTemplate

    client = FakeLLMClient()
    client.enqueue(_response("Paris"))
    rotator = _make_rotator(client)
    model = RotatorChatModel(rotator=rotator)

    prompt = ChatPromptTemplate.from_messages(
        [
            ("system", "Answer concisely."),
            ("user", "{question}"),
        ]
    )
    chain = prompt | model | StrOutputParser()
    result = await chain.ainvoke({"question": "Capital of France?"})

    assert result == "Paris"


@pytest.mark.asyncio
async def test_default_routing():
    client = FakeLLMClient()
    client.enqueue(_response())
    config = RotatorConfig(
        providers=[
            {
                "name": "test_provider",
                "client_type": "openai",
                "priority": 1,
                "model_groups": [
                    {"name": "premium", "tier": 1, "models": ["gpt-4o"]},
                    {"name": "cheap", "tier": 3, "models": ["gpt-4o-mini"]},
                ],
                "keys": [{"token": "sk-test", "alias": "key1"}],
            }
        ]
    )
    rotator = LLMRotator(config=config, clients={"openai": client})
    model = RotatorChatModel(
        rotator=rotator,
        default_routing=RoutingContext(tier=3),
    )

    await model.ainvoke([HumanMessage(content="Hi")])

    call = client.calls[0]
    assert call["model"] == "gpt-4o-mini"
