# -*- coding: utf-8 -*-
"""
LLM Client — tương tác với vLLM server
-------------------------------------
- Dùng HTTP endpoint (OpenAI-compatible) của vLLM.
- Nhận prompt đã được build từ PromptBuilder.
- Gọi model sinh câu trả lời.
"""

import requests
from typing import Optional, Dict, Any
from src.core.config import VLLM_URL, VLLM_MODEL, LLM_TEMPERATURE, LLM_MAX_TOKENS


class LLMClient:
    """
    Gọi model vLLM theo giao thức OpenAI-compatible.
    """
    def __init__(
        self,
        base_url: str = VLLM_URL,
        model: str = VLLM_MODEL,
        temperature: float = LLM_TEMPERATURE,
        max_tokens: int = LLM_MAX_TOKENS
    ):
        self.base_url = base_url.rstrip("/")
        self.model = model
        self.temperature = temperature
        self.max_tokens = max_tokens

    def generate(self, prompt: str, system_prompt: Optional[str] = None) -> str:
        """
        Gửi prompt tới vLLM API và trả về câu trả lời.
        """
        url = f"{self.base_url}/v1/chat/completions"

        messages = []
        if system_prompt:
            messages.append({"role": "system", "content": system_prompt})
        messages.append({"role": "user", "content": prompt})

        payload: Dict[str, Any] = {
            "model": self.model,
            "messages": messages,
            "temperature": self.temperature,
            "max_tokens": self.max_tokens,
            "stream": False,
        }

        try:
            response = requests.post(url, json=payload, timeout=60)
            response.raise_for_status()
            data = response.json()
            return data["choices"][0]["message"]["content"].strip()
        except Exception as e:
            print(f"[LLMClient] ❌ Error calling vLLM: {e}")
            return f"[Error] {e}"


# ---- Test nhanh ----
if __name__ == "__main__":
    client = LLMClient()
    prompt = "Viết đoạn mô tả ngắn về hành tinh Sao Hỏa."
    answer = client.generate(prompt)
    print("🪐 Kết quả từ LLM:")
    print(answer)