"""两级语义缓存（T-P6，M2）：L1 精确 + L2 n-gram 倒排 + singleflight + SSE 回放。

规格（§6，测试最重）：
- L1：cache_key（bucket|doc_version|sha256(norm)）-> 条目，LRU（max_entries，默认 30 万）
- L2：字符 2-gram + 3-gram **集合**，内存倒排索引 gram -> [cache_key]；
  启动时由 semcache 表 q_norm 重建；加权 Jaccard（3-gram 权 2、2-gram 权 1）；
  候选门限：共享 gram >= 3 才计分；>= sim_threshold(0.92) 命中；
  **L2 命中累计 promote_frequency(5) 次晋升 L1**
- TTL：ttl_ts 过期不可见；命中即续期（滑动过期）
- 持久化：semcache 表（put 同步写，索引内存维护；调用方 to_thread，D-P10）
"""
from __future__ import annotations

import asyncio
import json
import re
import time
from collections import OrderedDict
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

_WHITESPACE = re.compile(r"\s+")


def grams(text: str) -> set:
    """字符 2-gram + 3-gram 集合（中文天然适配，无需分词）。"""
    t = _WHITESPACE.sub("", (text or "").lower())
    out: set = set()
    for n in (2, 3):
        for i in range(len(t) - n + 1):
            out.add(t[i:i + n])
    return out or ({t} if t else set())


def weighted_jaccard(ga: set, gb: set) -> float:
    """加权 Jaccard：交集中每个 3-gram 权 2、2-gram 权 1，除以并集加权。"""
    if not ga or not gb:
        return 0.0
    inter = ga & gb
    if not inter:
        return 0.0
    w_inter = sum(2 if g in (3,) or len(g) == 3 else 1 for g in inter)
    w_union = sum(2 if len(g) == 3 else 1 for g in (ga | gb))
    return w_inter / w_union if w_union else 0.0


class CacheEntry:
    __slots__ = ("answer", "model", "q_norm", "g", "created_ts", "ttl_ts",
                 "doc_version", "hits")

    def __init__(self, answer: str, model: str, q_norm: str,
                 created_ts: float, ttl_ts: float, doc_version: int):
        self.answer = answer
        self.model = model
        self.q_norm = q_norm
        self.g = grams(q_norm)
        self.created_ts = created_ts
        self.ttl_ts = ttl_ts
        self.doc_version = doc_version
        self.hits = 0


class SemanticCache:
    """L1 精确 + L2 倒排（内存），sqlite 持久化（store 的 semcache 表）。"""

    def __init__(self, store, max_entries: int = 300000,
                 sim_threshold: float = 0.92, promote_frequency: int = 5,
                 now: Optional[float] = None):
        self.store = store
        self.max_entries = max(1, int(max_entries))
        self.sim_threshold = float(sim_threshold)
        self.promote_frequency = max(1, int(promote_frequency))
        self._l1: "OrderedDict[str, CacheEntry]" = OrderedDict()
        self._inverted: Dict[str, set] = {}
        self._clock = now            # 假时钟注入（None = time.time）
        self.hits_exact = 0
        self.hits_semantic = 0
        self._rebuild()

    # ---------- 时钟 ----------
    def _now(self) -> float:
        return time.time() if self._clock is None else self._clock

    # ---------- 启动重建 ----------
    def _rebuild(self) -> None:
        try:
            rows = self.store.semcache_rows(limit=self.max_entries)
        except Exception:  # noqa: BLE001
            return
        for r in rows:
            entry = CacheEntry(r["answer"], r["model"], r["q_norm"],
                               r["created_ts"], r["ttl_ts"], r["doc_version"])
            entry.hits = r["hits"]
            self._index(r["cache_key"], entry, promote=False)

    def _index(self, cache_key: str, entry: CacheEntry, promote: bool = True) -> None:
        self._l1[cache_key] = entry
        self._l1.move_to_end(cache_key)
        for g in entry.g:
            self._inverted.setdefault(g, set()).add(cache_key)
        if promote:
            self._evict()

    def _evict(self) -> None:
        """LRU 驱逐（超 max_entries 淘汰最久未用，含倒排回收）。"""
        while len(self._l1) > self.max_entries:
            key, entry = self._l1.popitem(last=False)
            for g in entry.g:
                bucket = self._inverted.get(g)
                if bucket is not None:
                    bucket.discard(key)
                    if not bucket:
                        self._inverted.pop(g, None)

    # ---------- 查询 ----------
    def lookup(self, cache_key: str, norm_text: str,
               doc_version: int = 1) -> Optional[Dict[str, Any]]:
        """L1 精确 -> L2 语义（§7 签名）。返回 {answer, model, level} 或 None。"""
        now = self._now()
        # L1
        entry = self._l1.get(cache_key)
        if entry is not None:
            if entry.ttl_ts < now:
                self._invalidate_key(cache_key)
            else:
                self.hits_exact += 1
                self._l1.move_to_end(cache_key)
                return {"answer": entry.answer, "model": entry.model,
                        "level": "exact"}
        # L2
        g = grams(norm_text)
        if len(g) < 3:
            return None
        candidates: Dict[str, int] = {}
        for gram in g:
            for key in self._inverted.get(gram, ()):
                candidates[key] = candidates.get(key, 0) + 1
        best_key = None
        best_score = 0.0
        for key, shared in candidates.items():
            if shared < 3:
                continue                       # 候选门限：共享 gram >= 3
            cand = self._l1.get(key)
            if cand is None or cand.ttl_ts < now or cand.doc_version != doc_version:
                continue
            score = weighted_jaccard(g, cand.g)
            if score > best_score:
                best_score = score
                best_key = key
        if best_key is None or best_score < self.sim_threshold:
            return None
        cand = self._l1[best_key]
        cand.hits += 1
        self.hits_semantic += 1
        promoted = False
        if cand.hits >= self.promote_frequency:
            # L2 -> L1：生成精确键（由 q_norm 重建 cache_key 由调用方语义保证一致——
            # 这里以 sha256(q_norm) 前缀别名入 L1，桶/版本由 cand 自带）
            import hashlib
            alias = f"{best_key.split('|')[0]}|{cand.doc_version}|promoted:" \
                    f"{hashlib.sha256(cand.q_norm.encode()).hexdigest()[:16]}"
            self._index(alias, cand, promote=True)
            try:
                self.store.promote_semcache(best_key, alias, cand.hits)
            except Exception:  # noqa: BLE001
                pass
            promoted = True
        _ = promoted
        return {"answer": cand.answer, "model": cand.model, "level": "semantic"}

    def _invalidate_key(self, cache_key: str) -> None:
        entry = self._l1.pop(cache_key, None)
        if entry:
            for g in entry.g:
                bucket = self._inverted.get(g)
                if bucket is not None:
                    bucket.discard(cache_key)
                    if not bucket:
                        self._inverted.pop(g, None)

    # ---------- 写入 ----------
    def put(self, cache_key: str, q_norm: str, answer: str, model: str,
            doc_version: int = 1, ttl_hours: int = 72) -> None:
        """写 L1 + 倒排 + sqlite 持久化（§7 签名）。"""
        now = self._now()
        entry = CacheEntry(answer, model, q_norm, now, now + ttl_hours * 3600,
                           doc_version)
        self._index(cache_key, entry)
        try:
            self.store.put_semcache(cache_key, "default", q_norm, answer, model,
                                    int(now), int(now + ttl_hours * 3600),
                                    doc_version)
        except Exception:  # noqa: BLE001
            pass  # 持久化失败不影响内存缓存（可重建）

    def stats(self) -> Dict[str, Any]:
        return {"entries": len(self._l1), "hits_exact": self.hits_exact,
                "hits_semantic": self.hits_semantic}


