"""Embedding Client

支持 Ollama 和 OpenAI 兼容 API（llama.cpp server / LM Studio / vLLM）。
单条/批量 embed，连接池复用，可用性检测，优雅降级，指数退避重试。

Embedder Protocol 定义通用接口，create_embedder() 工厂函数按配置选择后端。
"""

from __future__ import annotations

import asyncio
import logging
from typing import Protocol, runtime_checkable

import httpx

from .resilience import ServiceCooldown

logger = logging.getLogger(__name__)


@runtime_checkable
class Embedder(Protocol):
    """Embedding 后端协议 — 所有 embedder 实现此接口

    OllamaEmbedder（远程 HTTP）和 FastEmbedEmbedder（本地 ONNX）均遵循此协议。
    """

    @property
    def available(self) -> bool: ...

    @property
    def unavailable_reason(self) -> str: ...

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

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

    async def embed_one(self, text: str) -> list[float]: ...

    async def embed_batch(
        self, texts: list[str], batch_size: int = 32
    ) -> list[list[float] | None]: ...

    async def check_availability(self) -> bool: ...

    async def close(self) -> None: ...


def create_embedder(config) -> Embedder:
    """工厂函数 — 根据 config.embed_backend 创建 embedder 实例

    - "onnx": FastEmbedEmbedder（本地 ONNX，~5ms/query，无需外部服务）
    - "ollama": OllamaEmbedder（Ollama HTTP API，需要 ollama serve）
    """
    backend = getattr(config, "embed_backend", "ollama")

    if backend == "onnx":
        try:
            from .embed_onnx import FastEmbedEmbedder
            return FastEmbedEmbedder(
                model=getattr(config, "onnx_model", "BAAI/bge-small-en-v1.5"),
                dim=getattr(config, "onnx_dim", 384),
            )
        except ImportError:
            logger.warning(
                "fastembed 未安装，回退到 Ollama 后端。"
                "安装: pip install fastembed"
            )
            # 降级到 Ollama

    return OllamaEmbedder(
        model=config.embedding_model,
        ollama_url=config.ollama_url,
        dim=config.embedding_dim,
        api=config.embed_api,
    )


