from __future__ import annotations

import base64
import sys
from typing import Any
from types import SimpleNamespace

import pytest

from runtime.planner_config import PlannerConfig
from runtime.tool_calling_client import (
    AnthropicToolCallingClient,
    OpenAIToolCallingClient,
    ToolCallDecision,
    ToolCallUnavailable,
    build_client,
)
from runtime.tool_specs import FINISH_TASK_SPEC, TAP_SPEC
from tests.fakes import PNG_10X20


class FakeMessages:
    def __init__(
        self,
        *,
        response: object | None = None,
        responses: list[object] | None = None,
        error: Exception | None = None,
    ) -> None:
        self.response = response
        self._responses = list(responses) if responses is not None else None
        self.error = error
        self.calls: list[dict[str, Any]] = []

    def create(self, **kwargs: Any) -> object:
        self.calls.append(kwargs)
        if self.error:
            raise self.error
        if self._responses is not None:
            return self._responses.pop(0)
        return self.response


class FakeTransport:
    def __init__(self, messages: FakeMessages) -> None:
        self.messages = messages


class FakeCompletions:
    def __init__(
        self,
        *,
        response: object | None = None,
        responses: list[object] | None = None,
        error: Exception | None = None,
    ) -> None:
        self.response = response
        self._responses = list(responses) if responses is not None else None
        self.error = error
        self.calls: list[dict[str, Any]] = []

    def create(self, **kwargs: Any) -> object:
        self.calls.append(kwargs)
        if self.error:
            raise self.error
        if self._responses is not None:
            return self._responses.pop(0)
        return self.response


class FakeChat:
    def __init__(self, completions: FakeCompletions) -> None:
        self.completions = completions


class FakeOpenAITransport:
    def __init__(self, completions: FakeCompletions) -> None:
        self.chat = FakeChat(completions)


# --- Anthropic ---------------------------------------------------------


def test_anthropic_tool_calling_client_sends_auto_tool_choice_request() -> None:
    messages = FakeMessages(
        response={
            "content": [{"type": "tool_use", "name": "tap", "input": {"x": 1, "y": 2}}]
        }
    )
    client = AnthropicToolCallingClient(
        model="test-model", transport=FakeTransport(messages)
    )

    decision = client.decide(
        system_prompt="system",
        user_prompt="user",
        screenshot=None,
        tools=[TAP_SPEC, FINISH_TASK_SPEC],
        timeout=2.5,
    )

    assert decision == ToolCallDecision(
        tool_name="tap",
        arguments={"x": 1, "y": 2},
        system_prompt="system",
        user_prompt="user",
    )
    assert len(messages.calls) == 1
    call = messages.calls[0]
    assert call["model"] == "test-model"
    assert call["timeout"] == 2.5
    assert call["tool_choice"] == {"type": "auto", "disable_parallel_tool_use": True}
    assert call["tools"] == [
        {
            "name": "tap",
            "description": TAP_SPEC.description,
            "input_schema": TAP_SPEC.parameters,
        },
        {
            "name": "finish_task",
            "description": FINISH_TASK_SPEC.description,
            "input_schema": FINISH_TASK_SPEC.parameters,
        },
    ]
    assert call["system"][0]["text"] == "system"
    assert call["system"][0]["cache_control"] == {"type": "ephemeral"}
    assert call["messages"] == [
        {"role": "user", "content": [{"type": "text", "text": "user"}]}
    ]


def test_anthropic_tool_calling_client_includes_image_block_when_screenshot_present() -> (
    None
):
    messages = FakeMessages(
        response={
            "content": [
                {
                    "type": "tool_use",
                    "name": "finish_task",
                    "input": {"success": True, "reason": "done"},
                }
            ]
        }
    )
    client = AnthropicToolCallingClient(
        model="test-model", transport=FakeTransport(messages)
    )

    client.decide(
        system_prompt="system",
        user_prompt="user",
        screenshot=PNG_10X20,
        tools=[FINISH_TASK_SPEC],
        timeout=1,
    )

    content = messages.calls[0]["messages"][0]["content"]
    assert content[0] == {
        "type": "image",
        "source": {
            "type": "base64",
            "media_type": "image/png",
            "data": base64.b64encode(PNG_10X20).decode("ascii"),
        },
    }
    assert content[1] == {"type": "text", "text": "user"}


