import os

import httpx

from exchange import Message, Tool
from exchange.content import Text, ToolResult, ToolUse
from exchange.providers.base import Provider, Usage
from tenacity import retry, wait_fixed, stop_after_attempt
from exchange.providers.utils import raise_for_status, retry_if_status, encode_image
from exchange.langfuse_wrapper import observe_wrapper


GOOGLE_HOST = "https://generativelanguage.googleapis.com/v1beta"

retry_procedure = retry(
    wait=wait_fixed(2),
    stop=stop_after_attempt(2),
    retry=retry_if_status(codes=[429], above=500),
    reraise=True,
)


class GoogleProvider(Provider):
    """Provides chat completions for models hosted by Google, including Gemini and other experimental models."""

    PROVIDER_NAME = "google"
    REQUIRED_ENV_VARS = ["GOOGLE_API_KEY"]
    instructions_url = "https://ai.google.dev/gemini-api/docs/api-key"

    def __init__(self, client: httpx.Client) -> None:
        self.client = client

    @classmethod
    def from_env(cls: type["GoogleProvider"]) -> "GoogleProvider":
        cls.check_env_vars(cls.instructions_url)
        url = os.environ.get("GOOGLE_HOST", GOOGLE_HOST)
        key = os.environ.get("GOOGLE_API_KEY")
        client = httpx.Client(
            base_url=url,
            headers={
                "Content-Type": "application/json",
            },
            params={"key": key},
            timeout=httpx.Timeout(60 * 10),
        )
        return cls(client)

    @staticmethod
    def get_usage(data: dict) -> Usage:  # noqa: ANN401
        usage = data.get("usageMetadata")
        input_tokens = usage.get("promptTokenCount")
        output_tokens = usage.get("candidatesTokenCount")
        total_tokens = usage.get("totalTokenCount")

        if total_tokens is None and input_tokens is not None and output_tokens is not None:
            total_tokens = input_tokens + output_tokens

        return Usage(
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            total_tokens=total_tokens,
        )

    @staticmethod
    def google_response_to_message(response: dict) -> Message:
        candidates = response.get("candidates", [])
        if candidates:
            # Only use first candidate for now
            candidate = candidates[0]
            content_parts = candidate.get("content", {}).get("parts", [])
            content = []
            for part in content_parts:
                if "text" in part:
                    content.append(Text(text=part["text"]))
                elif "functionCall" in part:
                    content.append(
                        ToolUse(
                            id=part["functionCall"].get("name", ""),
                            name=part["functionCall"].get("name", ""),
                            parameters=part["functionCall"].get("args", {}),
                        )
                    )
            return Message(role="assistant", content=content)

        # If no valid candidates were found, return an empty message
        return Message(role="assistant", content=[])

    @staticmethod
    def tools_to_google_spec(tools: tuple[Tool, ...]) -> dict[str, list[dict[str, any]]]:
        if not tools:
            return {}
        converted_tools = []
        for tool in tools:
            converted_tool: dict[str, any] = {
                "name": tool.name,
                "description": tool.description or "",
            }
            if tool.parameters["properties"]:
                converted_tool["parameters"] = tool.parameters
            converted_tools.append(converted_tool)
        return {"functionDeclarations": converted_tools}

    @staticmethod
    def messages_to_google_spec(messages: list[Message]) -> list[dict[str, any]]:
        messages_spec = []
        for message in messages:
            role = "user" if message.role == "user" else "model"
            converted = {"role": role, "parts": []}
            for content in message.content:
                if isinstance(content, Text):
                    converted["parts"].append({"text": content.text})
                elif isinstance(content, ToolUse):
                    converted["parts"].append({"functionCall": {"name": content.name, "args": content.parameters}})
                elif isinstance(content, ToolResult):
                    if content.output.startswith('"image:'):
                        image_path = content.output.replace('"image:', "").replace('"', "")
                        converted["parts"].append(
                            {
                                "inline_data": {
                                    "mime_type": "image/png",
                                    "data": f"{encode_image(image_path)}",
                                }
                            }
                        )
                    else:
                        converted["parts"].append(
                            {"functionResponse": {"name": content.tool_use_id, "response": {"content": content.output}}}
                        )
            messages_spec.append(converted)

        if not messages_spec:
            messages_spec.append({"role": "user", "parts": [{"text": "Ignore"}]})

        return messages_spec

    @observe_wrapper(as_type="generation")
    def complete(
        self,
        model: str,
        system: str,
        messages: list[Message],
        tools: list[Tool] = None,
        **kwargs: dict[str, any],
    ) -> tuple[Message, Usage]:
        tools_set = set()
        unique_tools = []
        for tool in tools:
            if tool.name not in tools_set:
                unique_tools.append(tool)
                tools_set.add(tool.name)

        payload = dict(
            system_instruction={"parts": [{"text": system}]},
            contents=self.messages_to_google_spec(messages),
            tools=self.tools_to_google_spec(tuple(unique_tools)),
            **kwargs,
        )
        payload = {k: v for k, v in payload.items() if v}
        response = self._post(payload, model)
        message = self.google_response_to_message(response)
        usage = self.get_usage(response)
        return message, usage

    @retry_procedure
    def _post(self, payload: dict, model: str) -> httpx.Response:
        response = self.client.post("models/" + model + ":generateContent", json=payload)
        return raise_for_status(response).json()
