"""Google Gemini provider — OpenAI-compatible API."""

from __future__ import annotations

import logging
from typing import Any

from qanot.providers.openai import OpenAIProvider, _anthropic_tools_to_openai, _convert_messages

logger = logging.getLogger(__name__)

# Gemini pricing per million tokens (March 2026)
GEMINI_PRICING = {
    # Gemini 3.x series (latest)
    "gemini-3.1-pro-preview": {"input": 2.0, "output": 12.0},
    "gemini-3.1-flash-lite": {"input": 0.25, "output": 1.50},
    "gemini-3-flash-preview": {"input": 0.15, "output": 0.60},
    # Gemini 2.x series (stable)
    "gemini-2.5-pro": {"input": 1.25, "output": 10.0},
    "gemini-2.5-flash": {"input": 0.15, "output": 0.60},
    "gemini-2.0-flash": {"input": 0.10, "output": 0.40},
}

# JSON Schema keywords that Gemini does not support in tool definitions.
_UNSUPPORTED_SCHEMA_KEYS = frozenset({"patternProperties", "additionalProperties", "$ref"})


def _strip_unsupported_keys(obj: Any) -> Any:
    """Recursively remove unsupported JSON Schema keywords from a structure."""
    if isinstance(obj, dict):
        has_unsupported = any(k in _UNSUPPORTED_SCHEMA_KEYS for k in obj)
        if has_unsupported:
            return {
                k: _strip_unsupported_keys(v)
                for k, v in obj.items()
                if k not in _UNSUPPORTED_SCHEMA_KEYS
            }
        # No keys to strip; only recurse into values if needed
        needs_recursion = any(isinstance(v, (dict, list)) for v in obj.values())
        if not needs_recursion:
            return obj
        return {k: _strip_unsupported_keys(v) for k, v in obj.items()}
    if isinstance(obj, list):
        return [_strip_unsupported_keys(item) for item in obj]
    return obj


def _ensure_user_first(messages: list[dict]) -> list[dict]:
    """Prepend a synthetic user turn if the first message is not from a user.

    Gemini requires the conversation to start with a user message.
    If the first message is a system message followed by a non-user message,
    or if there are no user messages at the start, a synthetic user turn is
    inserted right after any system messages.
    """
    if not messages:
        return messages

    # Find the index of the first non-system message
    insert_idx = 0
    for i, msg in enumerate(messages):
        if msg.get("role") != "system":
            insert_idx = i
            break
    else:
        # All messages are system messages
        insert_idx = len(messages)

    # Check if the first non-system message is already a user message
    if insert_idx < len(messages) and messages[insert_idx].get("role") == "user":
        return messages

    synthetic = {"role": "user", "content": "(session start)"}
    return messages[:insert_idx] + [synthetic] + messages[insert_idx:]


def _sanitize_gemini_tools(tools: list[dict]) -> list[dict]:
    """Convert Anthropic-style tools to OpenAI format and strip unsupported keys."""
    openai_tools = _anthropic_tools_to_openai(tools)
    return _strip_unsupported_keys(openai_tools)


class GeminiProvider(OpenAIProvider):
    """Google Gemini provider via OpenAI-compatible API.

    Gemini supports the OpenAI chat completions format through
    generativelanguage.googleapis.com, so we reuse OpenAIProvider
    and only override base_url, pricing, and message/tool conversion.
    """

    def __init__(
        self,
        api_key: str,
        model: str = "gemini-2.5-flash",
        base_url: str = "https://generativelanguage.googleapis.com/v1beta/openai/",
    ):
        import openai
        self.client = openai.AsyncOpenAI(api_key=api_key, base_url=base_url)
        self.model = model

    def _calc_cost(self, inp: int, out: int) -> float:
        try:
            price_in, price_out = self._cached_pricing
        except AttributeError:
            prices = GEMINI_PRICING.get(self.model)
            if prices is not None:
                price_in, price_out = prices["input"], prices["output"]
            else:
                price_in, price_out = 0.15, 0.60
            self._cached_pricing = (price_in, price_out)
        return inp * price_in / 1_000_000 + out * price_out / 1_000_000

    def _prepare_messages(self, messages: list[dict], system: str | None) -> list[dict]:
        """Convert messages and ensure conversation starts with a user turn."""
        return _ensure_user_first(_convert_messages(messages, system))

    def _prepare_tools(self, tools: list[dict]) -> list[dict]:
        """Convert tools and strip Gemini-unsupported JSON Schema keys."""
        return _sanitize_gemini_tools(tools)
