"""
Gemini Provider Module

Google Gemini/AI Studio API provider:
- Gemini 1.5 Pro, Gemini 1.5 Flash, Gemini 2.0
- Tool/function calling support
"""

import json
import time
from typing import Optional

import httpx

from ai_coder.llm.interface import (
    LLMProvider,
    LLMConfig,
    Message,
    ToolCall,
    CompletionResponse,
    LLMError,
    RateLimitError,
    AuthenticationError,
    register_provider,
)


@register_provider("gemini")
class GeminiProvider(LLMProvider):
    """
    Google Gemini API provider.
    
    Supports:
    - Gemini 2.0 Flash
    - Gemini 1.5 Pro
    - Gemini 1.5 Flash
    """

    # Model name mapping - use -latest suffix for stable versions
    MODEL_MAP = {
        "gemini-pro": "gemini-1.5-pro-latest",
        "gemini-flash": "gemini-1.5-flash-latest",
        "gemini-2": "gemini-2.0-flash-001",
        "gemini-2.0": "gemini-2.0-flash-001",
        "gemini-2.0-flash": "gemini-2.0-flash-001",
        "gemini-2.0-flash-exp": "gemini-2.0-flash-001",
        "gemini-1.5-pro": "gemini-1.5-pro-latest",
        "gemini-1.5-flash": "gemini-1.5-flash-latest",
    }

    def __init__(self, config: LLMConfig):
        super().__init__(config)
        
        self.api_key = config.api_key
        if not self.api_key:
            raise AuthenticationError("GEMINI_API_KEY is required")
        
        # Use provided base or default to Google AI Studio
        self.base_url = config.api_base or "https://generativelanguage.googleapis.com/v1beta"
        self.base_url = self.base_url.rstrip("/")
        
        self.client = httpx.Client(timeout=config.timeout)

    def complete(
        self,
        messages: list[Message],
        tools: Optional[list[dict]] = None,
    ) -> CompletionResponse:
        """Generate completion using Gemini API."""
        
        # Map model name if needed
        model = self.MODEL_MAP.get(self.config.model, self.config.model)
        
        # Convert messages to Gemini format
        contents, system_instruction = self._format_messages(messages)
        
        payload = {
            "contents": contents,
            "generationConfig": {
                "temperature": self.config.temperature,
                "maxOutputTokens": self.config.max_tokens,
            }
        }
        
        if system_instruction:
            payload["systemInstruction"] = {"parts": [{"text": system_instruction}]}
        
        # Convert tools to Gemini format
        if tools:
            payload["tools"] = [{"functionDeclarations": self._convert_tools(tools)}]

        url = f"{self.base_url}/models/{model}:generateContent?key={self.api_key}"

        for attempt in range(self.config.max_retries):
            try:
                response = self.client.post(url, json=payload)
                
                if response.status_code == 401 or response.status_code == 403:
                    raise AuthenticationError("Gemini API authentication failed")
                
                if response.status_code == 429:
                    if attempt < self.config.max_retries - 1:
                        wait_time = 2 ** attempt
                        time.sleep(wait_time)
                        continue
                    raise RateLimitError("Rate limit exceeded")
                
                if response.status_code >= 400:
                    error_text = response.text[:500]
                    raise LLMError(f"Gemini API error ({response.status_code}): {error_text}")
                
                return self._parse_response(response.json())

            except httpx.TimeoutException:
                if attempt < self.config.max_retries - 1:
                    time.sleep(1)
                    continue
                raise LLMError("Request timed out")
            
            except httpx.RequestError as e:
                if attempt < self.config.max_retries - 1:
                    time.sleep(1)
                    continue
                raise LLMError(f"Connection error: {e}")

        raise LLMError("Max retries exceeded")

    def _format_messages(self, messages: list[Message]) -> tuple[list[dict], str]:
        """Convert messages to Gemini format."""
        contents = []
        system_instruction = ""
        
        for msg in messages:
            if msg.role == "system":
                system_instruction = msg.content
            elif msg.role == "user":
                contents.append({
                    "role": "user",
                    "parts": [{"text": msg.content}]
                })
            elif msg.role == "assistant":
                parts = []
                if msg.content:
                    parts.append({"text": msg.content})
                if msg.tool_calls:
                    for tc in msg.tool_calls:
                        func = tc.get("function", tc)
                        parts.append({
                            "functionCall": {
                                "name": func.get("name", tc.get("name", "")),
                                "args": func.get("arguments", tc.get("arguments", {}))
                            }
                        })
                if parts:
                    contents.append({"role": "model", "parts": parts})
            elif msg.role == "tool":
                contents.append({
                    "role": "user",
                    "parts": [{
                        "functionResponse": {
                            "name": msg.name or "function",
                            "response": {"result": msg.content}
                        }
                    }]
                })
        
        return contents, system_instruction

    def _convert_tools(self, openai_tools: list[dict]) -> list[dict]:
        """Convert OpenAI tool format to Gemini format."""
        gemini_tools = []
        
        for tool in openai_tools:
            if tool.get("type") == "function":
                func = tool["function"]
                gemini_tools.append({
                    "name": func["name"],
                    "description": func.get("description", ""),
                    "parameters": func.get("parameters", {"type": "object", "properties": {}})
                })
        
        return gemini_tools

    def _parse_response(self, data: dict) -> CompletionResponse:
        """Parse Gemini API response."""
        candidates = data.get("candidates", [])
        if not candidates:
            return CompletionResponse(content="No response generated")
        
        candidate = candidates[0]
        content_parts = candidate.get("content", {}).get("parts", [])
        
        text_parts = []
        tool_calls = []
        
        for part in content_parts:
            if "text" in part:
                text_parts.append(part["text"])
            elif "functionCall" in part:
                fc = part["functionCall"]
                tool_calls.append(ToolCall(
                    id=f"call_{len(tool_calls)}",
                    name=fc.get("name", ""),
                    arguments=fc.get("args", {}),
                ))
        
        content = "\n".join(text_parts) if text_parts else None
        
        usage = None
        if "usageMetadata" in data:
            meta = data["usageMetadata"]
            usage = {
                "prompt_tokens": meta.get("promptTokenCount", 0),
                "completion_tokens": meta.get("candidatesTokenCount", 0),
                "total_tokens": meta.get("totalTokenCount", 0),
            }
        
        return CompletionResponse(
            content=content,
            tool_calls=tool_calls,
            finish_reason=candidate.get("finishReason", "STOP"),
            usage=usage,
            raw_response=data,
        )

    def count_tokens(self, text: str) -> int:
        """Estimate token count for Gemini models."""
        # Gemini: ~4 chars per token
        return len(text) // 4