def test_anthropic_tool_calling_client_passes_custom_base_url_to_sdk(
    monkeypatch,
) -> None:
    constructed: list[dict[str, Any]] = []

    class RecordingAnthropic:
        def __init__(self, **kwargs: str) -> None:
            constructed.append(kwargs)

    monkeypatch.setitem(
        sys.modules,
        "anthropic",
        SimpleNamespace(Anthropic=RecordingAnthropic),
    )
    client = AnthropicToolCallingClient(
        model="test-model",
        api_key="managed-api-key",
        base_url="https://anthropic-proxy.example",
    )

    assert isinstance(client._client(), RecordingAnthropic)
    assert constructed == [
        {
            "api_key": "managed-api-key",
            "base_url": "https://anthropic-proxy.example",
        }
    ]


def test_anthropic_tool_calling_client_wraps_transport_errors() -> None:
    messages = FakeMessages(error=TimeoutError("timed out"))
    client = AnthropicToolCallingClient(
        model="test-model", transport=FakeTransport(messages)
    )

    with pytest.raises(ToolCallUnavailable):
        client.decide(
            system_prompt="s",
            user_prompt="u",
            screenshot=None,
            tools=[TAP_SPEC],
            timeout=1,
        )


@pytest.mark.parametrize(
    "response",
    [
        {"content": []},
        {"content": [{"type": "text", "text": "no tool call"}]},
        {"content": [{"type": "tool_use", "name": "tap", "input": "not-a-dict"}]},
    ],
)
def test_anthropic_tool_calling_client_wraps_malformed_responses(
    response: object,
) -> None:
    messages = FakeMessages(response=response)
    client = AnthropicToolCallingClient(
        model="test-model", transport=FakeTransport(messages)
    )

    with pytest.raises(ToolCallUnavailable):
        client.decide(
            system_prompt="s",
            user_prompt="u",
            screenshot=None,
            tools=[TAP_SPEC],
            timeout=1,
        )


def test_anthropic_client_retries_with_forced_tool_choice_when_model_omits_tool_call() -> (
    None
):
    messages = FakeMessages(
        responses=[
            {"content": [{"type": "text", "text": "just thinking out loud"}]},
            {
                "content": [
                    {"type": "tool_use", "name": "tap", "input": {"x": 1, "y": 2}}
                ]
            },
        ]
    )
    client = AnthropicToolCallingClient(
        model="test-model", transport=FakeTransport(messages)
    )

    decision = client.decide(
        system_prompt="s", user_prompt="u", screenshot=None, tools=[TAP_SPEC], timeout=1
    )

    assert decision.tool_name == "tap"
    assert len(messages.calls) == 2
    assert messages.calls[0]["tool_choice"] == {
        "type": "auto",
        "disable_parallel_tool_use": True,
    }
    assert messages.calls[1]["tool_choice"] == {
        "type": "any",
        "disable_parallel_tool_use": True,
    }


def test_anthropic_client_retry_drops_thinking_param_since_incompatible_with_forced_tool_choice() -> (
    None
):
    messages = FakeMessages(
        responses=[
            {"content": [{"type": "thinking", "thinking": "hmm, no tool yet"}]},
            {
                "content": [
                    {"type": "tool_use", "name": "tap", "input": {"x": 1, "y": 2}}
                ]
            },
        ]
    )
    client = AnthropicToolCallingClient(
        model="test-model",
        transport=FakeTransport(messages),
        thinking_budget_tokens=1024,
    )

    client.decide(
        system_prompt="s", user_prompt="u", screenshot=None, tools=[TAP_SPEC], timeout=1
    )

    assert "thinking" in messages.calls[0]
    assert "betas" in messages.calls[0]
    assert "thinking" not in messages.calls[1]
    assert "betas" not in messages.calls[1]


