"""RAG 管道：Embedding（本地模型，ModelScope 下载，离线可用）"""
import sys
import threading
import math
import logging
import time
import os
from functools import lru_cache
from pathlib import Path
from ..config import MODELS_DIR, EMBEDDING_MODEL

_lock = threading.Lock()
_encoder = None
_encoder_name = None
_encoder_load_ms = 0.0
_metric_lock = threading.Lock()
_last_embedding_metric: dict = {}
_embedding_dimension = ""
logger = logging.getLogger("yawen.rag.embed")

# This is deliberately a separate value from the model name.  Operators can
# bump it when replacing files in the model directory without changing the
# configured model identifier (which also invalidates query-vector caches).
EMBEDDING_MODEL_VERSION = os.environ.get("EMBEDDING_MODEL_VERSION", "1")


def _set_metric(row: dict) -> None:
    """Keep the most recent bounded embedding sample for the RAG envelope.

    The RAG worker is single threaded today, but protecting this assignment
    makes the helper safe when a worker is configured with concurrent requests.
    No source text is retained in telemetry.
    """
    global _last_embedding_metric
    clean = {
        "process_start_ms": 0.0,
        "ipc_ms": 0.0,
        "model_load_ms": 0.0,
        "tokenizer_ms": 0.0,
        "encode_ms": 0.0,
        "postprocess_ms": 0.0,
        "embedding_source": "sentence-transformers",
        "embedding_count": 0,
        "model": str(EMBEDDING_MODEL),
        "model_version": str(EMBEDDING_MODEL_VERSION),
        "pid": os.getpid(),
        "cache_hit": False,
    }
    clean.update({key: value for key, value in row.items() if value is not None})
    for key in ("process_start_ms", "ipc_ms", "model_load_ms", "tokenizer_ms",
                "encode_ms", "postprocess_ms"):
        try:
            clean[key] = round(float(clean[key]), 2)
        except (TypeError, ValueError):
            clean[key] = 0.0
    clean["embedding_count"] = int(clean.get("embedding_count") or 0)
    clean["pid"] = int(clean.get("pid") or os.getpid())
    with _metric_lock:
        _last_embedding_metric = clean


def last_embedding_metric() -> dict:
    """Return phase timings for the latest encode without exposing input text."""
    with _metric_lock:
        return dict(_last_embedding_metric)


def query_cache_snapshot() -> dict:
    """Expose the independent query-vector LRU counters for diagnostics."""
    info = _cached_query.cache_info()
    total = info.hits + info.misses
    return {
        "size": info.currsize, "max_size": info.maxsize,
        "hits": info.hits, "misses": info.misses,
        "hit_rate": round(info.hits / total, 4) if total else 0.0,
        "model": str(EMBEDDING_MODEL), "model_version": str(EMBEDDING_MODEL_VERSION),
        "dimension": int(_embedding_dimension or 0),
        "pid": os.getpid(),
    }


def clear_query_cache() -> None:
    """Clear only query vectors; retrieval candidate pools remain untouched."""
    _cached_query.cache_clear()


def set_transport_timings(*, process_start_ms: float = 0.0, ipc_ms: float = 0.0) -> None:
    """Attach process/IPC timings measured by the parent worker transport.

    ``embed.py`` cannot observe subprocess creation or pipe transfer itself;
    the worker pool can call this after a request and the next envelope then
    contains the complete set of phases.  Keeping this optional preserves the
    existing embedding API for in-process callers and tests.
    """
    metric = last_embedding_metric()
    metric.update(process_start_ms=process_start_ms, ipc_ms=ipc_ms)
    _set_metric(metric)


def ensure_model_downloaded(model_name: str) -> Path:
    """模型目录存在即用；否则从 ModelScope 下载（HF 被墙）"""
    safe = model_name.replace("/", "--")
    local = MODELS_DIR / safe
    if (local / "config.json").exists():
        return local
    try:
        from modelscope import snapshot_download
        print(f"[embed] 从 ModelScope 下载 {model_name} ...", file=sys.stderr)
        snapshot_download(model_name, local_dir=str(local))
        return local
    except Exception as e:
        raise RuntimeError(f"模型下载失败（ModelScope）: {e}")


def get_encoder():
    """单例加载 SentenceTransformer（本地路径）"""
    global _encoder, _encoder_name, _encoder_load_ms
    if _encoder is not None and _encoder_name == EMBEDDING_MODEL:
        return _encoder
    with _lock:
        if _encoder is not None and _encoder_name == EMBEDDING_MODEL:
            return _encoder
        started = time.perf_counter()
        from sentence_transformers import SentenceTransformer
        local = ensure_model_downloaded(EMBEDDING_MODEL)
        print(f"[embed] 加载 {local} ...", file=sys.stderr)
        _encoder = SentenceTransformer(str(local))
        _encoder_name = EMBEDDING_MODEL
        _encoder_load_ms = (time.perf_counter() - started) * 1000
        return _encoder


def prewarm() -> bool:
    """Load the embedding model and run one tiny encode as a best-effort warmup.

    Model loading is intentionally kept out of the API startup path.  The
    caller should run this function from a daemon thread so a slow model
    download/load never delays health checks or request handling.  Returning a
    boolean keeps this useful for startup diagnostics without making warmup
    failures fatal to the service.
    """
    started = time.perf_counter()
    try:
        encoder = get_encoder()
        # A real encode initializes tokenizer/runtime kernels as well as the
        # SentenceTransformer object.  The result is discarded; query vectors
        # continue to use the normal cached path.
        encoder.encode(["系统启动预热"], normalize_embeddings=True,
                       show_progress_bar=False, batch_size=1)
        logger.info("Embedding 模型预热完成：%s，耗时 %.3f 秒",
                    EMBEDDING_MODEL, time.perf_counter() - started)
        return True
    except Exception:
        logger.exception("Embedding 模型预热失败：%s，服务将按需加载", EMBEDDING_MODEL)
        return False


