app/llm.py

from __future__ import annotations

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

from tenacity import retry, stop_after_attempt, wait_exponential
import httpx
import google.generativeai as genai
from loguru import logger

from .gemini_client import GeminiClient, GeminiResponseError
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,
)

# --- Retry decorator cho các lỗi tạm thời của LLM (network, server-side) ---
retry_on_llm_transient_error = retry(
    stop=stop_after_attempt(4),  # 1 lần gọi gốc + 3 lần thử lại
    wait=wait_exponential(multiplier=5, min=10, max=60),  # Chờ 10s, 20s, 40s
    # Sửa lỗi: Điều kiện retry phải rõ ràng và an toàn
    # Chỉ retry nếu có exception, và exception đó không phải là GeminiResponseError.
    retry=lambda retry_state: (
        # Phải có exception mới xét đến retry
        retry_state.outcome.failed
        and
        # Exception đó không được là GeminiResponseError (lỗi logic)
        not isinstance(retry_state.outcome.exception(), GeminiResponseError)
    ),
    before_sleep=lambda retry_state: logger.warning(
        f"[LLM][RETRY] LLM call failed with transient error, retrying... "
        f"Attempt: {retry_state.attempt_number}, Error: {retry_state.outcome.exception()}"
    ),
)


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()
        self.gemini_model_name = config.get(
            "model", "gemini-1.5-flash"
        )  # Lấy model name từ config

        # Khởi tạo một model riêng để đếm token, tránh ảnh hưởng đến logic xoay vòng key của GeminiClient
        self.gemini_token_counter_model = None  # Khởi tạo là None
        try:  # noqa
            settings = get_settings()
            api_keys = settings.gemini_api_keys
            if api_keys:
                # Chỉ sử dụng key đầu tiên cho việc đếm token.
                first_key = api_keys.split(",")[0].strip()
                genai.configure(api_key=first_key)
                self.gemini_token_counter_model = genai.GenerativeModel(
                    self.gemini_model_name
                )  # noqa
                logger.info(
                    f"[LLM] Initialized GeminiClient and token counter for model '{self.gemini_model_name}'."
                )
            else:
                logger.warning(
                    "[LLM] Gemini API keys not found. Token counting will be disabled."
                )
        except Exception as e:
            logger.error(f"[LLM] Failed to initialize Gemini token counter: {e}")

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

    @timing_decorator_async
    async def generate_text(
        self,
        prompt: str,
        system_prompt: Optional[str] = None,
        **kwargs,  # noqa
    ) -> str:
        """
        Tạo text từ prompt sử dụng LLM.
        """
        logger.debug(
            # 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.debug(
                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.debug(
            f"[LLM] generate_text - provider: {self.provider}\n\t response: {_safe_truncate(str(response))}"
        )
        try:
            logger.debug(
                f"[LLM][RAW_RESPONSE] {json.dumps(response, ensure_ascii=False, indent=2) if hasattr(response, 'json') else str(response)}"
            )
        except Exception:  # noqa
            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")

    @retry_on_llm_transient_error
    async def _generate_gemini(self, prompt: str, **kwargs) -> str:
        # Thêm log để đếm số token trước khi gửi
        if self.gemini_token_counter_model:
            try:
                # Sử dụng phương thức bất đồng bộ để không block event loop
                token_count_response = (
                    await self.gemini_token_counter_model.count_tokens_async(prompt)
                )  # noqa
                logger.debug(
                    f"[LLM][GEMINI_TOKENS] Prompt token count: {token_count_response.total_tokens}"
                )
            except Exception as e:
                logger.warning(
                    f"[LLM][GEMINI_TOKENS] Could not count tokens for prompt: {e}"
                )
        else:
            # Thêm log này để rõ ràng hơn khi token counter không hoạt động
            logger.debug(
                "[LLM][GEMINI_TOKENS] Token counter not available, skipping count."
            )

        loop = asyncio.get_event_loop()
        # Đảm bảo kwargs được truyền nếu GeminiClient hỗ trợ
        # Decorator sẽ xử lý retry cho các lỗi tạm thời (network, server).
        # GeminiResponseError (lỗi nội dung) sẽ được raise lên cho message_processor xử lý.
        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.debug(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.debug(
                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:
        ```json
        {{
            "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 CỤM TỪ KHÓA NGẮN GỌN** là thuật ngữ pháp lý hoặc khái niệm cốt lõi để tìm kiếm trong văn bản luật.
          - **QUY TẮC 1 (Trích xuất & Chuẩn hóa)**: Xác định các hành vi vi phạm chính và chuyển đổi chúng thành cụm từ khóa pháp lý ngắn gọn. **KHÔNG** dùng cả câu mô tả đầy đủ hành vi.
            - Tốt: "vượt đèn đỏ" -> ["không chấp hành hiệu lệnh đèn tín hiệu giao thông"]
            - Xấu: "vượt đèn đỏ" -> ["Điều khiển xe ô tô không chấp hành hiệu lệnh của đèn tín hiệu giao thông"]
          - **QUY TẮC 2 (Suy luận & Mở rộng)**: Dựa vào câu hỏi, suy luận các từ khóa liên quan.
            - Ví dụ: hỏi về "biển hạn chế tốc độ tối đa" -> suy luận thêm ["biển báo cấm", "biển báo P.127"].
            - Ví dụ: hỏi về "hiệu lực biển báo khi qua ngã tư" -> suy luận thêm ["hiệu lực của biển báo", "nơi đường giao nhau"].
          - **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 chính xác hơn.
          - **QUY TẮC 4 (CẤM)**: Danh sách `tu_khoa` **CHỈ** chứa các thuật ngữ pháp lý hoặc khái niệm. **KHÔNG** được chứa:
            - Từ ngữ đời thường (ví dụ: "vượt đèn đỏ", "say rượu").
            - Các câu hỏi hoặc cụm từ chứa ý định hỏi (ví dụ: "mức phạt bao nhiêu", "phạt tiền").
            - Các câu diễn giải dài dòng.

        - **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:

        **VÍ DỤ 1 (Xử lý ngữ cảnh):**
        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:
        ```json
        {{
            "muc_dich": "hỏi về mức phạt",
            "phuong_tien": "Xe máy",
            "tu_khoa": ["đi vào đường cao tốc", "xe máy đ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?"
        }}
        ```

        **VÍ DỤ 2 (Suy luận từ khóa):**
        Lịch sử trò chuyện:
        ""

        Câu hỏi mới nhất: "qua ngã 3, ngã 4 thì biển báo hạn chế tốc độ tối đa (nền trắng, viền đỏ) có hết hiệu lực không hay chỉ khi gặp biển báo 'Hết tốc độ tối đa cho phép' thì mới hết hiệu lực?"

        Kết quả JSON mong muốn:
        ```json
        {{
            "muc_dich": "hỏi về quy tắc giao thông",
            "phuong_tien": "",
            "tu_khoa": [
                "hiệu lực biển báo cấm",
                "biển báo P.127",
                "biển báo hết tốc độ tối đa cho phép",
                "biển báo DP.134",
                "nơi đường giao nhau"
            ],
            "cau_hoi": "Hiệu lực của biển báo hạn chế tốc độ tối đa (P.127) khi đi qua nơi đường giao nhau (ngã ba, ngã tư) như thế nào và khi nào thì hết hiệu lực?"
        }}
        ```

        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.debug(
            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())