# --- OpenAI --------------------------------------------------------------


def test_openai_tool_calling_client_sends_auto_tool_choice_request() -> None:
    completions = FakeCompletions(
        response={
            "choices": [
                {
                    "message": {
                        "tool_calls": [
                            {
                                "function": {
                                    "name": "tap",
                                    "arguments": '{"x": 1, "y": 2}',
                                }
                            }
                        ]
                    }
                }
            ]
        }
    )
    client = OpenAIToolCallingClient(
        model="test-model", transport=FakeOpenAITransport(completions)
    )

    decision = client.decide(
        system_prompt="system",
        user_prompt="user",
        screenshot=None,
        tools=[TAP_SPEC, FINISH_TASK_SPEC],
        timeout=2.5,
    )

    assert decision == ToolCallDecision(
        tool_name="tap",
        arguments={"x": 1, "y": 2},
        system_prompt="system",
        user_prompt="user",
    )
    assert len(completions.calls) == 1
    call = completions.calls[0]
    assert call["model"] == "test-model"
    assert call["timeout"] == 2.5
    assert call["max_completion_tokens"] == 1024
    assert "max_tokens" not in call
    assert call["tool_choice"] == "auto"
    assert call["parallel_tool_calls"] is False
    assert call["tools"] == [
        {
            "type": "function",
            "function": {
                "name": "tap",
                "description": TAP_SPEC.description,
                "parameters": TAP_SPEC.parameters,
            },
        },
        {
            "type": "function",
            "function": {
                "name": "finish_task",
                "description": FINISH_TASK_SPEC.description,
                "parameters": FINISH_TASK_SPEC.parameters,
            },
        },
    ]
    assert call["messages"] == [
        {"role": "system", "content": "system"},
        {"role": "user", "content": "user"},
    ]


def test_openai_tool_calling_client_includes_image_block_when_screenshot_present() -> (
    None
):
    completions = FakeCompletions(
        response={
            "choices": [
                {
                    "message": {
                        "tool_calls": [
                            {
                                "function": {
                                    "name": "finish_task",
                                    "arguments": '{"success": true, "reason": "done"}',
                                }
                            }
                        ]
                    }
                }
            ]
        }
    )
    client = OpenAIToolCallingClient(
        model="test-model", transport=FakeOpenAITransport(completions)
    )

    client.decide(
        system_prompt="system",
        user_prompt="user",
        screenshot=PNG_10X20,
        tools=[FINISH_TASK_SPEC],
        timeout=1,
    )

    user_message = completions.calls[0]["messages"][1]
    assert user_message["role"] == "user"
    assert user_message["content"][0] == {"type": "text", "text": "user"}
    encoded = base64.b64encode(PNG_10X20).decode("ascii")
    assert user_message["content"][1] == {
        "type": "image_url",
        "image_url": {"url": f"data:image/png;base64,{encoded}"},
    }


def test_openai_tool_calling_client_accepts_arguments_already_as_dict() -> None:
    completions = FakeCompletions(
        response={
            "choices": [
                {
                    "message": {
                        "tool_calls": [
                            {"function": {"name": "tap", "arguments": {"x": 1, "y": 2}}}
                        ]
                    }
                }
            ]
        }
    )
    client = OpenAIToolCallingClient(
        model="test-model", transport=FakeOpenAITransport(completions)
    )

    decision = client.decide(
        system_prompt="s", user_prompt="u", screenshot=None, tools=[TAP_SPEC], timeout=1
    )

    assert decision == ToolCallDecision(
        tool_name="tap",
        arguments={"x": 1, "y": 2},
        system_prompt="s",
        user_prompt="u",
    )


