import logging
import os
import urllib.request

from app.config import EMBEDDING_MODEL

logger = logging.getLogger("learn_rag")
QUERY_PREFIX = "为这个句子生成表示以用于检索相关文章："


def configure_download() -> None:
    """Hugging Face 不可达时改用国内镜像，避免首次下载卡住。"""
    if os.getenv("HF_ENDPOINT"):
        return
    try:
        with urllib.request.urlopen("https://huggingface.co", timeout=4) as response:
            if response.status < 400:
                return
    except Exception:
        logger.info("无法连接 huggingface.co，嵌入模型改从 hf-mirror.com 下载")
        os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"


class Embedder:
    def __init__(self) -> None:
        configure_download()
        from fastembed import TextEmbedding

        logger.info("正在加载嵌入模型 %s", EMBEDDING_MODEL)
        self.model = TextEmbedding(model_name=EMBEDDING_MODEL)
        probe = next(iter(self.model.embed(["书店"])))
        self.dimension = len(probe)
        logger.info("嵌入模型就绪，维度 %s", self.dimension)

    def embed_documents(self, texts: list[str]) -> list[list[float]]:
        if not texts:
            return []
        return [vector.tolist() for vector in self.model.embed(texts)]

    def embed_query(self, text: str) -> list[float]:
        query_embed = getattr(self.model, "query_embed", None)
        if query_embed is not None:
            try:
                return next(iter(query_embed([text]))).tolist()
            except Exception:
                logger.info("query_embed 不可用，改用查询前缀")
        prefixed = QUERY_PREFIX + text
        return next(iter(self.model.embed([prefixed]))).tolist()