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 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


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,
    ) -> None:
        self.model = model
        self._transport = transport
        self.max_tokens = max_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,
            )
            return _decision_from_anthropic_response(response)
        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,
    ) -> Any:
        client = self._client()
        kwargs = {
            "model": self.model,
            "max_tokens": self.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": "any", "disable_parallel_tool_use": True},
        }
        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

        self._transport = anthropic.Anthropic()
        return self._transport


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

    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,
            )
            return _decision_from_openai_response(response)
        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,
    ) -> 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",
            "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

        self._transport = OpenAI()
        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)


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 _decision_from_anthropic_response(response: Any) -> ToolCallDecision:
    content = _value(response, "content")
    if not isinstance(content, list):
        raise ValueError("anthropic tool-call response missing content list")
    for block in content:
        if _value(block, "type") != "tool_use":
            continue
        name = _value(block, "name")
        arguments = _value(block, "input")
        if isinstance(name, str) and isinstance(arguments, dict):
            return ToolCallDecision(
                tool_name=name,
                arguments=arguments,
                usage=_anthropic_usage(response),
            )
    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 _decision_from_openai_response(response: Any) -> 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"))
    return ToolCallDecision(
        tool_name=name,
        arguments=arguments,
        usage=_openai_usage(response),
    )


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 _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