# -*- coding: utf-8 -*-
"""通义千问 DashScope 文本向量客户端(text-embedding-v3,1024 维)。

要点:
- Key 复用 PriceBot 的 QWEN_API_KEY;可用环境变量 DASHSCOPE_API_KEY / QWEN_API_KEY 覆盖。
- 接口走 OpenAI 兼容端点;单次批量上限 10(硬限),用线程池并行多批。
- name -> 向量 落盘缓存(cache/emb_<model>.jsonl):重复加载同一批菜名零成本、零计费。
"""
from __future__ import annotations

import json
import os
import sys
import threading
import time
from concurrent.futures import ThreadPoolExecutor, as_completed

import numpy as np
import requests

DASHSCOPE_URL = "https://dashscope.aliyuncs.com/compatible-mode/v1/embeddings"
# 默认复用 PriceBot 的通义千问 Key(本工具在 dish-images 下、不进 git);优先用环境变量覆盖。
DEFAULT_KEY = "sk-f66e689a9b0c43239e299137f68c453c"
BATCH = 10          # text-embedding-v3 单次最多 10 条
# 免费额度对"并发齐发"很敏感(冷却后多线程同时重试会互撞、永远挤不进窗口),默认单线程顺序发,
# 由下方自适应间隔 _delay 控速最稳。Key 额度充足时可设环境变量 EMBED_WORKERS=4 之类提速。
MAX_WORKERS = max(1, int(os.environ.get("EMBED_WORKERS", "1")))


def get_key(override: str | None = None) -> str:
    return (override or os.environ.get("DASHSCOPE_API_KEY")
            or os.environ.get("QWEN_API_KEY") or DEFAULT_KEY)


_tls = threading.local()


def _sess() -> requests.Session:
    """线程本地 session,trust_env=False:忽略系统代理(HTTP(S)_PROXY),直连 dashscope。
    Clash 等代理偶尔把请求甩到海外节点,会被阿里云按来源风控拒(HTTP 400 Access denied);
    dashscope.aliyuncs.com 是大陆公网端点,直连最稳(与 PriceBot 调千问一致)。"""
    s = getattr(_tls, "s", None)
    if s is None:
        s = _tls.s = requests.Session()
        s.trust_env = False
    return s


# ---------- 全局自适应限速:稳态间隔 + 冷却闸,收敛到免费额度可持续的速率 ----------
# dashscope 免费额度是「滚动每分钟预算」式:瞬时小量能过,持续高频会返回 400,且文案写成
# Arrearage/overdue(其实是"免费窗口用尽且无余额续费",约 1 分钟自动回血)。故这是【可重试】
# 信号:命中就放大间隔(降速)+ 拉长全员冷却(等回血),成功就缓慢提速,自动逼近可持续速率。
_pace_lock = threading.Lock()
_next_at = 0.0
_resume_at = 0.0
_delay = 0.6                    # 当前稳态请求间隔(秒),单线程顺序发
_cool = 5.0                     # 当前冷却时长(秒),命中放大、成功回落
DELAY_MIN, DELAY_MAX = 0.3, 4.0
COOL_MIN, COOL_MAX = 5.0, 45.0


def _pace() -> None:
    """按当前稳态间隔 _delay 给每个请求排队发牌(有效速率 ≈ 1/_delay)。"""
    global _next_at
    with _pace_lock:
        now = time.time()
        slot = max(now, _next_at)
        _next_at = slot + _delay
        wait = slot - now
    if wait > 0:
        time.sleep(wait)


def _gate_wait() -> None:
    """命中限速后的全员冷却,期间所有线程一起等(等额度回血)。"""
    while True:
        with _pace_lock:
            wait = _resume_at - time.time()
        if wait <= 0:
            return
        time.sleep(min(wait, 2.0) + 0.05)


def _rate_hit() -> None:
    """命中免费额度限速:降速 + 拉长冷却等回血。"""
    global _delay, _cool, _resume_at
    with _pace_lock:
        _delay = min(_delay * 1.5, DELAY_MAX)
        _cool = min(_cool * 1.4, COOL_MAX)
        _resume_at = max(_resume_at, time.time() + _cool)


def _ok() -> None:
    """成功:缓慢提速、缩短冷却,收敛到刚好不撞墙。"""
    global _delay, _cool
    with _pace_lock:
        _delay = max(_delay * 0.9, DELAY_MIN)
        _cool = max(_cool * 0.9, COOL_MIN)


class ArrearsError(RuntimeError):
    """重试多轮仍被免费额度拒(疑似额度耗尽/欠费),供上层提示"稍后重试/换 Key"。"""


def _is_rate_signal(status: int, text: str) -> bool:
    """免费额度限速信号(可重试):含真限速(429/503/throttling)与"伪欠费"(Arrearage/overdue,
    实为免费窗口用尽、约 1 分钟回血)。两者都靠退避+等待解决。"""
    if status in (429, 503):
        return True
    low = text.lower()
    return ("throttling" in low or "rate limit" in low or "requests rate" in low
            or "arrearage" in low or "overdue" in low or "good standing" in low
            or "欠费" in text or "余额不足" in text)


