from __future__ import annotations

import base64
import json
from dataclasses import dataclass
from typing import Any, Protocol

from runtime.planner_config import PlannerConfig
from runtime.tool_specs import ACTION_TOOL_NAMES, ToolSpec


class ToolCallUnavailable(Exception):
    """Internal signal for expected tool-calling transport/response failures."""


@dataclass(frozen=True)
class ToolCallUsage:
    input_tokens: int | None = None
    output_tokens: int | None = None
    total_tokens: int | None = None


@dataclass(frozen=True)
class ToolCallDecision:
    tool_name: str
    arguments: dict[str, Any]
    usage: ToolCallUsage | None = None
    # The actual prompts sent to the LLM for this call. Populated by the
    # built-in Anthropic/OpenAI clients; empty for clients (e.g.
    # CloudProxyToolCallingClient) that don't surface them.
    system_prompt: str = ""
    user_prompt: str = ""
    # Pre-tool text block emitted by the model (rationale / reflection).
    # None when the model omits a text block before the tool call.
    text_output: str | None = None
    # Extended thinking block (Anthropic) or reasoning_content (OpenAI o-series).
    # None when not enabled or not present in the response.
    thinking: str | None = None
    # Required structured metadata for device actions. These fields are removed
    # from ``arguments`` before the Runtime invokes the physical device tool.
    purpose: str | None = None
    expected_outcome: str | None = None


class ToolCallingClient(Protocol):
    def decide(
        self,
        *,
        system_prompt: str,
        user_prompt: str,
        screenshot: bytes | None,
        tools: list[ToolSpec],
        timeout: float,
    ) -> ToolCallDecision: ...


class AnthropicToolCallingClient:
    def __init__(
        self,
        *,
        model: str,
        transport: Any | None = None,
        max_tokens: int = 1024,
        api_key: str | None = None,
        base_url: str | None = None,
        thinking_budget_tokens: int | None = None,
    ) -> None:
        self.model = model
        self._transport = transport
        self.max_tokens = max_tokens
        self._api_key = api_key
        self._base_url = base_url
        self._thinking_budget_tokens = thinking_budget_tokens

    def decide(
        self,
        *,
        system_prompt: str,
        user_prompt: str,
        screenshot: bytes | None,
        tools: list[ToolSpec],
        timeout: float,
    ) -> ToolCallDecision:
        try:
            response = self._create_message(
                system_prompt,
                user_prompt,
                screenshot,
                tools,
                timeout=timeout,
                forced=False,
            )
            if not _anthropic_response_has_tool_use(response):
                # tool_choice="any"/"tool" forces Claude to skip any preceding
                # text block, so the reflection instructed by the system
                # prompt requires tool_choice="auto". That leaves a small
                # chance the model responds without calling a tool at all;
                # retry once with tool_choice="any" to guarantee progress.
                # The forced retry cannot carry rationale/thinking (Anthropic
                # rejects thinking combined with forced tool_choice).
                response = self._create_message(
                    system_prompt,
                    user_prompt,
                    screenshot,
                    tools,
                    timeout=timeout,
                    forced=True,
                )
            return _decision_from_anthropic_response(
                response,
                system_prompt=system_prompt,
                user_prompt=user_prompt,
            )
        except ToolCallUnavailable:
            raise
        except Exception as exc:
            raise ToolCallUnavailable(str(exc)) from exc

    def _create_message(
        self,
        system_prompt: str,
        user_prompt: str,
        screenshot: bytes | None,
        tools: list[ToolSpec],
        *,
        timeout: float,
        forced: bool,
    ) -> Any:
        client = self._client()
        # Forced tool_choice is incompatible with extended thinking.
        budget = self._thinking_budget_tokens if not forced else None
        # Enforce max_tokens >= budget + 1 when thinking is enabled.
        max_tokens = self.max_tokens
        if budget is not None:
            max_tokens = max(max_tokens, budget + 1)
        tool_choice_type = "any" if forced else "auto"
        kwargs: dict[str, Any] = {
            "model": self.model,
            "max_tokens": max_tokens,
            "timeout": timeout,
            "system": [
                {
                    "type": "text",
                    "text": system_prompt,
                    "cache_control": {"type": "ephemeral"},
                }
            ],
            "messages": [
                {
                    "role": "user",
                    "content": _anthropic_content(user_prompt, screenshot),
                }
            ],
            "tools": [_anthropic_tool(spec) for spec in tools],
            "tool_choice": {
                "type": tool_choice_type,
                "disable_parallel_tool_use": True,
            },
        }
        if budget is not None:
            kwargs["thinking"] = {"type": "enabled", "budget_tokens": budget}
            kwargs["betas"] = ["interleaved-thinking-2025-05-14"]
        messages = getattr(client, "messages", None)
        if messages is not None:
            return messages.create(**kwargs)
        return client.create(**kwargs)

    def _client(self) -> Any:
        if self._transport is not None:
            return self._transport

        try:
            import anthropic
        except Exception as exc:
            raise ToolCallUnavailable("anthropic SDK is unavailable") from exc

        client_kwargs: dict[str, str] = {}
        if self._api_key is not None:
            client_kwargs["api_key"] = self._api_key
        if self._base_url is not None:
            client_kwargs["base_url"] = self._base_url
        self._transport = anthropic.Anthropic(**client_kwargs)
        return self._transport


