from typing import List, Dict, Any, Optional, Union
import httpx
import json
from loguru import logger
from tenacity import retry, stop_after_attempt, wait_exponential
import os
from .gemini_client import GeminiClient
from .config import get_settings

from .utils import timing_decorator_async, timing_decorator_sync, call_endpoint_with_retry

class LLMClient:
    """
    Client để tương tác với các mô hình ngôn ngữ lớn (LLM).
    Hỗ trợ nhiều provider: OpenAI, HuggingFace, local models, etc.
    """
    
    def __init__(self, provider: str = "openai", **kwargs):
        """
        Khởi tạo LLMClient.
        
        Args:
            provider (str): Loại provider ("openai", "huggingface", "local", "custom")
            **kwargs: Các tham số cấu hình khác
        """
        self.provider = provider.lower()
        self._client = httpx.AsyncClient(timeout=60.0)
        
        # Cấu hình theo provider
        if self.provider == "openai":
            self._setup_openai(kwargs)
        elif self.provider == "huggingface":
            self._setup_huggingface(kwargs)
        elif self.provider == "local":
            self._setup_local(kwargs)
        elif self.provider == "custom":
            self._setup_custom(kwargs)
        elif self.provider == "hfs":
            self._setup_HFS(kwargs)
        elif self.provider == "gemini":
            self._setup_gemini(kwargs)
        else:
            raise ValueError(f"Unsupported provider: {provider}")
    
    def _setup_openai(self, config: Dict[str, Any]):
        """Cấu hình cho OpenAI."""
        self.api_key = config.get("api_key") or os.getenv("OPENAI_API_KEY") or ""
        self.base_url = config.get("base_url", "https://api.openai.com/v1")
        self.model = config.get("model", "gpt-3.5-turbo")
        self.max_tokens = config.get("max_tokens", 1000)
        self.temperature = config.get("temperature", 0.7)
        
        if not self.api_key:
            raise ValueError("OpenAI API key is required")
    
    def _setup_huggingface(self, config: Dict[str, Any]):
        """Cấu hình cho HuggingFace."""
        self.api_key = config.get("api_key", "")
        self.base_url = config.get("base_url", "https://api-inference.huggingface.co")
        self.model = config.get("model", "microsoft/DialoGPT-medium")
        self.max_tokens = config.get("max_tokens", 1000)
        self.temperature = config.get("temperature", 0.7)
    
    def _setup_local(self, config: Dict[str, Any]):
        """Cấu hình cho local model."""
        self.api_key = ""
        self.base_url = config.get("base_url", "http://localhost:8000")
        self.model = config.get("model", "default")
        self.max_tokens = config.get("max_tokens", 1000)
        self.temperature = config.get("temperature", 0.7)
    
    def _setup_custom(self, config: Dict[str, Any]):
        """Cấu hình cho custom provider."""
        self.api_key = config.get("api_key", "")
        self.base_url = config.get("base_url")
        self.model = config.get("model", "default")
        self.max_tokens = config.get("max_tokens", 1000)
        self.temperature = config.get("temperature", 0.7)
        if not self.base_url:
            raise ValueError("Custom provider requires base_url")

    def _setup_HFS(self, config: Dict[str, Any]):
        """Cấu hình cho custom provider."""
        self.api_key = config.get("api_key", "")
        self.base_url = config.get("base_url")
        if not self.base_url:
            raise ValueError("Custom provider requires base_url")
    
    def _setup_gemini(self, config: Dict[str, Any]):
        """Cấu hình cho Gemini."""
        self.gemini_client = GeminiClient()
    
    @timing_decorator_async
    async def generate_text(
        self, 
        prompt: str, 
        system_prompt: Optional[str] = None,
        **kwargs
    ) -> str:
        """
        Tạo text từ prompt sử dụng LLM.
        
        Args:
            prompt (str): Prompt đầu vào
            system_prompt (str, optional): System prompt
            **kwargs: Các tham số bổ sung
            
        Returns:
            str: Text được tạo ra
        """
        logger.info(f"[LLM] generate_text - provider: {self.provider} \n\t prompt: {prompt}")
        try:
            result = None
            if self.provider == "openai":
                result = await self._generate_openai(prompt, system_prompt, **kwargs)
            elif self.provider == "huggingface":
                result = await self._generate_huggingface(prompt, **kwargs)
            elif self.provider == "local":
                result = await self._generate_local(prompt, **kwargs)
            elif self.provider == "custom":
                result = await self._generate_custom(prompt, **kwargs)
            elif self.provider == "hfs":
                result = await self._generate_hfs(prompt, **kwargs)
            elif self.provider == "gemini":
                result = await self._generate_gemini(prompt, **kwargs)
            else:
                raise ValueError(f"Unsupported provider: {self.provider}")
            logger.info(f"[LLM] generate_text - provider: {self.provider}\n\t result: {result}")
            return result
        except Exception as e:
            logger.error(f"[LLM] Error generating text with {self.provider}: {e}")
            raise
    
    async def _generate_openai(self, prompt: str, system_prompt: Optional[str] = None, **kwargs) -> str:
        url = f"{self.base_url}/chat/completions"
        payload = {"model": kwargs.get("model", self.model), "messages": [{"role": "system", "content": system_prompt or ""}, {"role": "user", "content": prompt}], "max_tokens": kwargs.get("max_tokens", self.max_tokens), "temperature": kwargs.get("temperature", self.temperature), "stream": False}
        headers = {"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"}
        response = await call_endpoint_with_retry(self._client, url, payload, headers=headers)
        if response is not None:
            data = response.json()
            return data["choices"][0]["message"]["content"]
        else:
            logger.error("OpenAI API response is None")
            raise RuntimeError("OpenAI API response is None")
    
    async def _generate_huggingface(self, prompt: str, **kwargs) -> str:
        url = f"{self.base_url}/generate"
        payload = {"inputs": prompt}
        response = await call_endpoint_with_retry(self._client, url, payload)
        if response is not None:
            data = response.json()
            return data[0]["generated_text"]
        else:
            logger.error("HuggingFace API response is None")
            raise RuntimeError("HuggingFace API response is None")
    
    async def _generate_local(self, prompt: str, **kwargs) -> str:
        url = f"{self.base_url}/generate"
        payload = {"prompt": prompt}
        response = await call_endpoint_with_retry(self._client, url, payload)
        if response is not None:
            data = response.json()
            return data.get("text", "")
        else:
            logger.error("Local API response is None")
            raise RuntimeError("Local API response is None")
    
    async def _generate_custom(self, prompt: str, **kwargs) -> str:
        url = f"{self.base_url}/custom"
        payload = {"prompt": prompt}
        response = await call_endpoint_with_retry(self._client, url, payload)
        if response is not None:
            data = response.json()
            return data.get("text", "")
        else:
            logger.error("Custom API response is None")
            raise RuntimeError("Custom API response is None")
    
    async def _generate_hfs(self, prompt: str, **kwargs) -> str:
        endpoint = f"{self.base_url}/purechat"
        payload = {"prompt": prompt}
        headers = {}
        if self.api_key:
            headers["Authorization"] = f"Bearer {self.api_key}"
        response = await call_endpoint_with_retry(self._client, endpoint, payload, 3, 500, headers=headers)
        logger.info(f"[LLM] generate_text - provider: {self.provider}\n\t response: {response}")
        try:
            import json as _json
            logger.info(f"[LLM][RAW_RESPONSE] { _json.dumps(response, ensure_ascii=False, indent=2) }")
        except Exception:
            logger.info(f"[LLM][RAW_RESPONSE] {str(response)}")
        if response is not None:
            data = response.json()
            if 'response' in data:
                return data['response']
            elif 'result' in data:
                return data['result']
            elif 'data' in data and isinstance(data['data'], list):
                return data['data'][0]
            return str(data)
        else:
            logger.error("HFS API response is None")
            raise RuntimeError("HFS API response is None")

    async def _generate_gemini(self, prompt: str, **kwargs) -> str:
        import asyncio
        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(None, self.gemini_client.generate_text, prompt)

    @timing_decorator_async
    async def chat(
        self, 
        messages: List[Dict[str, str]], 
        **kwargs
    ) -> str:
        """
        Chat với LLM sử dụng conversation history.
        
        Args:
            messages (List[Dict]): List các message với format [{"role": "user", "content": "..."}]
            **kwargs: Các tham số bổ sung
            
        Returns:
            str: Response từ LLM
        """
        logger.info(f"[LLM] chat - messages: {messages}")
        if self.provider == "openai":
            return await self._chat_openai(messages, **kwargs)
        else:
            # Với các provider khác, convert messages thành prompt
            prompt = self._messages_to_prompt(messages)
            return await self.generate_text(prompt, **kwargs)
    
    async def _chat_openai(self, messages: List[Dict[str, str]], **kwargs) -> str:
        """Chat với OpenAI API."""
        payload = {
            "model": kwargs.get("model", self.model),
            "messages": messages,
            "max_tokens": kwargs.get("max_tokens", self.max_tokens),
            "temperature": kwargs.get("temperature", self.temperature),
            "stream": False
        }
        
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }
        
        response = await self._client.post(
            f"{self.base_url}/chat/completions",
            headers=headers,
            json=payload
        )
        response.raise_for_status()
        
        data = response.json()
        return data["choices"][0]["message"]["content"]
    
    def _messages_to_prompt(self, messages: List[Dict[str, str]]) -> str:
        """Convert conversation messages thành prompt string."""
        prompt = ""
        for msg in messages:
            role = msg.get("role", "user")
            content = msg.get("content", "")
            
            if role == "system":
                prompt += f"System: {content}\n\n"
            elif role == "user":
                prompt += f"User: {content}\n"
            elif role == "assistant":
                prompt += f"Assistant: {content}\n"
        
        prompt += "Assistant: "
        return prompt
    
    @timing_decorator_async
    async def classify_text(
        self, 
        text: str, 
        categories: List[str],
        **kwargs
    ) -> Dict[str, Any]:
        """
        Phân loại text vào các categories.
        
        Args:
            text (str): Text cần phân loại
            categories (List[str]): List các categories
            **kwargs: Các tham số bổ sung
            
        Returns:
            Dict: Kết quả phân loại
        """
        prompt = f"""
        Phân loại text sau vào một trong các categories: {', '.join(categories)}
        
        Text: {text}
        
        Trả về kết quả theo format JSON:
        {{
            "category": "tên_category",
            "confidence": 0.95,
            "reasoning": "lý do phân loại"
        }}
        """
        
        response = await self.generate_text(prompt, **kwargs)
        
        try:
            # Tìm JSON trong response
            import re
            json_match = re.search(r'\{.*\}', response, re.DOTALL)
            if json_match:
                result = json.loads(json_match.group())
                return result
            else:
                return {
                    "category": "unknown",
                    "confidence": 0.0,
                    "reasoning": "Không thể parse JSON response"
                }
        except json.JSONDecodeError:
            return {
                "category": "unknown", 
                "confidence": 0.0,
                "reasoning": f"JSON parse error: {response}"
            }
    
    @timing_decorator_async
    async def extract_entities(
        self, 
        text: str,
        entity_types: Optional[List[str]] = None,
        **kwargs
    ) -> List[Dict[str, Any]]:
        """
        Trích xuất entities từ text.
        
        Args:
            text (str): Text cần trích xuất
            entity_types (List[str]): Các loại entity cần tìm
            **kwargs: Các tham số bổ sung
            
        Returns:
            List[Dict]: List các entities được tìm thấy
        """
        if entity_types is None:
            entity_types = ["PERSON", "ORGANIZATION", "LOCATION", "MONEY", "DATE"]
        
        prompt = f"""
        Trích xuất các entities từ text sau. Tìm các entities thuộc types: {', '.join(entity_types)}
        
        Text: {text}
        
        Trả về kết quả theo format JSON:
        [
            {{
                "text": "tên entity",
                "type": "loại entity",
                "start": 0,
                "end": 10
            }}
        ]
        """
        
        response = await self.generate_text(prompt, **kwargs)
        
        try:
            import re
            # Log toàn bộ response dưới dạng text dễ đọc
            try:
                import json as _json
                logger.info(f"[LLM][RAW_RESPONSE] { _json.dumps(response, ensure_ascii=False, indent=2) }")
            except Exception:
                logger.info(f"[LLM][RAW_RESPONSE] {str(response)}")
            # Ưu tiên parse object JSON nếu có
            json_match_obj = re.search(r'\{[\s\S]*?\}', response)
            json_match_list = re.search(r'\[[\s\S]*?\]', response)
            if json_match_list:
                entities = json.loads(json_match_list.group())
                return entities
            elif json_match_obj:
                entity = json.loads(json_match_obj.group())
                return [entity]
            else:
                return []
        except json.JSONDecodeError:
            logger.error(f"Error parsing entities JSON: {response}")
            return []

    @timing_decorator_async
    async def analyze(
        self, 
        text: str,
        **kwargs
    ) -> List[Dict[str, Any]]:
        """
        Trích xuất entities từ text.
        
        Args:
            text (str): Text cần trích xuất
            **kwargs: Các tham số bổ sung
            
        Returns:
            List[Dict]: List các entities được tìm thấy
        """
                
        prompt = f"""
            Bạn là một AI chuyên phân tích ngữ nghĩa câu hỏi về luật giao thông. Với mỗi câu đầu vào, hãy trích xuất các thông tin sau dưới dạng JSON:

            {{
                "muc_dich": "mục đích của câu hỏi",
                "phuong_tien": "loại phương tiện giao thông (nếu có)",
                "hanh_vi_vi_pham": "hành vi vi phạm pháp luật giao thông"
            }}

            Nếu một trường không có trong câu hỏi, hãy để giá trị là chuỗi rỗng "". Ví dụ: nếu câu không đề cập phương tiện, thì "phuong_tien": ""

            Câu hỏi:
            \"{text}\"
            """.strip()
        
        response = await self.generate_text(prompt, **kwargs)
        
        logger.info(f"[LLM][RAW] Kết quả trả về từ generate_text: {response}")

        try:
            import re
            # Log toàn bộ response dưới dạng text dễ đọc
            try:
                import json as _json
                logger.info(f"[LLM][RAW_RESPONSE] { _json.dumps(response, ensure_ascii=False, indent=2) }")
            except Exception:
                logger.info(f"[LLM][RAW_RESPONSE] {str(response)}")
            # Ưu tiên parse object JSON nếu có
            json_match_obj = re.search(r'\{[\s\S]*?\}', response)
            json_match_list = re.search(r'\[[\s\S]*?\]', response)
            if json_match_list:
                entities = json.loads(json_match_list.group())
                return entities
            elif json_match_obj:
                entity = json.loads(json_match_obj.group())
                return [entity]
            else:
                return []
        except json.JSONDecodeError:
            logger.error(f"Error parsing entities JSON: {response}")
            return []
    
    async def close(self):
        """Đóng client connection."""
        await self._client.aclose()


# Factory function để tạo LLMClient dễ dàng
def create_llm_client(provider: str = "openai", **kwargs) -> LLMClient:
    """
    Factory function để tạo LLMClient.
    
    Args:
        provider (str): Loại provider
        **kwargs: Các tham số cấu hình
        
    Returns:
        LLMClient: Instance của LLMClient
    """
    return LLMClient(provider, **kwargs)


# Ví dụ sử dụng
if __name__ == "__main__":
    import asyncio
    
    async def test_llm():
        # Test với OpenAI
        settings = get_settings()
        llm_client = create_llm_client(
            provider=settings.llm_provider,
            model=settings.llm_model,
            # ... các config khác nếu cần ...
        )
        
        # Generate text
        response = await llm_client.generate_text("Xin chào, bạn có khỏe không?")
        print(f"Response: {response}")
        
        # Chat
        messages = [
            {"role": "user", "content": "Bạn có thể giúp tôi không?"}
        ]
        chat_response = await llm_client.chat(messages)
        print(f"Chat response: {chat_response}")
        
        await llm_client.close()
    
    asyncio.run(test_llm())