"""
Ollama LLM API 服務模組
支援一般請求與串流模式
"""

import requests
import json
from typing import Generator, Optional
from config import Config


class LLMService:
    """Ollama API 服務封裝"""

    def __init__(self, api_url: str = None, default_model: str = None):
        self.api_url = api_url or Config.OLLAMA_API_URL
        self.default_model = default_model or Config.OLLAMA_DEFAULT_MODEL

    def get_available_models(self) -> list:
        """取得可用模型列表"""
        try:
            response = requests.get(f"{self.api_url}/v1/models", timeout=10)
            response.raise_for_status()
            models = response.json()
            return [m['id'] for m in models.get('data', [])]
        except requests.RequestException as e:
            raise LLMServiceError(f"無法取得模型列表: {e}")

    def chat(
        self,
        messages: list,
        model: str = None,
        temperature: float = 0.7,
        system_prompt: str = None
    ) -> str:
        """
        發送聊天請求 (非串流)

        Args:
            messages: 訊息列表 [{"role": "user", "content": "..."}]
            model: 模型名稱
            temperature: 溫度參數 (0-1)
            system_prompt: 系統提示詞

        Returns:
            AI 回應內容
        """
        model = model or self.default_model

        # 加入系統提示詞
        if system_prompt:
            messages = [{"role": "system", "content": system_prompt}] + messages

        payload = {
            "model": model,
            "messages": messages,
            "temperature": temperature,
            "stream": False
        }

        try:
            response = requests.post(
                f"{self.api_url}/v1/chat/completions",
                json=payload,
                timeout=60
            )
            response.raise_for_status()
            result = response.json()
            return result['choices'][0]['message']['content']
        except requests.RequestException as e:
            raise LLMServiceError(f"聊天請求失敗: {e}")

    def chat_stream(
        self,
        messages: list,
        model: str = None,
        temperature: float = 0.7,
        system_prompt: str = None
    ) -> Generator[str, None, None]:
        """
        發送聊天請求 (串流模式)

        Args:
            messages: 訊息列表
            model: 模型名稱
            temperature: 溫度參數
            system_prompt: 系統提示詞

        Yields:
            串流回應的每個片段
        """
        model = model or self.default_model

        if system_prompt:
            messages = [{"role": "system", "content": system_prompt}] + messages

        payload = {
            "model": model,
            "messages": messages,
            "temperature": temperature,
            "stream": True
        }

        try:
            response = requests.post(
                f"{self.api_url}/v1/chat/completions",
                json=payload,
                stream=True,
                timeout=120
            )
            response.raise_for_status()

            for line in response.iter_lines():
                if line:
                    if line.startswith(b"data: "):
                        data_str = line[6:].decode('utf-8')
                        if data_str.strip() != "[DONE]":
                            try:
                                data = json.loads(data_str)
                                if 'choices' in data:
                                    delta = data['choices'][0].get('delta', {})
                                    if 'content' in delta:
                                        yield delta['content']
                            except json.JSONDecodeError:
                                continue
        except requests.RequestException as e:
            raise LLMServiceError(f"串流請求失敗: {e}")

    def simple_query(self, prompt: str, system_prompt: str = None) -> str:
        """
        簡單查詢 (單一問題)

        Args:
            prompt: 使用者問題
            system_prompt: 系統提示詞

        Returns:
            AI 回應
        """
        messages = [{"role": "user", "content": prompt}]
        return self.chat(messages, system_prompt=system_prompt)


class LLMServiceError(Exception):
    """LLM 服務錯誤"""
    pass


# 全域實例
llm_service = LLMService()