EMBEDDING_POLICY = "full-token-windows-v1"


def text_windows(text: str, encoder=None) -> list[str]:
    """Cover every source character without sending over-limit windows to BGE."""
    enc = encoder or get_encoder()
    tokenizer = enc.tokenizer
    maximum = int(enc.max_seq_length)
    def fits(value):
        return len(tokenizer.encode(value, add_special_tokens=True, truncation=False)) <= maximum
    if fits(text):
        return [text]
    # Character boundaries preserve exact original text; a binary search uses the
    # actual model tokenizer rather than assuming that one character is one token.
    windows, offset = [], 0
    while offset < len(text):
        remaining = text[offset:]
        if fits(remaining):
            windows.append(remaining)
            break
        low, high = 1, min(len(remaining), maximum * 4)
        while low < high:
            middle = (low + high + 1) // 2
            if fits(remaining[:middle]):
                low = middle
            else:
                high = middle - 1
        if not fits(remaining[:low]):
            raise ValueError("单个字符超过向量模型输入上限")
        windows.append(remaining[:low])
        offset += low
    return windows


def embed_texts(texts: list[str]) -> list[list[float]]:
    if not texts:
        return []
    # ``get_encoder`` includes model download and construction only on the
    # first call in each resident worker.  Subsequent calls report zero load
    # time, which makes accidental worker restarts visible in telemetry.
    enc = get_encoder()
    global _encoder_load_ms
    model_load_ms = _encoder_load_ms
    # Consume the one-shot load measurement.  A resident encoder must report
    # zero on later requests; a non-zero value then reliably indicates a worker
    # restart or an explicit model replacement.
    _encoder_load_ms = 0.0
    # At most 64 originals and 64 encoded windows in one model batch. Long
    # paragraphs are pooled in the existing vector space, keeping chunk IDs and
    # consumption records stable while representing their tails as well.
    output = []
    tokenizer_ms = encode_ms = postprocess_ms = 0.0
    flat_count = 0
    for start in range(0, len(texts), 64):
        tokenizer_started = time.perf_counter()
        groups = [text_windows(str(text), enc) for text in texts[start:start + 64]]
        flat = [window for group in groups for window in group]
        tokenizer_ms += (time.perf_counter() - tokenizer_started) * 1000
        flat_count += len(flat)
        encode_started = time.perf_counter()
        vectors = enc.encode(flat, normalize_embeddings=True, show_progress_bar=False, batch_size=64)
        encode_ms += (time.perf_counter() - encode_started) * 1000
        postprocess_started = time.perf_counter()
        cursor = 0
        for group in groups:
            selected = vectors[cursor:cursor + len(group)]
            cursor += len(group)
            if len(group) == 1:
                output.append([float(v) for v in selected[0]])
                continue
            weights = [max(1, len(enc.tokenizer.encode(window, add_special_tokens=False, truncation=False)))
                       for window in group]
            combined = [sum(float(vector[i]) * weight for vector, weight in zip(selected, weights))
                        for i in range(len(selected[0]))]
            norm = math.sqrt(sum(value * value for value in combined)) or 1.0
            output.append([value / norm for value in combined])
        postprocess_ms += (time.perf_counter() - postprocess_started) * 1000
    _set_metric({
        "model_load_ms": model_load_ms,
        "tokenizer_ms": tokenizer_ms,
        "encode_ms": encode_ms,
        "postprocess_ms": postprocess_ms,
        "embedding_count": len(texts),
        # ``flat_count`` is useful when long documents are split into windows;
        # retain it as a diagnostic without changing the public return value.
        "window_count": flat_count,
        "embedding_dimension": len(output[0]) if output else 0,
        "cache_hit": False,
    })
    return output


@lru_cache(maxsize=256)
def _cached_query(model: str, policy: str, version: str, dimension: str, text: str) -> tuple[float, ...]:
    """Cache only query vectors, keyed by model identity and policy version."""
    return tuple(embed_texts([text])[0])


def embed_query(text: str) -> list[float]:
    # lru_cache does not expose whether this invocation hit without comparing
    # its counters.  The comparison is process local and avoids retaining the
    # query itself in telemetry.
    before = _cached_query.cache_info()
    global _embedding_dimension
    # Keep the cache identity stable across the first call: the actual vector
    # dimension is discovered only after encoding.  Model/version already
    # identify the encoder contract; ``auto`` prevents the first "unknown"
    # dimension from creating a second key for the same query.
    result = list(_cached_query(EMBEDDING_MODEL, EMBEDDING_POLICY,
                                EMBEDDING_MODEL_VERSION, "auto", str(text)))
    if not _embedding_dimension and result:
        _embedding_dimension = str(len(result))
    after = _cached_query.cache_info()
    cache_hit = after.hits > before.hits
    metric = last_embedding_metric()
    metric.update(
        embedding_count=1,
        cache_hit=cache_hit,
        embedding_source="query-vector-lru",
        model=EMBEDDING_MODEL,
        model_version=EMBEDDING_MODEL_VERSION,
        pid=os.getpid(),
        embedding_dimension=len(result) if result else 0,
    )
    if cache_hit:
        # A cache hit performs no tokenizer, model inference, or serialization;
        # do not leak timings from the previous (possibly indexing) call.
        metric.update(model_load_ms=0.0, tokenizer_ms=0.0,
                      encode_ms=0.0, postprocess_ms=0.0, window_count=0)
    _set_metric(metric)
    return result