"""
LiteLLM wrapper for public-facing deployments.

Provides the same interface as OpenAIAsGenerativeModel but routes through LiteLLM
to support multiple providers (Google, OpenAI, Anthropic, Cohere, etc.)

Replaces:
- google_wrapper.py
- google_wrapper_embeddings.py  
- llama_wrapper.py
"""

import os
import io
import base64
import json
import logging
from types import SimpleNamespace
from typing import Any, Dict, List, Optional, Union

try:
    from PIL import Image
except ImportError:
    Image = None

try:
    import litellm
    LITELLM_AVAILABLE = True
except ImportError:
    LITELLM_AVAILABLE = False
    litellm = None


def _ensure_gemini_api_key():
    """
    Ensure GEMINI_API_KEY is set for LiteLLM.
    Falls back to GOOGLE_API_KEY for backward compatibility.
    """
    if not os.environ.get("GEMINI_API_KEY"):
        google_key = os.environ.get("GOOGLE_API_KEY")
        if google_key:
            os.environ["GEMINI_API_KEY"] = google_key

def _normalize_model_name(model: str) -> str:
    """
    Add provider prefix if not present.
    
    LiteLLM requires prefixes like 'gemini/', 'openai/', etc. to route correctly.
    This helper auto-detects common model patterns and adds the appropriate prefix.
    """
    if not model:
        return model
        
    if '/' in model:
        # Already has prefix
        return model
    
    model_lower = model.lower()
    
    # Google Gemini models
    if 'gemini' in model_lower:
        return f"gemini/{model}"
    
    # OpenAI models
    if model_lower.startswith(('gpt-', 'o1-', 'o3-', 'text-embedding', 'davinci', 'curie', 'babbage', 'ada')):
        return f"openai/{model}"
    
    # Anthropic models
    if model_lower.startswith('claude'):
        return f"anthropic/{model}"
    
    # Default: return as-is, let LiteLLM figure it out
    return model

def _check_litellm():
    """Raise ImportError if LiteLLM is not available."""
    if not LITELLM_AVAILABLE:
        raise ImportError(
            "LiteLLM is required for public deployments. "
            "Install with: pip install litellm"
        )
    
    _ensure_gemini_api_key()
    
    # Suppress verbose logging
    import logging
    logging.getLogger("LiteLLM").setLevel(logging.WARNING)
    logging.getLogger("litellm").setLevel(logging.WARNING)
    
    litellm.suppress_debug_info = True


