"""
本地 MLX Qwen3 Embedding（Apple Silicon）。

默认模型目录：项目根 ``Qwen3-Embedding-4B-mxfp8``，可用 ``VOC_EMBED_MODEL_PATH`` 覆盖。

在 M4 / 16GB、约 8GB 可用内存下实测（mxfp8，500 字符/条）：
  - 加载峰值约 1.5GB
  - batch 1–16 稳定；默认 ``VOC_EMBED_BATCH_SIZE=16``，串行推理（勿多线程并行加载同一 MLX 模型）

需 Python 3.10+ 与 ``mlx-embeddings>=0.1.0``（推荐项目 ``310py`` 虚拟环境）。
"""
from __future__ import annotations

import logging
import os
import sys
import types
from pathlib import Path
from typing import List, Sequence

logger = logging.getLogger("local_embedding")

PROJECT_ROOT = Path(__file__).resolve().parent
DEFAULT_MODEL_PATH = PROJECT_ROOT / "Qwen3-Embedding-4B-mxfp8"
# 实测短句 batch=32 仍 <0.5GB增量；长句 500 字符 batch=16 约 1.5GB 峰值
DEFAULT_BATCH_SIZE = 16
DEFAULT_MAX_TEXT_CHARS = 10_000


def _apply_hf_hub_shim() -> None:
    if "huggingface_hub.utils._errors" in sys.modules:
        return
    try:
        from huggingface_hub.errors import RepositoryNotFoundError
    except ImportError:
        try:
            from huggingface_hub.utils._errors import RepositoryNotFoundError  # type: ignore
        except ImportError:
            RepositoryNotFoundError = Exception  # type: ignore[misc, assignment]
    mod = types.ModuleType("huggingface_hub.utils._errors")
    mod.RepositoryNotFoundError = RepositoryNotFoundError
    sys.modules["huggingface_hub.utils._errors"] = mod


def _resolve_model_path() -> Path:
    raw = os.environ.get("VOC_EMBED_MODEL_PATH", "").strip()
    p = Path(raw).expanduser() if raw else DEFAULT_MODEL_PATH
    if not p.is_dir():
        raise FileNotFoundError(f"本地 embedding 模型目录不存在: {p}")
    return p.resolve()


def _resolve_batch_size(explicit: int | None = None) -> int:
    if explicit is not None and explicit > 0:
        return explicit
    env = os.environ.get("VOC_EMBED_BATCH_SIZE", "").strip()
    if env.isdigit() and int(env) > 0:
        return int(env)
    return DEFAULT_BATCH_SIZE


def _resolve_max_chars() -> int:
    env = os.environ.get("VOC_EMBED_MAX_TEXT_CHARS", "").strip()
    if env.isdigit() and int(env) > 0:
        return int(env)
    return DEFAULT_MAX_TEXT_CHARS


def _truncate(text: str, max_chars: int) -> str:
    t = (text or "").strip()
    if len(t) <= max_chars:
        return t
    return t[: max_chars - 3] + "..."


_MODEL = None
_PROCESSOR = None
_DIMENSIONS: int | None = None


def _load():
    global _MODEL, _PROCESSOR
    if _MODEL is not None:
        return _MODEL, _PROCESSOR
    _apply_hf_hub_shim()
    from mlx_embeddings import load

    path = _resolve_model_path()
    logger.info("加载本地 embedding: %s", path)
    _MODEL, _PROCESSOR = load(str(path))
    return _MODEL, _PROCESSOR


def embedding_dimensions() -> int:
    global _DIMENSIONS
    if _DIMENSIONS is not None:
        return _DIMENSIONS
    vecs = embed_texts(["dimension probe"])
    _DIMENSIONS = len(vecs[0])
    return _DIMENSIONS


def embed_texts(
    texts: Sequence[str],
    *,
    batch_size: int | None = None,
    max_chars: int | None = None,
) -> List[List[float]]:
    """返回与输入等长的浮点向量列表（已 L2 归一化）。"""
    if not texts:
        return []
    from mlx_embeddings import generate

    model, processor = _load()
    bs = _resolve_batch_size(batch_size)
    cap = max_chars if max_chars is not None else _resolve_max_chars()
    cleaned = [_truncate(t, cap) for t in texts]

    out_all: List[List[float]] = []
    n = len(cleaned)
    n_chunks = (n + bs - 1) // bs
    log_every = max(1, n_chunks // 20)

    for chunk_idx, i in enumerate(range(0, n, bs)):
        chunk = cleaned[i : i + bs]
        output = generate(model, processor, texts=chunk)
        emb = output.text_embeds
        for row in emb:
            out_all.append([float(x) for x in row.tolist()])

        done = chunk_idx + 1
        if done == 1 or done == n_chunks or done % log_every == 0:
            rows_done = min(done * bs, n)
            logger.info(
                "向量化进度 %s/%s 批（%s/%s 条）",
                done,
                n_chunks,
                rows_done,
                n,
            )
    global _DIMENSIONS
    if _DIMENSIONS is None and out_all:
        _DIMENSIONS = len(out_all[0])
    return out_all