from __future__ import annotations

import asyncio
import json
from typing import Any

import httpx

from opensquilla.provider.ollama import _OLLAMA_DEFAULT_NUM_CTX, OllamaProvider
from opensquilla.provider.selector import ProviderConfig, _build_provider
from opensquilla.provider.types import (
    ChatConfig,
    ContentBlockImage,
    ContentBlockText,
    ContentBlockToolResult,
    ContentBlockToolUse,
    DoneEvent,
    Message,
    TextDeltaEvent,
    ToolDefinition,
    ToolInputSchema,
    ToolUseEndEvent,
)


def _ndjson(*chunks: dict[str, Any]) -> bytes:
    return b"".join((json.dumps(chunk) + "\n").encode() for chunk in chunks)


_DEFAULT_BODY = _ndjson(
    {"model": "llama3", "message": {"role": "assistant", "content": "ok"}},
    {
        "model": "llama3",
        "message": {"role": "assistant", "content": ""},
        "done": True,
        "prompt_eval_count": 3,
        "eval_count": 2,
    },
)


def _patch_stream(monkeypatch: Any, captured: dict[str, Any], body: bytes = _DEFAULT_BODY) -> None:
    def handler(request: httpx.Request) -> httpx.Response:
        captured["url"] = str(request.url)
        captured["headers"] = request.headers
        captured["payload"] = json.loads(request.content.decode("utf-8"))
        return httpx.Response(200, content=body)

    transport = httpx.MockTransport(handler)
    real_async_client = httpx.AsyncClient

    def patched_async_client(*args: Any, **kwargs: Any) -> httpx.AsyncClient:
        kwargs["transport"] = transport
        return real_async_client(*args, **kwargs)

    monkeypatch.setattr("opensquilla.provider.ollama.httpx.AsyncClient", patched_async_client)


def _collect(
    provider: OllamaProvider,
    messages: list[Message],
    cfg: ChatConfig | None = None,
    tools: list[ToolDefinition] | None = None,
) -> list[Any]:
    async def _run() -> list[Any]:
        return [
            event
            async for event in provider.chat(messages, tools=tools, config=cfg or ChatConfig())
        ]

    return asyncio.run(_run())


def test_tool_use_replayed_as_tool_calls_and_correlated_tool_message(monkeypatch: Any) -> None:
    captured: dict[str, Any] = {}
    _patch_stream(monkeypatch, captured)
    provider = OllamaProvider(model="llama3")
    messages = [
        Message(role="user", content="look it up"),
        Message(
            role="assistant",
            content=[ContentBlockToolUse(id="call_1", name="lookup", input={"q": "cache"})],
        ),
        Message(
            role="user",
            content=[ContentBlockToolResult(tool_use_id="call_1", content="cache is warm")],
        ),
        Message(role="user", content="continue"),
    ]

    _collect(provider, messages)

    msgs = captured["payload"]["messages"]
    # Assistant turn keeps its tool_calls instead of being dropped.
    assert msgs[1]["role"] == "assistant"
    assert msgs[1]["tool_calls"] == [{"function": {"name": "lookup", "arguments": {"q": "cache"}}}]
    # Tool result becomes its own tool-role message, correlated by tool_name.
    assert msgs[2] == {"role": "tool", "content": "cache is warm", "tool_name": "lookup"}
    assert msgs[3] == {"role": "user", "content": "continue"}


def test_parallel_tool_results_expand_to_separate_tool_messages(monkeypatch: Any) -> None:
    captured: dict[str, Any] = {}
    _patch_stream(monkeypatch, captured)
    provider = OllamaProvider(model="llama3")
    messages = [
        Message(
            role="assistant",
            content=[
                ContentBlockToolUse(id="call_a", name="foo", input={"x": 1}),
                ContentBlockToolUse(id="call_b", name="bar", input={"y": 2}),
            ],
        ),
        Message(
            role="user",
            content=[
                ContentBlockToolResult(tool_use_id="call_a", content="ra"),
                ContentBlockToolResult(tool_use_id="call_b", content="rb"),
            ],
        ),
    ]

    _collect(provider, messages)

    msgs = captured["payload"]["messages"]
    assert msgs[0]["tool_calls"] == [
        {"function": {"name": "foo", "arguments": {"x": 1}}},
        {"function": {"name": "bar", "arguments": {"y": 2}}},
    ]
    # Both results survive as distinct tool messages (the old code dropped all but one).
    assert msgs[1] == {"role": "tool", "content": "ra", "tool_name": "foo"}
    assert msgs[2] == {"role": "tool", "content": "rb", "tool_name": "bar"}


def test_tool_result_list_content_is_json_encoded(monkeypatch: Any) -> None:
    captured: dict[str, Any] = {}
    _patch_stream(monkeypatch, captured)
    provider = OllamaProvider(model="llama3")
    messages = [
        Message(
            role="user",
            content=[ContentBlockToolResult(tool_use_id="call_x", content=[{"k": "v"}])],
        ),
    ]

    _collect(provider, messages)

    assert captured["payload"]["messages"][0]["content"] == json.dumps([{"k": "v"}])