def test_openai_tool_calling_client_wraps_transport_errors() -> None:
    completions = FakeCompletions(error=TimeoutError("timed out"))
    client = OpenAIToolCallingClient(
        model="test-model", transport=FakeOpenAITransport(completions)
    )

    with pytest.raises(ToolCallUnavailable):
        client.decide(
            system_prompt="s",
            user_prompt="u",
            screenshot=None,
            tools=[TAP_SPEC],
            timeout=1,
        )


@pytest.mark.parametrize(
    "response",
    [
        {"choices": []},
        {"choices": [{"message": {"tool_calls": []}}]},
        {
            "choices": [
                {
                    "message": {
                        "tool_calls": [
                            {"function": {"name": "tap", "arguments": "not-json"}}
                        ]
                    }
                }
            ]
        },
    ],
)
def test_openai_tool_calling_client_wraps_malformed_responses(response: object) -> None:
    completions = FakeCompletions(response=response)
    client = OpenAIToolCallingClient(
        model="test-model", transport=FakeOpenAITransport(completions)
    )

    with pytest.raises(ToolCallUnavailable):
        client.decide(
            system_prompt="s",
            user_prompt="u",
            screenshot=None,
            tools=[TAP_SPEC],
            timeout=1,
        )


def test_openai_client_retries_with_forced_tool_choice_when_model_omits_tool_call() -> (
    None
):
    completions = FakeCompletions(
        responses=[
            {"choices": [{"message": {"content": "just chatting, no tool"}}]},
            {
                "choices": [
                    {
                        "message": {
                            "tool_calls": [
                                {
                                    "function": {
                                        "name": "tap",
                                        "arguments": '{"x": 1, "y": 2}',
                                    }
                                }
                            ]
                        }
                    }
                ]
            },
        ]
    )
    client = OpenAIToolCallingClient(
        model="test-model", transport=FakeOpenAITransport(completions)
    )

    decision = client.decide(
        system_prompt="s", user_prompt="u", screenshot=None, tools=[TAP_SPEC], timeout=1
    )

    assert decision.tool_name == "tap"
    assert len(completions.calls) == 2
    assert completions.calls[0]["tool_choice"] == "auto"
    assert completions.calls[1]["tool_choice"] == "required"


# --- build_client ----------------------------------------------------------


def test_build_client_selects_provider_and_resolves_default_model() -> None:
    anthropic_client = build_client(PlannerConfig(provider="anthropic", model=""))
    assert isinstance(anthropic_client, AnthropicToolCallingClient)
    assert anthropic_client.model == "claude-sonnet-5"

    openai_client = build_client(PlannerConfig(provider="openai", model=""))
    assert isinstance(openai_client, OpenAIToolCallingClient)
    assert openai_client.model == "gpt-5.6"


def test_build_client_honors_explicit_model_override() -> None:
    client = build_client(PlannerConfig(provider="openai", model="gpt-5.6-custom"))
    assert client.model == "gpt-5.6-custom"


# --- thinking / text_output capture -----------------------------------------


def test_anthropic_client_captures_thinking_block() -> None:
    messages = FakeMessages(
        response={
            "content": [
                {"type": "thinking", "thinking": "I should tap the button."},
                {"type": "tool_use", "name": "tap", "input": {"x": 10, "y": 20}},
            ]
        }
    )
    client = AnthropicToolCallingClient(
        model="test-model", transport=FakeTransport(messages)
    )
    decision = client.decide(
        system_prompt="s", user_prompt="u", screenshot=None, tools=[TAP_SPEC], timeout=1
    )
    assert decision.thinking == "I should tap the button."
    assert decision.text_output is None