class OllamaEmbedder:
    """Embedding 客户端，支持 Ollama 和 OpenAI 兼容 API"""

    def __init__(
        self,
        model: str,
        ollama_url: str,
        dim: int = 768,
        api: str = "ollama",
    ):
        self._model = model
        self._ollama_url = ollama_url.rstrip("/")
        self._dim = dim
        self._api = api  # "ollama" | "openai_compat"
        self._client: httpx.AsyncClient | None = None
        self._client_lock = asyncio.Lock()
        self._available: bool | None = None  # None = 未检测
        self._unavailable_reason: str = ""  # 不可用的具体原因
        self._cooldown = ServiceCooldown(
            f"ollama_embed_{model}", base_cooldown=60.0, max_cooldown=300.0,
        )

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

    # 并发配置
    _CONCURRENT_BATCHES = 4  # embed_batch 并发批次数

    # 连接配置（提取为类常量，避免魔法数字）
    _DEFAULT_TIMEOUT = httpx.Timeout(connect=3.0, read=30.0, write=10.0, pool=10.0)
    _DEFAULT_LIMITS = httpx.Limits(max_connections=10, max_keepalive_connections=5)

    def _create_client(self) -> httpx.AsyncClient:
        """创建新的 httpx 客户端（无锁，仅供内部调用）"""
        # httpx 0.28+ 中 proxy=None 会从环境变量读取代理（ALL_PROXY 等），
        # 导致本地 Ollama 请求走 SOCKS 代理而失败。
        # 显式提供无代理 transport 彻底绕过环境变量检测。
        return httpx.AsyncClient(
            base_url=self._ollama_url,
            timeout=self._DEFAULT_TIMEOUT,
            transport=httpx.AsyncHTTPTransport(limits=self._DEFAULT_LIMITS),
        )

    async def _get_client(self) -> httpx.AsyncClient:
        if self._client is not None:  # 快速路径，无锁
            return self._client
        async with self._client_lock:
            if self._client is None:  # 双重检查
                self._client = self._create_client()
            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 = self._create_client()
            return self._client

    # ── API 适配层 ──

    def _embed_endpoint(self) -> str:
        """返回 embed 请求的路径"""
        if self._api == "openai_compat":
            return "/v1/embeddings"
        return "/api/embed"

    def _embed_payload(self, input_data: str | list[str]) -> dict:
        """构建 embed 请求体（Ollama 和 OpenAI 格式相同）"""
        return {"model": self._model, "input": input_data}

    def _parse_embeddings(self, data: dict) -> list[list[float]]:
        """从响应中提取 embedding 向量列表"""
        if self._api == "openai_compat":
            # OpenAI 格式: {"data": [{"embedding": [...], "index": 0}, ...]}
            items = data.get("data")
            if not items or not isinstance(items, list):
                raise ValueError(f"OpenAI 响应缺少有效的 data 数组: {str(data)[:200]}")
            for i, item in enumerate(items):
                if not isinstance(item, dict) or "embedding" not in item:
                    raise ValueError(f"data[{i}] 缺少 embedding 字段")
            # 按 index 排序确保顺序正确（部分服务可能不返回 index）
            if all("index" in item for item in items):
                items.sort(key=lambda x: x["index"])
            return [item["embedding"] for item in items]
        # Ollama 格式: {"embeddings": [[...], [...]]}
        return data["embeddings"]

    # ── 核心方法 ──

    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(
                    self._embed_endpoint(),
                    json=self._embed_payload(text),
                )
                resp.raise_for_status()
                embeddings = self._parse_embeddings(resp.json())
                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 None:
            raise RuntimeError("重试循环结束但无异常，逻辑错误")
        raise last_exc

    async def _embed_single_batch(
        self, batch: list[str], batch_num: int
    ) -> list[list[float]]:
        """单批次嵌入，带指数退避重试（内部方法）"""
        last_exc: Exception | None = None
        for attempt in range(self._MAX_RETRIES):
            client = await self._get_client()
            try:
                resp = await client.post(
                    self._embed_endpoint(),
                    json=self._embed_payload(batch),
                )
                resp.raise_for_status()
                return self._parse_embeddings(resp.json())
            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
        if last_exc is not None:
            logger.warning("embed_batch 第 %d 批重试耗尽（%d 次）", batch_num, self._MAX_RETRIES)
            raise last_exc
        raise RuntimeError("重试循环结束但无异常，逻辑错误")

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

        Ollama /api/embed 的 input 支持数组，单次请求批量 embed。
        OpenAI /v1/embeddings 同样支持数组 input。
        超过 batch_size 则分批，多批并发发送（受 _CONCURRENT_BATCHES 限制）。

        允许部分成功：失败批次对应位置填 None，调用方自行处理断层。
        chunk 没有 embedding 仍可入库（BM25 搜索可用，向量搜索缺失）。
        """
        if not texts:
            return []

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

        # 单批次直接调用
        if len(batches) == 1:
            try:
                return await self._embed_single_batch(batches[0], 0)
            except Exception as e:
                logger.warning("embed_batch 单批失败: %s，返回 None 填充", e)
                return [None] * len(batches[0])

        # 多批次并发，允许部分失败
        sem = asyncio.Semaphore(self._CONCURRENT_BATCHES)

        async def _limited(batch: list[str], num: int) -> list[list[float]]:
            async with sem:
                return await self._embed_single_batch(batch, num)

        results = await asyncio.gather(
            *[_limited(b, i) for i, b in enumerate(batches)],
            return_exceptions=True,
        )

        all_embeddings: list[list[float] | None] = []
        failed_count = 0
        for i, batch_result in enumerate(results):
            if isinstance(batch_result, Exception):
                logger.warning("embed_batch 第 %d 批失败: %s，该批向量置空", i, batch_result)
                all_embeddings.extend([None] * len(batches[i]))
                failed_count += 1
            else:
                all_embeddings.extend(batch_result)

        if failed_count:
            logger.warning(
                "embed_batch 部分成功: %d/%d 批失败，%d/%d 条文本无向量",
                failed_count, len(batches),
                sum(1 for e in all_embeddings if e is None), len(texts),
            )

        return all_embeddings

    async def check_availability(self) -> bool:
        """检测 embedding 服务是否可用

        Ollama 模式：两步检测（/api/tags + 实际 embed）
        OpenAI 兼容模式：直接 embed 短文本测试

        失败后通过 ServiceCooldown 指数退避冷却，避免每次搜索都等连接超时。
        """
        # 冷却中直接返回（指数退避，避免频繁检测）
        if self._cooldown.is_cooling:
            return False

        # 锁外执行网络请求（避免阻塞其他操作）
        client = await self._get_client()
        result = False
        reason = ""

        if self._api == "openai_compat":
            return await self._check_openai_compat(client)

        # ── Ollama 模式 ──

        # 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(
                name == self._model or name.startswith(f"{self._model}:")
                for name in model_names
            )
            if not model_found:
                reason = (
                    f"模型 {self._model} 未下载，"
                    f"请运行: ollama pull {self._model}"
                )
                logger.warning(
                    "Ollama 可用但模型 %s 未下载（已有: %s）",
                    self._model, ", ".join(model_names) or "无",
                )
                self._set_unavailable(reason)
                return False
        except (httpx.ConnectError, httpx.TimeoutException) as e:
            reason = "Ollama 未启动，请运行: ollama serve"
            logger.warning("Ollama 不可用: %s", e)
            self._set_unavailable(reason)
            return False
        except httpx.HTTPError as e:
            reason = f"Ollama 检测异常: {e}"
            logger.warning("Ollama 检测异常: %s", e)
            self._set_unavailable(reason)
            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]:
                reason = (
                    f"模型 {self._model} 返回空向量，可能已损坏，"
                    f"请运行: ollama rm {self._model} && ollama pull {self._model}"
                )
                self._set_unavailable(reason)
                return False
        except httpx.TimeoutException as e:
            reason = (
                f"模型 {self._model} 加载超时（15s），"
                "可能系统内存不足或模型过大"
            )
            logger.warning("Ollama embed 测试超时: %s", e)
            self._set_unavailable(reason)
            return False
        except httpx.HTTPError as e:
            reason = f"Ollama embed 测试失败: {e}"
            logger.warning("Ollama embed 测试失败: %s", e)
            self._set_unavailable(reason)
            return False

        # 全部通过
        self._set_available()
        logger.info("Ollama 检测通过: 模型 %s 可用", self._model)
        return True

    async def _check_openai_compat(self, client: httpx.AsyncClient) -> bool:
        """OpenAI 兼容模式的可用性检测：直接 embed 短文本"""
        try:
            resp = await client.post(
                "/v1/embeddings",
                json={"model": self._model, "input": "ping"},
                timeout=15.0,
            )
            resp.raise_for_status()
            data = resp.json().get("data", [])
            if not data:
                self._set_unavailable(f"模型 {self._model} 响应缺少 data 字段")
                return False
            embedding = data[0].get("embedding")
            if not embedding or len(embedding) == 0:
                self._set_unavailable(f"模型 {self._model} 返回空向量")
                return False
        except (httpx.ConnectError, httpx.TimeoutException) as e:
            reason = f"Embedding 服务不可用 ({self._ollama_url}): {e}"
            logger.warning("OpenAI 兼容服务不可用: %s", e)
            self._set_unavailable(reason)
            return False
        except httpx.HTTPError as e:
            reason = f"Embedding 服务异常: {e}"
            logger.warning("OpenAI 兼容服务异常: %s", e)
            self._set_unavailable(reason)
            return False

        self._set_available()
        logger.info("OpenAI 兼容 embedding 检测通过: %s @ %s", self._model, self._ollama_url)
        return True

    def _set_unavailable(self, reason: str) -> None:
        """标记服务不可用 + 记录冷却"""
        self._available = False
        self._unavailable_reason = reason
        self._cooldown.record_failure()

    def _set_available(self) -> None:
        """标记服务可用 + 重置冷却"""
        self._available = True
        self._unavailable_reason = ""
        self._cooldown.record_success()

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

    @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

    @property
    def api(self) -> str:
        return self._api

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