import requests


class OllamaClient:
    def __init__(self, api_url, api_key=None, timeout=30):
        """
        初始化 OllamaClient。

        :param api_url: API 基本端點，例如 'http://192.168.88.82:11434'
        :param api_key: 如果需要身份驗證，填入 API 金鑰；否則為 None。
        :param timeout: 請求超時時間（秒）。
        """
        self.api_url = api_url.rstrip("/")  # 移除末尾的斜杠以避免拼接錯誤
        self.api_key = api_key
        self.timeout = timeout
        self.headers = {"Content-Type": "application/json"}
        if self.api_key:
            self.headers["Authorization"] = f"Bearer {self.api_key}"

    def process_text(self, text, model="qwen2.5:14b", additional_prompt=None):
        """
        發送文本到指定模型，並獲取處理後的結果。

        :param text: 要處理的文本。
        :param model: 指定要使用的模型（默認為 'qwen2.5:14b'）。
        :param additional_prompt: 可選的提示詞，將添加在文本之前。
        :return: 處理後的文本。
        """
        # 合成最終的 prompt
        final_prompt = f"{additional_prompt}\n{text}" if additional_prompt else text

        payload = {
            "model": model,
            "prompt": final_prompt,
            "stream": False,  # 確保回應為全量模式
        }
        try:
            response = requests.post(
                f"{self.api_url}/api/generate",
                json=payload,
                headers=self.headers,
                timeout=self.timeout,
            )
            response.raise_for_status()

            # 確保回應是 JSON 格式
            try:
                data = response.json()
            except ValueError as e:
                print(f"無法解析 Ollama LLM 的回應，非 JSON 格式: {e}")
                print(f"回應內容: {response.text}")
                return text

            # 提取回應中的處理結果
            processed_text = data.get("response")
            if not processed_text:
                print("Ollama LLM 回應中缺少 'response' 欄位。")
                return text
            return processed_text
        except requests.exceptions.RequestException as e:
            print(f"Ollama LLM 請求錯誤: {e}")
            return text

    def get_models(self):
        """
        獲取當前可用的模型完整資訊。

        :return: 包含模型完整資訊的列表（如果失敗，返回空列表）。
        """
        try:
            response = requests.get(
                f"{self.api_url}/api/tags", headers=self.headers, timeout=self.timeout
            )
            response.raise_for_status()
            data = response.json()
            models = data.get("models", [])
            if not models:
                print("未能從 Ollama LLM 獲取模型列表。")
                return []
            return models
        except requests.exceptions.RequestException as e:
            print(f"Ollama LLM 請求錯誤: {e}")
            return []
        except ValueError:
            print("無法解析 Ollama LLM 的回應。")
            return []