class SingleFlight:
    """同请求合并（§7）：dict[norm_hash -> Future]，上限 256，60s 超时降级。"""

    MAX = 256
    TIMEOUT_S = 60.0

    def __init__(self):
        self._inflight: Dict[str, asyncio.Future] = {}

    def try_claim(self, norm_hash: str):
        """返回 (future_or_None, slot)。future 非 None = 等待方；slot = 登记句柄。"""
        fut = self._inflight.get(norm_hash)
        if fut is not None:
            return fut, None
        if len(self._inflight) >= self.MAX:
            return None, None                        # 超限旁路（不合并）
        fut = asyncio.get_event_loop().create_future()
        self._inflight[norm_hash] = fut
        return None, (norm_hash, fut)

    def release(self, slot, result: Any = None, error: Any = None) -> None:
        if slot is None:
            return
        norm_hash, fut = slot
        self._inflight.pop(norm_hash, None)
        if not fut.done():
            if error is not None:
                fut.set_exception(error)
            else:
                fut.set_result(result)

    async def wait(self, fut, timeout_s: float = TIMEOUT_S):
        """等待方：超时 -> 降级直连（返回 None）。"""
        try:
            return await asyncio.wait_for(asyncio.shield(fut), timeout=timeout_s)
        except (asyncio.TimeoutError, Exception):
            return None


# ---------------- SSE 合成回放（§6） ----------------

def synth_sse_chunks(answer: str, chunk_size: int = 20,
                     model: str = "cached", request_id: str = "") -> List[bytes]:
    """缓存命中且 stream=true：合成合法 SSE（分块 delta + finish + [DONE]）。"""
    cid = f"chatcmpl-cached-{request_id or '0'}"
    created = int(time.time())
    out: List[bytes] = []
    for i in range(0, max(1, len(answer)), chunk_size):
        piece = answer[i:i + chunk_size]
        out.append(("data: " + json.dumps({
            "id": cid, "object": "chat.completion.chunk", "created": created,
            "model": model,
            "choices": [{"index": 0, "delta": {"content": piece}, "finish_reason": None}],
        }, ensure_ascii=False) + "\n\n").encode("utf-8"))
    out.append(("data: " + json.dumps({
        "id": cid, "object": "chat.completion.chunk", "created": created,
        "model": model,
        "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
    }, ensure_ascii=False) + "\n\n").encode("utf-8"))
    out.append(b"data: [DONE]\n\n")
    return out