def test_num_ctx_defaults_above_ollama_2048(monkeypatch: Any) -> None:
    captured: dict[str, Any] = {}
    _patch_stream(monkeypatch, captured)
    provider = OllamaProvider(model="llama3")

    _collect(provider, [Message(role="user", content="hi")], ChatConfig(max_tokens=512))

    options = captured["payload"]["options"]
    assert options["num_ctx"] == _OLLAMA_DEFAULT_NUM_CTX
    assert _OLLAMA_DEFAULT_NUM_CTX > 2048
    assert options["num_predict"] == 512


def test_num_ctx_is_configurable(monkeypatch: Any) -> None:
    captured: dict[str, Any] = {}
    _patch_stream(monkeypatch, captured)
    provider = OllamaProvider(model="llama3", num_ctx=32768)

    _collect(provider, [Message(role="user", content="hi")])

    assert captured["payload"]["options"]["num_ctx"] == 32768


def test_bearer_header_sent_when_api_key_set(monkeypatch: Any) -> None:
    captured: dict[str, Any] = {}
    _patch_stream(monkeypatch, captured)
    provider = OllamaProvider(
        model="gpt-oss:120b-cloud", base_url="https://ollama.com", api_key="secret"
    )

    _collect(provider, [Message(role="user", content="hi")])

    assert captured["url"] == "https://ollama.com/api/chat"
    assert captured["headers"]["authorization"] == "Bearer secret"


def test_no_authorization_header_without_api_key(monkeypatch: Any) -> None:
    captured: dict[str, Any] = {}
    _patch_stream(monkeypatch, captured)
    provider = OllamaProvider(model="llama3")

    _collect(provider, [Message(role="user", content="hi")])

    assert "authorization" not in captured["headers"]


def test_image_block_attached_to_message_images(monkeypatch: Any) -> None:
    captured: dict[str, Any] = {}
    _patch_stream(monkeypatch, captured)
    provider = OllamaProvider(model="llava")
    messages = [
        Message(
            role="user",
            content=[
                ContentBlockText(text="what is this"),
                ContentBlockImage(source_type="base64", media_type="image/png", data="QkFTRTY0"),
            ],
        ),
    ]

    _collect(provider, messages)

    msg = captured["payload"]["messages"][0]
    assert msg["content"] == "what is this"
    assert msg["images"] == ["QkFTRTY0"]


def test_system_prompt_is_first_message(monkeypatch: Any) -> None:
    captured: dict[str, Any] = {}
    _patch_stream(monkeypatch, captured)
    provider = OllamaProvider(model="llama3")

    _collect(provider, [Message(role="user", content="hi")], ChatConfig(system="be brief"))

    assert captured["payload"]["messages"][0] == {"role": "system", "content": "be brief"}


def test_stream_emits_text_and_done_with_usage(monkeypatch: Any) -> None:
    captured: dict[str, Any] = {}
    _patch_stream(monkeypatch, captured)
    provider = OllamaProvider(model="llama3")

    events = _collect(provider, [Message(role="user", content="hi")])

    text = next(e for e in events if isinstance(e, TextDeltaEvent))
    done = next(e for e in events if isinstance(e, DoneEvent))
    assert text.text == "ok"
    assert done.input_tokens == 3
    assert done.output_tokens == 2


def test_stream_tool_call_emits_tool_events(monkeypatch: Any) -> None:
    captured: dict[str, Any] = {}
    body = _ndjson(
        {
            "model": "llama3",
            "message": {
                "role": "assistant",
                "content": "",
                "tool_calls": [{"function": {"name": "lookup", "arguments": {"q": "hi"}}}],
            },
        },
        {
            "model": "llama3",
            "message": {"role": "assistant", "content": ""},
            "done": True,
            "prompt_eval_count": 4,
            "eval_count": 1,
        },
    )
    _patch_stream(monkeypatch, captured, body)
    provider = OllamaProvider(model="llama3")
    tool = ToolDefinition(
        name="lookup",
        description="Lookup a value.",
        input_schema=ToolInputSchema(properties={"q": {"type": "string"}}, required=["q"]),
    )

    events = _collect(provider, [Message(role="user", content="hi")], tools=[tool])

    tool_end = next(e for e in events if isinstance(e, ToolUseEndEvent))
    assert tool_end.tool_name == "lookup"
    assert tool_end.arguments == {"q": "hi"}
    assert captured["payload"]["tools"][0]["function"]["name"] == "lookup"


def test_selector_passes_api_key_to_ollama_provider() -> None:
    cfg = ProviderConfig(
        provider="ollama",
        model="gpt-oss:120b-cloud",
        api_key="sk-test",
        base_url="https://ollama.com",
    )

    provider = _build_provider(cfg)

    assert isinstance(provider, OllamaProvider)
    assert provider._headers() == {"Authorization": "Bearer sk-test"}


def test_selector_omits_api_key_for_local_ollama() -> None:
    cfg = ProviderConfig(provider="ollama", model="llama3")

    provider = _build_provider(cfg)

    assert isinstance(provider, OllamaProvider)
    assert provider._headers() == {}