app/llm.py

from __future__ import annotations

from typing import List, Dict, Any, Optional, Union
import json
import re
import os
import asyncio

import httpx
from loguru import logger

from .gemini_client import GeminiClient
from .config import get_settings
from .utils import (
    timing_decorator_async,
    timing_decorator_sync,  # kept for compatibility even if unused here
    call_endpoint_with_retry,
    _safe_truncate
)

def _parse_json_from_text(text: str) -> Optional[Union[List[Dict[str, Any]], Dict[str, Any]]]:
    """Best-effort JSON extractor from LLM free-form responses.

    Strategy:
    1) Try json.loads() on the whole string first.
    2) Fallback to regex to find the first JSON list/object snippet.
    """
    if not text:
        return None

    # 1) try direct load
    try:
        return json.loads(text)
    except Exception:
        pass

    # 2) find first JSON array or object
    match = re.search(r"(\[[\s\S]+?\]|\{[\s\S]+?\})", text)
    if match:
        try:
            return json.loads(match.group(1))
        except Exception:
            return None
    return None


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, custom, HFS, Gemini.
    """

    def __init__(self, provider: str = "openai", **kwargs):
        self.provider = provider.lower()
        self._client = httpx.AsyncClient(timeout=60.0)

        # Dispatch provider setup
        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}")

    # ---------- Provider setups ---------- #

    def _setup_openai(self, config: Dict[str, Any]):
        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]):
        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]):
        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]):
        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]):
        self.api_key = config.get("api_key", "")
        self.base_url = config.get("base_url")
        if not self.base_url:
            raise ValueError("HFS provider requires base_url")

    def _setup_gemini(self, config: Dict[str, Any]):
        # Sử dụng GeminiClient với RequestLimitManager (theo thiết kế của bạn)
        self.gemini_client = GeminiClient()
        logger.info("[LLM] Initialized GeminiClient with RequestLimitManager")

    # ---------- Core APIs ---------- #

    @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.
        """
        logger.info(
            f"[LLM] generate_text - provider: {self.provider}\n\t prompt: {_safe_truncate(prompt)}"
            # f"[LLM] generate_text - provider: {self.provider}\n\t prompt: {prompt}"
        )
        try:
            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: {_safe_truncate(result)}")
            return result
        except Exception as e:
            logger.exception(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"]
        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"]
        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", "")
        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", "")
        logger.error("Custom API response is None")
        raise RuntimeError("Custom API response is None")

    async def _generate_hfs(self, prompt: str, **kwargs) -> str:
        # Giữ nguyên chữ ký call_endpoint_with_retry như bạn đã dùng
        endpoint = f"{self.base_url}/purechat"
        payload = {"prompt": prompt}
        headers = {}
        if hasattr(self, "api_key") and 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: {_safe_truncate(str(response))}"
        )
        try:
            logger.info(
                f"[LLM][RAW_RESPONSE] {json.dumps(response, ensure_ascii=False, indent=2) if hasattr(response, 'json') else str(response)}"
            )
        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"]
            if "result" in data:
                return data["result"]
            if "data" in data and isinstance(data["data"], list) and data["data"]:
                return data["data"][0]
            return str(data)
        logger.error("HFS API response is None")
        raise RuntimeError("HFS API response is None")

    async def _generate_gemini(self, prompt: str, **kwargs) -> str:
        loop = asyncio.get_event_loop()
        # Đảm bảo kwargs được truyền nếu GeminiClient hỗ trợ
        return await loop.run_in_executor(None, lambda: self.gemini_client.generate_text(prompt, **kwargs))

    @timing_decorator_async
    async def chat(self, messages: List[Dict[str, str]], **kwargs) -> str:
        """Chat với LLM sử dụng conversation history."""
        logger.info(f"[LLM] chat - provider: {self.provider} - messages: {messages}")
        if self.provider == "openai":
            return await self._chat_openai(messages, **kwargs)
        # Convert messages -> prompt cho các provider khác
        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:
        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:
        prompt_lines = []
        for msg in messages:
            role = msg.get("role", "user")
            content = msg.get("content", "")
            if role == "system":
                prompt_lines.append(f"System: {content}\n")
            elif role == "user":
                prompt_lines.append(f"User: {content}\n")
            elif role == "assistant":
                prompt_lines.append(f"Assistant: {content}\n")
        prompt_lines.append("Assistant: ")
        return "".join(prompt_lines)

    # ---------- Utility tasks ---------- #

    @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."""
        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)

        result = _parse_json_from_text(response or "")
        if isinstance(result, dict):
            return result
        # fallback default
        return {
            "category": "unknown",
            "confidence": 0.0,
            "reasoning": f"Cannot parse JSON from response: {_safe_truncate(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."""
        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:
            logger.info(
                f"[LLM][RAW_RESPONSE][extract_entities] {_safe_truncate(response)}"
            )
            parsed = _parse_json_from_text(response or "")
            if isinstance(parsed, list):
                return parsed
            if isinstance(parsed, dict):
                return [parsed]
            return []
        except Exception as e:
            logger.error(f"Error parsing entities JSON: {e} | Raw: {response}")
            return []

    @timing_decorator_async
    async def analyze(
        self, text: str, conversation_context: str, **kwargs
    ) -> List[Dict[str, Any]]:
        """
        Phân tích câu hỏi về luật giao thông Việt Nam và chuẩn hóa thành JSON.
        """
        prompt = f"""
        Bạn là một chuyên gia phân tích ngôn ngữ tự nhiên (NLP) chuyên xử lý các câu hỏi về luật giao thông Việt Nam. Nhiệm vụ của bạn là đọc kỹ **lịch sử trò chuyện** và **câu hỏi mới nhất** của người dùng để trích xuất thông tin vào một cấu trúc JSON duy nhất. **Luôn chỉ trả về đối tượng JSON hợp lệ**, không thêm bất kỳ giải thích nào.

        Định dạng JSON bắt buộc:
        {{
            "muc_dich": "...",
            "phuong_tien": "...",
            "tu_khoa": [],
            "cau_hoi": "..."
        }}

        Hướng dẫn chi tiết cho từng trường:

        **muc_dich**: Phải là một trong các giá trị sau, dựa vào **câu hỏi mới nhất**:
        - "hỏi về mức phạt"
        - "hỏi về quy tắc giao thông"
        - "hỏi về báo hiệu đường bộ"
        - "hỏi về quy trình xử lý vi phạm giao thông"
        - "thông tin cá nhân của AI"
        - "khác"

        **phuong_tien**: Tên phương tiện được đề cập trong câu hỏi mới hoặc trong lịch sử gần nhất. Nếu không có, để chuỗi rỗng "".

        **tu_khoa**: **MỘT DANH SÁCH (LIST) các thuật ngữ pháp lý** ngắn gọn, chính xác nhất để tìm kiếm trong cơ sở dữ liệu luật.
        - **Quy tắc 1 (Chuyển đổi ngôn ngữ)**: Chuyển đổi ngôn ngữ đời thường của người dùng (ví dụ: "vượt đèn đỏ") thành thuật ngữ pháp lý chính xác (ví dụ: "Không chấp hành hiệu lệnh của đèn tín hiệu giao thông").
        - **Quy tắc 2 (Trích xuất nhiều từ khóa)**: Nếu câu hỏi phức tạp chứa nhiều hành vi, chi tiết, thông tin, hãy trích xuất nhiều từ khóa liên quan. Ví dụ: "vượt đèn đỏ khi đang say rượu" -> ["Không chấp hành hiệu lệnh của đèn tín hiệu giao thông", "Điều khiển xe trên đường mà trong máu hoặc hơi thở có nồng độ cồn"].
        - **Quy tắc 3 (Xử lý ngữ cảnh không hài lòng)**: Đọc kỹ lịch sử. Nếu người dùng hỏi lại hoặc thể hiện không hài lòng (ví dụ: "không phải", "ý tôi là..."), , và trong lịch sử có ghi chú (từ khóa đã dùng: ...) thì TUYỆT ĐỐI KHÔNG SỬ DỤNG LẠI các từ khóa đó. Hãy tạo ra một bộ từ khóa **HOÀN TOÀN MỚI** để tìm kiếm thông tin chính xác hơn.

        **cau_hoi**: Diễn đạt lại câu hỏi mới nhất của người dùng thành một câu hỏi hoàn chỉnh, kết hợp ngữ cảnh từ lịch sử nếu cần, sử dụng đúng thuật ngữ pháp lý.

        VÍ DỤ MẪU:

        Lịch sử trò chuyện:
        "Người dùng: xe máy đi vào đường cấm thì sao? (từ khóa đã dùng: đi vào khu vực cấm)
        Trợ lý: Mức phạt cho hành vi đi vào khu vực cấm là..."

        Câu hỏi mới nhất: "không phải, ý tôi là đi vào đường cao tốc cơ"

        Kết quả JSON mong muốn:
        {{
            "muc_dich": "hỏi về mức phạt",
            "phuong_tien": "Xe máy",
            "tu_khoa": ["Điều khiển xe đi vào đường cao tốc"],
            "cau_hoi": "Mức xử phạt cho hành vi xe máy đi vào đường cao tốc là bao nhiêu?"
        }}

        Bây giờ, hãy phân tích lịch sử và câu hỏi sau và chỉ trả về đối tượng JSON.

        Lịch sử trò chuyện:
        "{conversation_context}"

        Câu hỏi mới nhất:
        "{text}"
        """.strip()

        response = await self.generate_text(prompt, **kwargs)
        logger.info(f"[LLM][RAW][analyze] Kết quả trả về từ generate_text: {_safe_truncate(response)}")
        # logger.info(f"[LLM][RAW][analyze] Kết quả trả về từ generate_text: {response}")

        try:
            parsed = _parse_json_from_text(response or "")
            if isinstance(parsed, list):
                return parsed
            if isinstance(parsed, dict):
                return [parsed]
            return []
        except Exception as e:
            logger.error(f"Error parsing analyze JSON: {e} | Raw: {response}")
            return []

    async def close(self):
        """Đóng client connection."""
        if hasattr(self, "_client") and self._client:
            await self._client.aclose()


# Factory function để tạo LLMClient dễ dàng
def create_llm_client(provider: str = "openai", **kwargs) -> LLMClient:
    return LLMClient(provider, **kwargs)


# Ví dụ sử dụng
if __name__ == "__main__":
    async def test_llm():
        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())