class LiteLLMGenerativeModel:
    """
    LiteLLM-based model that matches the OpenAIAsGenerativeModel interface.
    
    Supports 100+ LLM providers through LiteLLM's unified API.
    
    Usage:
        # Google Gemini
        model = LiteLLMGenerativeModel("gemini/gemini-2.0-flash", api_key="...")
        
        # OpenAI
        model = LiteLLMGenerativeModel("gpt-4o", api_key="...")
        
        # Anthropic
        model = LiteLLMGenerativeModel("claude-sonnet-4-20250514", api_key="...")
        
        # Local Ollama
        model = LiteLLMGenerativeModel("ollama/llama3")
        
        # All use the same interface:
        response = model.generate_content(["Hello!", image])
        print(response.text)
    
    Model Name Format:
        LiteLLM uses provider prefixes for routing:
        - Google: "gemini/gemini-2.0-flash"
        - OpenAI: "gpt-4o" or "openai/gpt-4o"
        - Anthropic: "claude-sonnet-4-20250514" or "anthropic/claude-sonnet-4-20250514"
        - Azure: "azure/deployment-name"
        - Ollama: "ollama/llama3"
        
        See: https://docs.litellm.ai/docs/providers
    """
    
    def __init__(
        self,
        model: str,
        api_key: Optional[str] = None,
        base_url: Optional[str] = None,
        system_instruction: Optional[str] = None,
        tools: Optional[List] = None,
    ):
        """
        Initialize the LiteLLM model.
        
        Args:
            model: LiteLLM model string (e.g., "gemini/gemini-2.0-flash")
            api_key: API key for the provider
            base_url: Optional custom API base URL
            system_instruction: Default system message
            tools: Default tools for function calling
        """
        _check_litellm()
        
        self.model = _normalize_model_name(model)
        self.api_key = api_key
        self.base_url = base_url
        self.system_instruction = system_instruction
        self.tools = tools
    
    def generate_content(
        self,
        contents: Union[str, List],
        generation_config: Any = None,
        safety_settings: Any = None,
        stream: bool = False,
        tools: Optional[List] = None,
    ):
        """
        Generate content - matches OpenAIAsGenerativeModel interface.
        
        Args:
            contents: Prompt string or list of parts (text, images, dicts)
            generation_config: Generation parameters (temperature, max_tokens, etc.)
            safety_settings: Ignored (kept for interface compatibility)
            stream: If True, return streaming iterator
            tools: Tools for function calling (overrides default)
        
        Returns:
            Response object with .text and .candidates attributes
        """
        messages = self._build_messages(contents)
        params = self._build_params(generation_config, tools or self.tools)
        
        try:
            response = litellm.completion(
                model=self.model,
                messages=messages,
                api_key=self.api_key,
                api_base=self.base_url,
                stream=stream,
                **params
            )
            
            if stream:
                return self._handle_stream(response)
            
            return self._to_legacy_response(response)
            
        except Exception as e:
            logging.error(f"LiteLLM generation error: {e}")
            raise
    
    def _handle_stream(self, stream):
        """Yield legacy-format responses from stream."""
        for chunk in stream:
            yield self._to_legacy_response(chunk, is_stream=True)
    
    def start_chat(
        self,
        history: Optional[List] = None,
        enable_automatic_function_calling: bool = False
    ) -> 'LiteLLMChatSession':
        """
        Start a chat session.
        
        Args:
            history: Optional conversation history
            enable_automatic_function_calling: Enable auto function calling
        
        Returns:
            LiteLLMChatSession instance
        """
        return LiteLLMChatSession(
            model=self,
            history=history,
            enable_afc=enable_automatic_function_calling
        )
    
    def _build_messages(self, contents: Union[str, List]) -> List[Dict]:
        """Convert contents to LiteLLM/OpenAI message format."""
        messages = []
        
        if self.system_instruction:
            messages.append({"role": "system", "content": self.system_instruction})
        
        if isinstance(contents, str):
            messages.append({"role": "user", "content": contents})
        elif isinstance(contents, list):
            user_content = self._convert_parts(contents)
            messages.append({"role": "user", "content": user_content})
        else:
            messages.append({"role": "user", "content": str(contents)})
        
        return messages
    
    def _convert_parts(self, parts: List) -> Union[str, List[Dict]]:
        """Convert list of parts (text, images, dicts) to OpenAI content format."""
        converted = []
        
        for part in parts:
            if isinstance(part, str):
                converted.append({"type": "text", "text": part})
            
            elif Image and isinstance(part, Image.Image):
                buf = io.BytesIO()
                part.save(buf, format="PNG")
                b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
                converted.append({
                    "type": "image_url",
                    "image_url": {"url": f"data:image/png;base64,{b64}"}
                })
            
            elif isinstance(part, dict):
                if "mime_type" in part and "data" in part:
                    data = part["data"]
                    if isinstance(data, bytes):
                        b64 = base64.b64encode(data).decode("utf-8")
                    else:
                        b64 = data
                    converted.append({
                        "type": "image_url",
                        "image_url": {"url": f"data:{part['mime_type']};base64,{b64}"}
                    })
                elif "text" in part:
                    converted.append({"type": "text", "text": part["text"]})
                else:
                    converted.append({"type": "text", "text": str(part)})
            
            else:
                converted.append({"type": "text", "text": str(part)})
        
        # Simplify if only text parts
        if all(p.get("type") == "text" for p in converted):
            return " ".join(p["text"] for p in converted)
        
        return converted
    
    def _build_params(self, generation_config: Any, tools: Optional[List]) -> Dict:
        """Build LiteLLM parameters from generation config."""
        params = {}
        
        if generation_config:
            if isinstance(generation_config, dict):
                cfg = generation_config
            elif hasattr(generation_config, '__dict__'):
                cfg = vars(generation_config)
            else:
                cfg = {}
            
            mapping = {
                "temperature": "temperature",
                "top_p": "top_p",
                "top_k": "top_k",
                "max_output_tokens": "max_tokens",
                "stop_sequences": "stop",
                "presence_penalty": "presence_penalty",
                "frequency_penalty": "frequency_penalty",
            }
            
            for old, new in mapping.items():
                val = cfg.get(old)
                if val is not None:
                    params[new] = val
        
        if tools:
            params["tools"] = tools
        
        return params
    
    def _to_legacy_response(self, response, is_stream: bool = False):
        """Convert LiteLLM response to legacy format."""
        candidates = []
        
        choices = getattr(response, "choices", None) or []
        
        if not choices:
            logging.warning(f"LiteLLM response has no choices")
            return SimpleNamespace(
                text="",
                candidates=[SimpleNamespace(
                    content=SimpleNamespace(parts=[], role="model"),
                    finish_reason=1
                )]
            )
        
        for choice in choices:
            parts = []
            
            message = getattr(choice, "message", None)
            if message is None:
                message = getattr(choice, "delta", None)
            
            text = ""
            if message:
                content = getattr(message, "content", None)
                if content:
                    text = self._clean_text(str(content))
            
            tool_calls = getattr(message, "tool_calls", None) if message else None
            if tool_calls:
                for tc in tool_calls:
                    func = getattr(tc, "function", None)
                    if func:
                        try:
                            args = json.loads(getattr(func, "arguments", "{}"))
                        except json.JSONDecodeError:
                            args = {}
                        parts.append(SimpleNamespace(
                            function_call=SimpleNamespace(
                                name=getattr(func, "name", ""),
                                args=args
                            ),
                            text=None
                        ))
            
            if text:
                parts.append(SimpleNamespace(text=text, function_call=None))
            
            finish_reason = getattr(choice, "finish_reason", "stop")
            finish_map = {"stop": 1, "length": 0, "tool_calls": 2, "content_filter": 3}
            fr_int = finish_map.get(finish_reason, 1) if finish_reason else 1
            
            content = SimpleNamespace(parts=parts, role="model")
            candidates.append(SimpleNamespace(content=content, finish_reason=fr_int))
        
        first_text = ""
        if candidates and candidates[0].content.parts:
            for part in candidates[0].content.parts:
                if hasattr(part, "text") and part.text:
                    first_text += part.text
        
        return SimpleNamespace(text=first_text, candidates=candidates)
    
    def _clean_text(self, text: str) -> str:
        """
        Extract JSON from LLM response, handling preamble text and code fences.
        """
        if not text:
            return ""
        
        text = text.strip()
        
        import re
        
        # Pattern 1: ```json ... ``` anywhere in text
        json_match = re.search(r'```json\s*([\s\S]*?)\s*```', text)
        if json_match:
            return json_match.group(1).strip()
        
        # Pattern 2: ``` ... ``` anywhere in text
        code_match = re.search(r'```\s*([\s\S]*?)\s*```', text)
        if code_match:
            extracted = code_match.group(1).strip()
            if extracted.startswith('{') or extracted.startswith('['):
                return extracted
        
        # Pattern 3: Raw JSON object anywhere in text
        json_obj_match = re.search(r'(\{[\s\S]*\})', text)
        if json_obj_match:
            potential = json_obj_match.group(1)
            if '"' in potential and ':' in potential:
                return potential
        
        return text


