"""
Ollama 本地模型提供商
"""
from typing import Optional

from loguru import logger

from .base import BaseLLMProvider, LLMResponse
from ..types import LLMConfig


class OllamaProvider(BaseLLMProvider):
    """Ollama 本地模型提供商"""

    def __init__(self, config: LLMConfig):
        self.config = config

    async def call(self, prompt: str, system_prompt: Optional[str] = None) -> LLMResponse:
        """
        调用 Ollama API

        Args:
            prompt: 用户提示
            system_prompt: 系统提示

        Returns:
            LLMResponse: 响应结果
        """
        try:
            import aiohttp

            payload = {
                "model": self.config.ollama_model,
                "prompt": prompt,
                "stream": False,
                "options": {
                    "temperature": self.config.temperature,
                    "num_predict": self.config.max_tokens,
                }
            }
            if system_prompt:
                payload["system"] = system_prompt

            logger.debug(f"调用 Ollama API: model={self.config.ollama_model}")
            async with aiohttp.ClientSession() as session:
                async with session.post(
                    f"{self.config.ollama_base_url}/api/generate",
                    json=payload,
                    timeout=aiohttp.ClientTimeout(total=self.config.timeout),
                ) as response:
                    if response.status != 200:
                        error_text = await response.text()
                        return LLMResponse(content="", success=False, error=f"Ollama API 错误: {error_text}")

                    data = await response.json()
                    return LLMResponse(
                        content=data.get("response", ""),
                        success=True,
                        tokens_used=data.get("eval_count", 0),
                    )
        except ImportError:
            return LLMResponse(content="", success=False, error="请安装 aiohttp 包: pip install aiohttp")
        except Exception as e:
            logger.error(f"Ollama API 调用失败: {e}")
            return LLMResponse(content="", success=False, error=str(e))

    async def is_available(self) -> bool:
        """检查 Ollama 服务是否可用"""
        try:
            import aiohttp
            async with aiohttp.ClientSession() as session:
                async with session.get(
                    f"{self.config.ollama_base_url}/api/tags",
                    timeout=aiohttp.ClientTimeout(total=5),
                ) as response:
                    return response.status == 200
        except Exception:
            return False
