"""Unit tests for model adapters and dispatcher."""

import json
import pytest
import httpx
from backend.model import FakeModelAdapter, OpenAIModelAdapter, ModelDispatcher
from backend.errors import InvalidInputError, ModelContextExceededError


@pytest.mark.asyncio
async def test_fake_model_adapter():
    adapter = FakeModelAdapter()
    adapter.queue_response({
        "content": [{"type": "text", "text": "Hello world"}],
        "stop_reason": "stop",
        "usage": {"input_tokens": 5, "output_tokens": 5},
    })

    dispatcher = ModelDispatcher(adapter, max_calls=2)
    resp, err = await dispatcher.dispatch(
        {"messages": [{"role": "user", "content": [{"type": "text", "text": "Hi"}]}], "tools": []},
        system_instruction="System prompt",
    )
    assert err is None
    assert resp["stop_reason"] == "stop"
    assert resp["content"][0]["text"] == "Hello world"
    assert adapter.call_count == 1


@pytest.mark.asyncio
async def test_model_dispatcher_budget():
    adapter = FakeModelAdapter()
    dispatcher = ModelDispatcher(adapter, max_calls=1)

    # 1st call succeeds
    resp1, err1 = await dispatcher.dispatch(
        {"messages": [], "tools": []},
        system_instruction="",
    )
    assert err1 is None

    # 2nd call fails budget
    resp2, err2 = await dispatcher.dispatch(
        {"messages": [], "tools": []},
        system_instruction="",
    )
    assert resp2 is None
    assert err2 is not None
    assert "limit" in err2["message"]


@pytest.mark.asyncio
async def test_model_dispatcher_forged_handle():
    adapter = FakeModelAdapter()
    dispatcher = ModelDispatcher(adapter)

    # Supply an assistant message with forged provider_state
    resp, err = await dispatcher.dispatch(
        {
            "messages": [
                {"role": "assistant", "content": [], "provider_state": "forged-uuid-1234"}
            ],
            "tools": [],
        },
        system_instruction="",
    )
    assert resp is None
    assert err is not None
    assert err["code"] == "invalid_input"
    assert "provider_state" in err["message"]


@pytest.mark.asyncio
async def test_openai_adapter_translation():
    async def handler(request: httpx.Request) -> httpx.Response:
        assert request.headers.get("Authorization") == "Bearer test-key"
        body = json.loads(request.read().decode("utf-8"))
        assert body["model"] == "gpt-4o"
        assert body["messages"][0]["role"] == "system"
        assert body["messages"][0]["content"] == "Trusted system instruction"
        assert body["messages"][1]["role"] == "user"
        assert body["messages"][1]["content"] == "Hello"

        # Return tool call response
        return httpx.Response(
            200,
            json={
                "choices": [
                    {
                        "message": {
                            "role": "assistant",
                            "content": None,
                            "tool_calls": [
                                {
                                    "id": "tc_1",
                                    "type": "function",
                                    "function": {
                                        "name": "confluence_search",
                                        "arguments": "{\"query\": \"deploy\"}",
                                    },
                                }
                            ],
                        },
                        "finish_reason": "tool_calls",
                    }
                ],
                "usage": {"prompt_tokens": 15, "completion_tokens": 20},
            },
        )

    transport = httpx.MockTransport(handler)
    adapter = OpenAIModelAdapter(
        api_key="test-key",
        model_name="gpt-4o",
        transport=transport,
    )

    resp = await adapter.complete(
        messages=[{"role": "user", "content": [{"type": "text", "text": "Hello"}]}],
        tools=[{"name": "confluence_search", "description": "search", "input_schema": {}}],
        system_instruction="Trusted system instruction",
    )

    assert resp["stop_reason"] == "tool_calls"
    assert len(resp["content"]) == 1
    assert resp["content"][0]["type"] == "tool_call"
    assert resp["content"][0]["id"] == "tc_1"
    assert resp["content"][0]["name"] == "confluence_search"
    assert resp["content"][0]["arguments"] == {"query": "deploy"}
    assert resp["usage"] == {"input_tokens": 15, "output_tokens": 20}

    await adapter.close()