"""Ollama Embedding Client

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

from __future__ import annotations

import asyncio
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._client_lock = asyncio.Lock()
        self._available: bool | None = None  # None = 未检测
        self._last_check: float = 0.0  # 上次检测时间戳
        self._unavailable_reason: str = ""  # 不可用的具体原因

    # 重试配置
    _MAX_RETRIES = 3
    _BACKOFF_BASE = 1.0  # 退避基数（秒）：1, 2, 4

    async def _get_client(self) -> httpx.AsyncClient:
        if self._client is None:
            self._client = httpx.AsyncClient(
                base_url=self._ollama_url,
                timeout=httpx.Timeout(connect=3.0, read=30.0, write=10.0, pool=10.0),
                limits=httpx.Limits(max_connections=10, max_keepalive_connections=5),
                proxy=None,  # 本地 Ollama 请求不走代理
            )
        return self._client

    async def _rebuild_client(self) -> httpx.AsyncClient:
        """关闭旧连接并重建 client（连接池损坏时调用）"""
        async with self._client_lock:
            if self._client:
                try:
                    await self._client.aclose()
                except Exception as e:
                    logger.warning("关闭旧 client 失败: %s", e)
                self._client = None
            return await self._get_client()

    async def embed_one(self, text: str) -> list[float]:
        """单条文本 → 向量，查询时使用。带指数退避重试。

        Raises:
            httpx.HTTPError: 重试耗尽后仍然失败
        """
        last_exc: Exception | None = None
        for attempt in range(self._MAX_RETRIES):
            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.ConnectError, httpx.TimeoutException) as e:
                last_exc = e
                wait = self._BACKOFF_BASE * (2 ** attempt)
                logger.warning(
                    "embed_one 请求失败（第 %d/%d 次）: %s，%.1fs 后重试",
                    attempt + 1, self._MAX_RETRIES, e, wait,
                )
                await self._rebuild_client()
                await asyncio.sleep(wait)
            except httpx.HTTPStatusError as e:
                logger.warning("embed_one HTTP 状态错误: %s", e)
                raise
            except (KeyError, IndexError) as e:
                logger.warning("embed_one 响应格式错误: %s", e)
                raise
        logger.warning("embed_one 重试耗尽（%d 次）", self._MAX_RETRIES)
        if last_exc is not None:
            raise last_exc
        raise RuntimeError("embed_one 失败但未捕获到异常")

    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 []

        all_embeddings: list[list[float]] = []

        for i in range(0, len(texts), batch_size):
            batch = texts[i : i + batch_size]
            batch_num = i // batch_size

            last_exc: Exception | None = None
            for attempt in range(self._MAX_RETRIES):
                client = await self._get_client()
                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)
                    last_exc = None
                    break
                except (httpx.ConnectError, httpx.TimeoutException) as e:
                    last_exc = e
                    wait = self._BACKOFF_BASE * (2 ** attempt)
                    logger.warning(
                        "embed_batch 第 %d 批失败（第 %d/%d 次）: %s，%.1fs 后重试",
                        batch_num, attempt + 1, self._MAX_RETRIES, e, wait,
                    )
                    await self._rebuild_client()
                    await asyncio.sleep(wait)
                except httpx.HTTPStatusError as e:
                    logger.warning("embed_batch 第 %d 批 HTTP 状态错误: %s", batch_num, e)
                    raise
                except (KeyError, IndexError) as e:
                    logger.warning("embed_batch 第 %d 批响应格式错误: %s", batch_num, e)
                    raise
            else:
                # for 循环未 break（重试全部耗尽）
                if last_exc is not None:
                    logger.warning("embed_batch 第 %d 批重试耗尽（%d 次）", batch_num, self._MAX_RETRIES)
                    raise last_exc

        return all_embeddings

    async def check_availability(self) -> bool:
        """检测 Ollama 是否可用 + 模型是否真正能 embed

        两步检测：
        1. /api/tags 检查模型是否已下载
        2. 实际 embed 一段短文本，确认模型能加载到内存并返回结果

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

        self._last_check = now
        self._unavailable_reason = ""
        client = await self._get_client()

        # Step 1: 检查 Ollama 服务 + 模型列表
        try:
            resp = await client.get("/api/tags", timeout=3.0)
            resp.raise_for_status()
            models = resp.json().get("models", [])
            model_names = [m.get("name", "") for m in models]
            model_found = any(
                self._model in name or name.startswith(self._model)
                for name in model_names
            )
            if not model_found:
                self._available = False
                self._unavailable_reason = (
                    f"模型 {self._model} 未下载，"
                    f"请运行: ollama pull {self._model}"
                )
                logger.warning(
                    "Ollama 可用但模型 %s 未下载（已有: %s）",
                    self._model,
                    ", ".join(model_names) or "无",
                )
                return False
        except (httpx.ConnectError, httpx.TimeoutException) as e:
            self._available = False
            self._unavailable_reason = (
                "Ollama 未启动，请运行: ollama serve"
            )
            logger.warning("Ollama 不可用: %s", e)
            return False
        except Exception as e:
            self._available = False
            self._unavailable_reason = f"Ollama 检测异常: {e}"
            logger.warning("Ollama 检测异常: %s", e)
            return False

        # Step 2: 实际 embed 短文本，确认模型能工作
        try:
            resp = await client.post(
                "/api/embed",
                json={"model": self._model, "input": "ping"},
                timeout=15.0,
            )
            resp.raise_for_status()
            embeddings = resp.json().get("embeddings", [])
            if not embeddings or not embeddings[0]:
                self._available = False
                self._unavailable_reason = (
                    f"模型 {self._model} 返回空向量，可能已损坏，"
                    f"请运行: ollama rm {self._model} && ollama pull {self._model}"
                )
                return False
        except (httpx.TimeoutException,) as e:
            self._available = False
            self._unavailable_reason = (
                f"模型 {self._model} 加载超时（15s），"
                "可能系统内存不足或模型过大"
            )
            logger.warning("Ollama embed 测试超时: %s", e)
            return False
        except Exception as e:
            self._available = False
            self._unavailable_reason = f"Ollama embed 测试失败: {e}"
            logger.warning("Ollama embed 测试失败: %s", e)
            return False

        self._available = True
        self._unavailable_reason = ""
        logger.info("Ollama 检测通过: 模型 %s 可用", self._model)
        return True

    def reset_cooldown(self) -> None:
        """重置冷却期，允许立即重新检测"""
        self._last_check = 0.0
        self._available = None

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

    @property
    def unavailable_reason(self) -> str:
        """不可用的具体原因（供错误提示使用）"""
        return self._unavailable_reason

    @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
