from typing import List, Dict, Any, Optional, Tuple
import os
import requests
import time
import json
import logging
import traceback
from .constants import MAX_HISTORY_LENGTH, load_config

logger = logging.getLogger(__name__)

class LLMClient:
    """
    Client for interacting with LLM models via KeyPool service.
    Handles conversation history and API requests.
    """
    
    def __init__(self, model_name: str, keypool_url: str=None, api_key: str=None):
        """
        Initialize LLM client.
        
        Args:
            model_name: Name of the model to use (e.g., "gemini-2.5-flash-lite")
                        Must match a model_name in KeyPool's keys.yaml
        """
        self.history: List[Dict[str, Any]] = []
        
        # Default timeout for requests
        self.timeout: int = 120
        
        # KeyPool Service Configuration
        config = load_config() if (keypool_url is None or api_key is None) else {}
        if model_name is None:
            model_name = (
                config.get("CURRENT_MODEL_NAME")
                or config.get("llm_api", {}).get("current_model_name", "")
                or ""
            )
        self.model_name = str(model_name).strip()
        if keypool_url is None:
            keypool_url = (
                config.get("KEY_POOL_URL")
                or config.get("llm_api", {}).get("key_pool_url", "")
                or ""
            )
        self.keypool_url = str(keypool_url).strip().rstrip("/")
        
        if api_key is None:
            api_key = (
                config.get("KEY_POOL_API_KEY")
                or config.get("llm_api", {}).get("key_pool_api_key", "")
                or os.environ.get("KEY_POOL_API_KEY", "")
                or ""
            )
        self.api_key = str(api_key).strip()
        
        logger.info(f"LLMClient initialized for model: {self.model_name}")

    def query(self, messages: List[Dict[str, Any]], extra_params: Optional[Dict[str, Any]] = None) -> Tuple[str, Any, Dict[str, Any]]:
        """
        Send a query to the LLM via KeyPool.
        
        Args:
            messages: List of message dictionaries
            extra_params: Optional dictionary of extra parameters to override defaults
            
        Returns:
            Tuple[str, Any, Dict[str, Any]]: (response_content, reasoning_content, token_usage)
        """
        try:
            # Prepare payload for KeyPool
            payload = {
                "model": self.model_name,
                "messages": messages,
                "stream": False
            }
            
            # Add extra parameters if provided
            if extra_params:
                payload.update(extra_params)
            
            # Build headers (OpenAI-compatible /v1/chat/completions endpoint)
            headers = {"Content-Type": "application/json"}
            if self.api_key:
                headers["Authorization"] = f"Bearer {self.api_key}"
            
            # Send request to KeyPool (OpenAI-compatible endpoint)
            response = requests.post(
                f"{self.keypool_url}/v1/chat/completions",
                json=payload,
                headers=headers,
                timeout=self.timeout
            )
            
            response.raise_for_status()
            result = response.json()
            
            content = ""
            reasoning = None
            usage = {}
            
            if "choices" in result and len(result["choices"]) > 0:
                message = result["choices"][0]["message"]
                content = message.get("content", "")
                reasoning = message.get("reasoning_content")
                
            if "usage" in result:
                usage = result["usage"]
            return content, reasoning, usage
                
        except Exception as e:
            logger.error(f"Error querying KeyPool: {e}")
            logger.error(traceback.format_exc())
            return "", None, {}

    def add_history(self, role: str, content: Any):
        """Add a message to history"""
        self.history.append({"role": role, "content": content})
        if len(self.history) > MAX_HISTORY_LENGTH:
            self.history.pop(0)

    def clear_history(self):
        """Clear conversation history"""
        self.history = []