class LiteLLMChatSession:
    """
    Chat session that maintains conversation history.
    
    Matches the interface of the legacy Google GenAI ChatSession.
    """
    
    def __init__(
        self,
        model: LiteLLMGenerativeModel,
        history: Optional[List] = None,
        enable_afc: bool = False
    ):
        self._model = model
        self._history = self._convert_history(history) if history else []
        self._enable_afc = enable_afc
    
    def _convert_history(self, legacy_history: List) -> List[Dict]:
        """Convert legacy history format to messages."""
        messages = []
        for entry in legacy_history:
            if isinstance(entry, dict):
                role = entry.get("role", "user")
                parts = entry.get("parts", [])
            else:
                role = getattr(entry, "role", "user")
                parts = getattr(entry, "parts", [])
            
            text_parts = []
            for part in parts:
                if isinstance(part, str):
                    text_parts.append(part)
                elif hasattr(part, "text"):
                    text_parts.append(part.text)
                else:
                    text_parts.append(str(part))
            
            messages.append({"role": role, "content": " ".join(text_parts)})
        
        return messages
    
    @property
    def history(self) -> List:
        """Return conversation history."""
        return self._history
    
    def send_message(
        self,
        content: Union[str, List],
        generation_config: Any = None,
        stream: bool = False
    ):
        """
        Send message and update history.
        
        Args:
            content: Message content (string or list of parts)
            generation_config: Optional generation config
            stream: If True, return streaming iterator
        
        Returns:
            Response object with .text and .candidates
        """
        messages = []
        
        if self._model.system_instruction:
            messages.append({"role": "system", "content": self._model.system_instruction})
        
        messages.extend(self._history)
        
        if isinstance(content, str):
            user_content = content
        else:
            user_content = self._model._convert_parts(content)
        
        messages.append({"role": "user", "content": user_content})
        
        params = self._model._build_params(generation_config, self._model.tools)
        
        response = litellm.completion(
            model=self._model.model,
            messages=messages,
            api_key=self._model.api_key,
            api_base=self._model.base_url,
            stream=stream,
            **params
        )
        
        if stream:
            return self._handle_stream(response, user_content)
        
        self._history.append({"role": "user", "content": user_content})
        
        legacy_response = self._model._to_legacy_response(response)
        
        if legacy_response.text:
            self._history.append({"role": "assistant", "content": legacy_response.text})
        
        return legacy_response
    
    def _handle_stream(self, stream, user_content):
        """Handle streaming response."""
        full_text = ""
        for chunk in stream:
            legacy_chunk = self._model._to_legacy_response(chunk, is_stream=True)
            full_text += legacy_chunk.text
            yield legacy_chunk
        
        self._history.append({"role": "user", "content": user_content})
        self._history.append({"role": "assistant", "content": full_text})