def test_anthropic_client_captures_text_block_as_text_output() -> None:
    messages = FakeMessages(
        response={
            "content": [
                {"type": "text", "text": "Previous step succeeded. Now tapping login."},
                {"type": "tool_use", "name": "tap", "input": {"x": 5, "y": 5}},
            ]
        }
    )
    client = AnthropicToolCallingClient(
        model="test-model", transport=FakeTransport(messages)
    )
    decision = client.decide(
        system_prompt="s", user_prompt="u", screenshot=None, tools=[TAP_SPEC], timeout=1
    )
    assert decision.text_output == "Previous step succeeded. Now tapping login."
    assert decision.thinking is None


def test_anthropic_client_captures_both_thinking_and_text_output() -> None:
    messages = FakeMessages(
        response={
            "content": [
                {"type": "thinking", "thinking": "Deep thought."},
                {"type": "text", "text": "Step succeeded. Tapping next."},
                {"type": "tool_use", "name": "tap", "input": {"x": 1, "y": 1}},
            ]
        }
    )
    client = AnthropicToolCallingClient(
        model="test-model", transport=FakeTransport(messages)
    )
    decision = client.decide(
        system_prompt="s", user_prompt="u", screenshot=None, tools=[TAP_SPEC], timeout=1
    )
    assert decision.thinking == "Deep thought."
    assert decision.text_output == "Step succeeded. Tapping next."


def test_anthropic_client_tool_only_response_has_none_thinking_and_text_output() -> (
    None
):
    messages = FakeMessages(
        response={
            "content": [
                {"type": "tool_use", "name": "tap", "input": {"x": 0, "y": 0}},
            ]
        }
    )
    client = AnthropicToolCallingClient(
        model="test-model", transport=FakeTransport(messages)
    )
    decision = client.decide(
        system_prompt="s", user_prompt="u", screenshot=None, tools=[TAP_SPEC], timeout=1
    )
    assert decision.thinking is None
    assert decision.text_output is None


def test_anthropic_client_separates_required_action_metadata_from_arguments() -> None:
    messages = FakeMessages(
        response={
            "content": [
                {
                    "type": "tool_use",
                    "name": "tap",
                    "input": {
                        "x": 10,
                        "y": 20,
                        "purpose": "Open the account screen.",
                        "expected_outcome": "The account screen is visible.",
                    },
                }
            ]
        }
    )
    client = AnthropicToolCallingClient(
        model="test-model", transport=FakeTransport(messages)
    )

    decision = client.decide(
        system_prompt="s", user_prompt="u", screenshot=None, tools=[TAP_SPEC], timeout=1
    )

    assert decision.arguments == {"x": 10, "y": 20}
    assert decision.purpose == "Open the account screen."
    assert decision.expected_outcome == "The account screen is visible."


def test_openai_client_captures_reasoning_content() -> None:
    completions = FakeCompletions(
        response={
            "choices": [
                {
                    "message": {
                        "reasoning_content": "I reasoned about this step.",
                        "tool_calls": [
                            {
                                "function": {
                                    "name": "tap",
                                    "arguments": '{"x": 1, "y": 2}',
                                }
                            }
                        ],
                    }
                }
            ]
        }
    )
    client = OpenAIToolCallingClient(
        model="test-model", transport=FakeOpenAITransport(completions)
    )
    decision = client.decide(
        system_prompt="s", user_prompt="u", screenshot=None, tools=[TAP_SPEC], timeout=1
    )
    assert decision.thinking == "I reasoned about this step."
    assert decision.text_output is None


def test_openai_client_captures_message_content_as_text_output() -> None:
    completions = FakeCompletions(
        response={
            "choices": [
                {
                    "message": {
                        "content": "Previous step succeeded. Now tapping login.",
                        "tool_calls": [
                            {
                                "function": {
                                    "name": "tap",
                                    "arguments": '{"x": 5, "y": 5}',
                                }
                            }
                        ],
                    }
                }
            ]
        }
    )
    client = OpenAIToolCallingClient(
        model="test-model", transport=FakeOpenAITransport(completions)
    )
    decision = client.decide(
        system_prompt="s", user_prompt="u", screenshot=None, tools=[TAP_SPEC], timeout=1
    )
    assert decision.text_output == "Previous step succeeded. Now tapping login."