class OpenAIToolCallingClient:
    def __init__(
        self,
        *,
        model: str,
        transport: Any | None = None,
        max_tokens: int = 1024,
        api_key: str | None = None,
        base_url: str | None = None,
    ) -> None:
        self.model = model
        self._transport = transport
        self.max_tokens = max_tokens
        self._api_key = api_key
        self._base_url = base_url

    def decide(
        self,
        *,
        system_prompt: str,
        user_prompt: str,
        screenshot: bytes | None,
        tools: list[ToolSpec],
        timeout: float,
    ) -> ToolCallDecision:
        try:
            response = self._create_completion(
                system_prompt,
                user_prompt,
                screenshot,
                tools,
                timeout=timeout,
                forced=False,
            )
            if not _openai_response_has_tool_call(response):
                # tool_choice="required" forces a function call and suppresses
                # any reflection text, so capturing rationale requires
                # tool_choice="auto". Retry once, forcing tool use, if the
                # model responds without calling a tool at all.
                response = self._create_completion(
                    system_prompt,
                    user_prompt,
                    screenshot,
                    tools,
                    timeout=timeout,
                    forced=True,
                )
            return _decision_from_openai_response(
                response,
                system_prompt=system_prompt,
                user_prompt=user_prompt,
            )
        except ToolCallUnavailable:
            raise
        except Exception as exc:
            raise ToolCallUnavailable(str(exc)) from exc

    def _create_completion(
        self,
        system_prompt: str,
        user_prompt: str,
        screenshot: bytes | None,
        tools: list[ToolSpec],
        *,
        timeout: float,
        forced: bool,
    ) -> Any:
        client = self._client()
        kwargs = {
            "model": self.model,
            "max_completion_tokens": self.max_tokens,
            "timeout": timeout,
            "messages": [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": _openai_content(user_prompt, screenshot)},
            ],
            "tools": [_openai_tool(spec) for spec in tools],
            "tool_choice": "required" if forced else "auto",
            "parallel_tool_calls": False,
        }
        chat = getattr(client, "chat", None)
        if chat is not None:
            return chat.completions.create(**kwargs)
        return client.create(**kwargs)

    def _client(self) -> Any:
        if self._transport is not None:
            return self._transport

        try:
            from openai import OpenAI
        except Exception as exc:
            raise ToolCallUnavailable("openai SDK is unavailable") from exc

        kwargs: dict[str, str] = {}
        if self._api_key is not None:
            kwargs["api_key"] = self._api_key
        if self._base_url is not None:
            kwargs["base_url"] = self._base_url
        self._transport = OpenAI(**kwargs)
        return self._transport


def build_client(config: PlannerConfig) -> ToolCallingClient:
    model = config.resolved_model()
    if config.provider == "openai":
        return OpenAIToolCallingClient(model=model)
    return AnthropicToolCallingClient(
        model=model,
        thinking_budget_tokens=config.thinking_budget_tokens,
    )


def _anthropic_content(
    user_prompt: str,
    screenshot: bytes | None,
) -> list[dict[str, Any]]:
    content: list[dict[str, Any]] = []
    if screenshot is not None:
        content.append(
            {
                "type": "image",
                "source": {
                    "type": "base64",
                    "media_type": "image/png",
                    "data": base64.b64encode(screenshot).decode("ascii"),
                },
            }
        )
    content.append({"type": "text", "text": user_prompt})
    return content


def _anthropic_tool(spec: ToolSpec) -> dict[str, Any]:
    return {
        "name": spec.name,
        "description": spec.description,
        "input_schema": spec.parameters,
    }


def _anthropic_response_has_tool_use(response: Any) -> bool:
    content = _value(response, "content")
    if not isinstance(content, list):
        return False
    return any(_value(block, "type") == "tool_use" for block in content)


def _decision_from_anthropic_response(
    response: Any,
    *,
    system_prompt: str = "",
    user_prompt: str = "",
) -> ToolCallDecision:
    content = _value(response, "content")
    if not isinstance(content, list):
        raise ValueError("anthropic tool-call response missing content list")
    thinking: str | None = None
    text_parts: list[str] = []
    for block in content:
        block_type = _value(block, "type")
        if block_type == "thinking" and thinking is None:
            raw = _value(block, "thinking")
            if isinstance(raw, str):
                thinking = raw
        elif block_type == "text":
            raw = _value(block, "text")
            if isinstance(raw, str):
                text_parts.append(raw)
        elif block_type == "tool_use":
            name = _value(block, "name")
            arguments = _value(block, "input")
            if isinstance(name, str) and isinstance(arguments, dict):
                executable_arguments, purpose, expected_outcome = (
                    _split_action_metadata(name, arguments)
                )
                return ToolCallDecision(
                    tool_name=name,
                    arguments=executable_arguments,
                    usage=_anthropic_usage(response),
                    system_prompt=system_prompt,
                    user_prompt=user_prompt,
                    text_output="\n".join(text_parts) if text_parts else None,
                    thinking=thinking,
                    purpose=purpose,
                    expected_outcome=expected_outcome,
                )
    raise ValueError("anthropic response did not include a tool_use block")


