"""embedding 服务客户端 — 调用 114:8003 Qwen3-Embedding 服务"""
import asyncio
import logging
from typing import List

import httpx

from app.config import EMBEDDING_URL, EMBEDDING_MODEL, EMBEDDING_BATCH_SIZE

logger = logging.getLogger(__name__)

_client: httpx.AsyncClient | None = None


def _get_client() -> httpx.AsyncClient:
    """复用 httpx 客户端"""
    global _client
    if _client is None or _client.is_closed:
        _client = httpx.AsyncClient(timeout=60)
    return _client


async def embed_batch(texts: List[str], batch_size: int = EMBEDDING_BATCH_SIZE) -> List[List[float]]:
    """批量向量化

    Args:
        texts: 文本列表
        batch_size: 每批大小,默认 32

    Returns:
        向量列表,顺序与输入一致

    Raises:
        RuntimeError: embedding 服务不可用或返回错误
    """
    if not texts:
        return []

    client = _get_client()
    all_vecs: List[List[float]] = []

    for i in range(0, len(texts), batch_size):
        batch = texts[i : i + batch_size]
        payload = {
            "model": EMBEDDING_MODEL,
            "input": batch,
        }

        # 重试 3 次,指数退避
        last_err = None
        for attempt in range(3):
            try:
                resp = await client.post(
                    f"{EMBEDDING_URL}/embeddings",
                    json=payload,
                )
                resp.raise_for_status()
                data = resp.json()
                # vLLM embeddings 响应格式: {data: [{embedding: [...]}], ...}
                vecs = [item["embedding"] for item in data["data"]]
                all_vecs.extend(vecs)
                last_err = None
                break
            except Exception as e:
                last_err = e
                wait = 2 ** attempt
                logger.warning(f"embedding 批次 {i//batch_size} 第 {attempt+1} 次失败: {e}, {wait}s 后重试")
                await asyncio.sleep(wait)

        if last_err is not None:
            raise RuntimeError(f"embedding 服务调用失败(重试 3 次): {last_err}")

    return all_vecs


async def embed_text(text: str) -> List[float]:
    """单条文本向量化"""
    vecs = await embed_batch([text])
    return vecs[0]