def test_openai_client_separates_required_action_metadata_from_arguments() -> None:
    completions = FakeCompletions(
        response={
            "choices": [
                {
                    "message": {
                        "tool_calls": [
                            {
                                "function": {
                                    "name": "tap",
                                    "arguments": (
                                        '{"x": 5, "y": 6, '
                                        '"purpose": "Open settings.", '
                                        '"expected_outcome": "Settings is visible."}'
                                    ),
                                }
                            }
                        ]
                    }
                }
            ]
        }
    )
    client = OpenAIToolCallingClient(
        model="test-model", transport=FakeOpenAITransport(completions)
    )

    decision = client.decide(
        system_prompt="s", user_prompt="u", screenshot=None, tools=[TAP_SPEC], timeout=1
    )

    assert decision.arguments == {"x": 5, "y": 6}
    assert decision.purpose == "Open settings."
    assert decision.expected_outcome == "Settings is visible."


def test_openai_client_no_reasoning_content_gives_none_thinking() -> None:
    completions = FakeCompletions(
        response={
            "choices": [
                {
                    "message": {
                        "tool_calls": [
                            {
                                "function": {
                                    "name": "tap",
                                    "arguments": '{"x": 1, "y": 2}',
                                }
                            }
                        ]
                    }
                }
            ]
        }
    )
    client = OpenAIToolCallingClient(
        model="test-model", transport=FakeOpenAITransport(completions)
    )
    decision = client.decide(
        system_prompt="s", user_prompt="u", screenshot=None, tools=[TAP_SPEC], timeout=1
    )
    assert decision.thinking is None


def test_anthropic_client_sends_thinking_param_when_budget_set() -> None:
    messages = FakeMessages(
        response={
            "content": [{"type": "tool_use", "name": "tap", "input": {"x": 1, "y": 2}}]
        }
    )
    client = AnthropicToolCallingClient(
        model="test-model",
        transport=FakeTransport(messages),
        thinking_budget_tokens=2048,
    )
    client.decide(
        system_prompt="s", user_prompt="u", screenshot=None, tools=[TAP_SPEC], timeout=1
    )
    call = messages.calls[0]
    assert call["thinking"] == {"type": "enabled", "budget_tokens": 2048}
    assert "interleaved-thinking-2025-05-14" in call["betas"]
    # max_tokens must be >= budget + 1
    assert call["max_tokens"] >= 2049


def test_anthropic_client_enforces_max_tokens_floor_for_thinking() -> None:
    messages = FakeMessages(
        response={
            "content": [{"type": "tool_use", "name": "tap", "input": {"x": 1, "y": 2}}]
        }
    )
    # max_tokens=1024, budget=4096 → max_tokens should be raised to 4097
    client = AnthropicToolCallingClient(
        model="test-model",
        transport=FakeTransport(messages),
        max_tokens=1024,
        thinking_budget_tokens=4096,
    )
    client.decide(
        system_prompt="s", user_prompt="u", screenshot=None, tools=[TAP_SPEC], timeout=1
    )
    assert messages.calls[0]["max_tokens"] == 4097


def test_anthropic_client_no_thinking_param_when_budget_not_set() -> None:
    messages = FakeMessages(
        response={
            "content": [{"type": "tool_use", "name": "tap", "input": {"x": 1, "y": 2}}]
        }
    )
    client = AnthropicToolCallingClient(
        model="test-model", transport=FakeTransport(messages)
    )
    client.decide(
        system_prompt="s", user_prompt="u", screenshot=None, tools=[TAP_SPEC], timeout=1
    )
    call = messages.calls[0]
    assert "thinking" not in call
    assert "betas" not in call