def _openai_content(
    user_prompt: str,
    screenshot: bytes | None,
) -> str | list[dict[str, Any]]:
    if screenshot is None:
        return user_prompt
    encoded = base64.b64encode(screenshot).decode("ascii")
    return [
        {"type": "text", "text": user_prompt},
        {
            "type": "image_url",
            "image_url": {"url": f"data:image/png;base64,{encoded}"},
        },
    ]


def _openai_tool(spec: ToolSpec) -> dict[str, Any]:
    return {
        "type": "function",
        "function": {
            "name": spec.name,
            "description": spec.description,
            "parameters": spec.parameters,
        },
    }


def _openai_response_has_tool_call(response: Any) -> bool:
    choices = _value(response, "choices")
    if not isinstance(choices, list) or not choices:
        return False
    message = _value(choices[0], "message")
    tool_calls = _value(message, "tool_calls")
    return isinstance(tool_calls, list) and len(tool_calls) > 0


def _decision_from_openai_response(
    response: Any,
    *,
    system_prompt: str = "",
    user_prompt: str = "",
) -> ToolCallDecision:
    choices = _value(response, "choices")
    if not isinstance(choices, list) or not choices:
        raise ValueError("openai tool-call response missing choices")
    message = _value(choices[0], "message")
    tool_calls = _value(message, "tool_calls")
    if not isinstance(tool_calls, list) or not tool_calls:
        raise ValueError("openai response did not include a tool call")
    function = _value(tool_calls[0], "function")
    name = _value(function, "name")
    if not isinstance(name, str):
        raise ValueError("openai tool call missing a function name")
    arguments = _decode_openai_arguments(_value(function, "arguments"))
    executable_arguments, purpose, expected_outcome = _split_action_metadata(
        name, arguments
    )
    # Extract reasoning_content from o-series models when present.
    raw_reasoning = _value(message, "reasoning_content")
    thinking: str | None = raw_reasoning if isinstance(raw_reasoning, str) else None
    # Pre-tool reflection text (rationale), when the model emits content
    # alongside the tool call under tool_choice="auto".
    raw_content = _value(message, "content")
    text_output: str | None = raw_content if isinstance(raw_content, str) else None
    return ToolCallDecision(
        tool_name=name,
        arguments=executable_arguments,
        usage=_openai_usage(response),
        system_prompt=system_prompt,
        user_prompt=user_prompt,
        thinking=thinking,
        text_output=text_output,
        purpose=purpose,
        expected_outcome=expected_outcome,
    )


def _anthropic_usage(response: Any) -> ToolCallUsage | None:
    usage = _value(response, "usage")
    input_tokens = _integer_value(usage, "input_tokens")
    output_tokens = _integer_value(usage, "output_tokens")
    if input_tokens is None and output_tokens is None:
        return None
    return ToolCallUsage(
        input_tokens=input_tokens,
        output_tokens=output_tokens,
        total_tokens=(input_tokens or 0) + (output_tokens or 0),
    )


def _openai_usage(response: Any) -> ToolCallUsage | None:
    usage = _value(response, "usage")
    input_tokens = _integer_value(usage, "prompt_tokens")
    output_tokens = _integer_value(usage, "completion_tokens")
    total_tokens = _integer_value(usage, "total_tokens")
    if input_tokens is None and output_tokens is None and total_tokens is None:
        return None
    return ToolCallUsage(
        input_tokens=input_tokens,
        output_tokens=output_tokens,
        total_tokens=total_tokens
        if total_tokens is not None
        else (input_tokens or 0) + (output_tokens or 0),
    )


def _decode_openai_arguments(raw_arguments: Any) -> dict[str, Any]:
    if isinstance(raw_arguments, dict):
        return raw_arguments
    if isinstance(raw_arguments, str):
        decoded = json.loads(raw_arguments)
        if isinstance(decoded, dict):
            return decoded
    raise ValueError("openai tool call arguments must decode to a JSON object")


def _split_action_metadata(
    tool_name: str,
    arguments: dict[str, Any],
) -> tuple[dict[str, Any], str | None, str | None]:
    executable_arguments = dict(arguments)
    if tool_name not in ACTION_TOOL_NAMES:
        return executable_arguments, None, None
    return (
        {
            name: value
            for name, value in executable_arguments.items()
            if name not in {"purpose", "expected_outcome"}
        },
        _metadata_text(executable_arguments.get("purpose")),
        _metadata_text(executable_arguments.get("expected_outcome")),
    )


def _metadata_text(value: Any) -> str | None:
    if not isinstance(value, str):
        return None
    return value.strip() or None


def _value(source: Any, key: str) -> Any:
    if isinstance(source, dict):
        return source.get(key)
    value = getattr(source, key, None)
    return None if callable(value) else value


def _integer_value(source: Any, key: str) -> int | None:
    value = _value(source, key)
    return value if isinstance(value, int) and value >= 0 else None