class LiteLLMEmbeddingModel:
    """
    LiteLLM-based embeddings matching OpenAIAsEmbeddingModel interface.
    
    Features:
    - Automatic provider prefix detection (gemini/, openai/, etc.)
    - Size-aware batching to avoid payload limits
    - Automatic retry with smaller batches on size errors
    
    Usage:
        # Google
        embedder = LiteLLMEmbeddingModel("gemini-embedding-001", api_key="...")
        
        # OpenAI
        embedder = LiteLLMEmbeddingModel("text-embedding-3-small", api_key="...")
        
        response = embedder.embed_content(
            model="...",
            content=["Hello", "World"]
        )
        print(response["embedding"])  # [[...], [...]]
    """
    
    # Max payload size ~30MB to stay under 40MB limit with overhead
    MAX_BATCH_BYTES = 30 * 1024 * 1024
    # Also limit by count to avoid other API limits
    MAX_BATCH_SIZE = 100
    
    def __init__(
        self,
        model: Optional[str] = None,
        api_key: Optional[str] = None,
        base_url: Optional[str] = None
    ):
        """
        Initialize the embedding model.
        
        Args:
            model: Default model name (can be overridden in embed_content)
            api_key: API key for the provider
            base_url: Optional custom API base URL
        """
        _check_litellm()
        
        self.model = _normalize_model_name(model) if model else None
        self.api_key = api_key
        self.base_url = base_url
    
    def embed_content(
        self,
        model: Optional[str] = None,
        content: Union[str, List[str]] = None,
        task_type: Optional[str] = None,
        **kwargs
    ) -> Dict[str, Any]:
        """
        Generate embeddings with automatic size-aware batching.
        
        Args:
            model: Model name (overrides default)
            content: Text or list of texts to embed
            task_type: Ignored (kept for interface compatibility)
        
        Returns:
            Dict with 'embedding' key containing vector(s)
        """
        effective_model = model or self.model
        if not effective_model:
            raise ValueError("Model must be specified")
        
        effective_model = _normalize_model_name(effective_model)
        
        is_single = isinstance(content, str)
        inputs = [content] if is_single else content
        
        try:
            # Create size-aware batches to avoid payload limits
            batches = self._create_size_aware_batches(inputs)
            
            all_embeddings = []
            for batch_idx, batch in enumerate(batches):
                try:
                    response = litellm.embedding(
                        model=effective_model,
                        input=batch,
                        api_key=self.api_key,
                        api_base=self.base_url
                    )
                    
                    batch_embeddings = [item["embedding"] for item in response.data]
                    all_embeddings.extend(batch_embeddings)
                    
                except Exception as e:
                    error_str = str(e).lower()
                    # Check for payload size errors
                    if "payload size" in error_str or "too large" in error_str or "exceeds the limit" in error_str:
                        logging.warning(f"Batch {batch_idx + 1} too large, splitting and retrying...")
                        # Recursively embed with smaller batches
                        sub_embeddings = self._embed_with_retry(effective_model, batch)
                        all_embeddings.extend(sub_embeddings)
                    else:
                        raise
            
            if is_single:
                return {"embedding": all_embeddings[0]}
            return {"embedding": all_embeddings}
            
        except Exception as e:
            logging.error(f"LiteLLM embedding error: {e}")
            raise
    
    def _create_size_aware_batches(self, inputs: List[str]) -> List[List[str]]:
        """
        Split inputs into batches that won't exceed payload limits.
        
        Uses both byte size and count limits to stay within API constraints.
        """
        batches = []
        current_batch = []
        current_size = 0
        
        for text in inputs:
            text_size = len(text.encode('utf-8'))
            
            # Check if adding this text would exceed limits
            would_exceed_size = current_size + text_size > self.MAX_BATCH_BYTES
            would_exceed_count = len(current_batch) >= self.MAX_BATCH_SIZE
            
            if (would_exceed_size or would_exceed_count) and current_batch:
                batches.append(current_batch)
                current_batch = []
                current_size = 0
            
            # Handle individual texts that are too large
            if text_size > self.MAX_BATCH_BYTES:
                logging.warning(
                    f"Single text too large ({text_size:,} bytes), truncating to fit limit..."
                )
                # Truncate to fit - estimate conservatively
                max_chars = (self.MAX_BATCH_BYTES * 3) // 4  # Account for UTF-8 overhead
                text = text[:max_chars]
                text_size = len(text.encode('utf-8'))
            
            current_batch.append(text)
            current_size += text_size
        
        if current_batch:
            batches.append(current_batch)
        
        return batches
    
    def _embed_with_retry(
        self, 
        model: str, 
        texts: List[str], 
        max_retries: int = 3
    ) -> List[List[float]]:
        """
        Recursively embed texts, splitting on size errors.
        
        If a batch fails due to size, split it in half and retry.
        """
        if not texts:
            return []
        
        # Base case: single text
        if len(texts) == 1:
            try:
                response = litellm.embedding(
                    model=model,
                    input=texts,
                    api_key=self.api_key,
                    api_base=self.base_url
                )
                return [item["embedding"] for item in response.data]
            except Exception as e:
                error_str = str(e).lower()
                if "payload size" in error_str or "too large" in error_str:
                    # Single text is too large - truncate it
                    logging.warning("Single text still too large, truncating further...")
                    truncated = texts[0][:len(texts[0]) // 2]
                    response = litellm.embedding(
                        model=model,
                        input=[truncated],
                        api_key=self.api_key,
                        api_base=self.base_url
                    )
                    return [item["embedding"] for item in response.data]
                raise
        
        # Try the full batch first
        try:
            response = litellm.embedding(
                model=model,
                input=texts,
                api_key=self.api_key,
                api_base=self.base_url
            )
            return [item["embedding"] for item in response.data]
        except Exception as e:
            error_str = str(e).lower()
            if "payload size" in error_str or "too large" in error_str or "exceeds the limit" in error_str:
                # Split in half and retry each half
                mid = len(texts) // 2
                left_half = texts[:mid]
                right_half = texts[mid:]
                
                logging.info(f"Splitting batch of {len(texts)} into {len(left_half)} + {len(right_half)}")
                
                left_embeddings = self._embed_with_retry(model, left_half, max_retries)
                right_embeddings = self._embed_with_retry(model, right_half, max_retries)
                
                return left_embeddings + right_embeddings
            raise