class Embedder:
    def __init__(self, model: str = "text-embedding-v4", cache_dir: str = "cache",
                 key: str | None = None):
        self.model = model
        self.key = get_key(key)
        os.makedirs(cache_dir, exist_ok=True)
        self.cache_path = os.path.join(cache_dir, f"emb_{model}.jsonl")
        self.cache: dict[str, list[float]] = {}
        self.last_failed = 0
        self.last_error = ""
        self._lock = threading.Lock()
        self._load_cache()

    # ---------- 缓存 ----------
    def _load_cache(self) -> None:
        if not os.path.exists(self.cache_path):
            return
        with open(self.cache_path, encoding="utf-8") as f:
            for line in f:
                try:
                    o = json.loads(line)
                    self.cache[o["t"]] = o["v"]
                except Exception:
                    pass

    def _append_cache(self, items: list[tuple[str, list[float]]]) -> None:
        with self._lock:
            with open(self.cache_path, "a", encoding="utf-8") as f:
                for t, v in items:
                    f.write(json.dumps({"t": t, "v": v}, ensure_ascii=False) + "\n")
                    self.cache[t] = v

    # ---------- 接口 ----------
    def _embed_batch(self, texts: list[str]) -> list[list[float]]:
        body = {"model": self.model, "input": texts, "encoding_format": "float"}
        last = None
        rate_hits = 0
        for _ in range(7):           # 配合冷却升级,总等待可超 1 分钟,足以熬过免费窗口回血
            _gate_wait()
            _pace()
            try:
                r = _sess().post(
                    DASHSCOPE_URL,
                    headers={"Authorization": f"Bearer {self.key}",
                             "Content-Type": "application/json"},
                    json=body, timeout=30)
                if r.status_code == 200:
                    _ok()
                    data = sorted(r.json()["data"], key=lambda d: d["index"])
                    return [d["embedding"] for d in data]
                txt = r.text[:200]
                last = f"HTTP {r.status_code}: {txt}"
                if _is_rate_signal(r.status_code, txt):
                    rate_hits += 1
                    _rate_hit()          # 降速 + 拉长冷却等回血后重试
                else:
                    time.sleep(1.0)      # 其它错误轻退避后重试
            except requests.RequestException as e:
                last = f"{type(e).__name__}: {e}"
                time.sleep(2)
        if rate_hits >= 4:               # 多轮仍被额度拒 → 让上层提示换Key/稍后重试
            raise ArrearsError(f"模型 {self.model} 免费额度反复受限(可能已用尽);"
                               f"可稍后重试,或在 embed.py 换 DASHSCOPE_API_KEY")
        raise RuntimeError(f"embed batch failed: {last}")

    def embed_corpus(self, names: list[str], progress=None) -> None:
        """把 names 全部编码进缓存。progress(done, total) 回调可选。失败的批跳过(不致命)。"""
        todo = [n for n in dict.fromkeys(names) if n not in self.cache]
        total = len(todo)
        self.last_failed = 0
        self.last_error = ""
        if progress:
            progress(0, total)
        if not total:
            return
        batches = [todo[i:i + BATCH] for i in range(0, total, BATCH)]
        done = 0
        ok_count = 0
        with ThreadPoolExecutor(max_workers=MAX_WORKERS) as ex:
            futs = {ex.submit(self._embed_batch, b): b for b in batches}
            for fut in as_completed(futs):
                b = futs[fut]
                try:
                    self._append_cache(list(zip(b, fut.result())))
                    ok_count += len(b)
                except Exception as e:  # noqa: BLE001
                    self.last_failed += len(b)
                    self.last_error = str(e)[:200]
                    if self.last_failed <= len(b) * 2:   # 只打前一两次,避免刷屏
                        print(f"[embed] 批失败({self.last_failed}): {e}", file=sys.stderr, flush=True)
                    # 一次都没成功 + 已连续失败 ≥2 批 → 判定 Key 不可用,中止省额度;
                    # 若前面已有成功(只是滚动额度临时受限),则继续磨,失败的批留待下次续抓。
                    if ok_count == 0 and self.last_failed >= 20:
                        print("[embed] 持续失败且零成功,中止向量编码(BM25 仍可用)",
                              file=sys.stderr, flush=True)
                        ex.shutdown(wait=False, cancel_futures=True)
                        break
                done += len(b)
                if progress:
                    progress(done, total)

    # ---------- 取用 ----------
    def matrix(self, names: list[str]) -> tuple[list[str], np.ndarray]:
        """返回 (有向量的 name 列表, L2 归一化矩阵 N×dim)。"""
        keep = [n for n in names if n in self.cache]
        if not keep:
            return [], np.zeros((0, 0), dtype=np.float32)
        m = np.array([self.cache[n] for n in keep], dtype=np.float32)
        norm = np.linalg.norm(m, axis=1, keepdims=True)
        norm[norm == 0] = 1.0
        return keep, m / norm

    def embed_query(self, q: str) -> np.ndarray:
        v = np.array(self._embed_batch([q])[0], dtype=np.float32)
        n = np.linalg.norm(v)
        return v / (n if n else 1.0)