"""Ollama Embedding Client

支持单条/批量 embed，连接池复用，可用性检测，优雅降级。
"""

from __future__ import annotations

import logging
import time

import httpx

logger = logging.getLogger(__name__)


class OllamaEmbedder:
    """Ollama embedding 客户端"""

    # 失败后多久重新检测（秒）
    _RETRY_INTERVAL = 60

    def __init__(self, model: str, ollama_url: str, dim: int = 768):
        self._model = model
        self._ollama_url = ollama_url.rstrip("/")
        self._dim = dim
        self._client: httpx.AsyncClient | None = None
        self._available: bool | None = None  # None = 未检测
        self._last_check: float = 0.0  # 上次检测时间戳

    async def _get_client(self) -> httpx.AsyncClient:
        if self._client is None:
            self._client = httpx.AsyncClient(
                base_url=self._ollama_url,
                timeout=30.0,
                limits=httpx.Limits(max_connections=10, max_keepalive_connections=5),
            )
        return self._client

    async def embed_one(self, text: str) -> list[float]:
        """单条文本 → 向量，查询时使用

        Raises:
            httpx.HTTPError: Ollama 请求失败
        """
        client = await self._get_client()
        try:
            resp = await client.post(
                "/api/embed", json={"model": self._model, "input": text}
            )
            resp.raise_for_status()
            embeddings = resp.json()["embeddings"]
            return embeddings[0]
        except (httpx.HTTPError, httpx.ConnectError, httpx.TimeoutException) as e:
            logger.warning("embed_one 请求失败: %s", e)
            raise
        except (KeyError, IndexError) as e:
            logger.warning("embed_one 响应格式错误: %s", e)
            raise

    async def embed_batch(
        self, texts: list[str], batch_size: int = 32
    ) -> list[list[float]]:
        """批量文本 → 向量列表，索引时使用

        Ollama /api/embed 的 input 支持数组，单次请求批量 embed。
        超过 batch_size 则分批调用。
        """
        if not texts:
            return []

        client = await self._get_client()
        all_embeddings: list[list[float]] = []

        for i in range(0, len(texts), batch_size):
            batch = texts[i : i + batch_size]
            try:
                resp = await client.post(
                    "/api/embed", json={"model": self._model, "input": batch}
                )
                resp.raise_for_status()
                batch_embeddings = resp.json()["embeddings"]
                all_embeddings.extend(batch_embeddings)
            except (httpx.HTTPError, httpx.ConnectError, httpx.TimeoutException) as e:
                logger.warning("embed_batch 第 %d 批请求失败: %s", i // batch_size, e)
                raise
            except (KeyError, IndexError) as e:
                logger.warning("embed_batch 第 %d 批响应格式错误: %s", i // batch_size, e)
                raise

        return all_embeddings

    async def check_availability(self) -> bool:
        """检测 Ollama 是否可用 + 模型是否已加载

        失败后 _RETRY_INTERVAL 秒内不重复检测，避免每次搜索都等连接超时。
        """
        now = time.monotonic()
        # 上次检测失败且未过冷却期 → 直接返回缓存结果
        if self._available is False and (now - self._last_check) < self._RETRY_INTERVAL:
            return False

        self._last_check = now
        client = await self._get_client()
        try:
            resp = await client.get("/api/tags", timeout=5.0)
            resp.raise_for_status()
            models = resp.json().get("models", [])
            model_names = [m.get("name", "") for m in models]
            # 模型名可能带 :latest 后缀
            self._available = any(
                self._model in name or name.startswith(self._model)
                for name in model_names
            )
            if not self._available:
                logger.info(
                    "Ollama 可用但模型 %s 未加载（已有: %s）",
                    self._model,
                    ", ".join(model_names) or "无",
                )
            return self._available
        except (httpx.ConnectError, httpx.TimeoutException) as e:
            logger.info("Ollama 不可用: %s", e)
            self._available = False
            return False
        except Exception as e:
            logger.warning("Ollama 检测异常: %s", e)
            self._available = False
            return False

    @property
    def available(self) -> bool:
        """是否可用（需先调用 check_availability）"""
        return self._available is True

    @property
    def model(self) -> str:
        return self._model

    @property
    def dim(self) -> int:
        return self._dim

    async def close(self) -> None:
        if self._client:
            await self._client.aclose()
            self._client = None
