"""广告分类 Predictor：把 laya-ad 微调的 mmBERT 二分类模型接入 Jev 兼容 API。

包含此前工作流验证的 CPU 优化：
- 动态微批：并发请求在 15ms 窗口内合批前向（单 worker 多线程下 100 并发整批 ≤1s SLO 的关键）
- CPU 线程数：LAYA_CPU_THREADS（默认物理核数 ≈ 逻辑核一半，Zen5 实测最优档）
- bf16 推理：Zen5 AVX-512 BF16 有效

支持的题目（其余标记 unsupported 并返回无信息中性答案，绝不冒充判断）：
- choice 且标签恰为 {ad, clean}（大小写不敏感）→ 分类器直接输出
- noul 且 instructions 提及广告（ad/advert/广告）→ "noul" 字段输出 P(is_ad)

usage.input_tokens 为真实 token 数；output_tokens = 题目数（每题一个决策位）。
"""
from __future__ import annotations

import json
import os
import queue
import re
import threading
import time
from pathlib import Path
from typing import Any, Callable, Protocol

import torch

from .quant_loader import detect_precision, load_ad_model

_AD_RE = re.compile(r"ad|advert|广告", re.IGNORECASE)
_AD_LABELS = {"ad", "clean"}


class Predictor(Protocol):
    """Jev 推理 Predictor 协议（与上游 LayaAdapter 相同的形状契约）。"""

    def predict(self, state: Any, questions: dict[str, Any], model: str | None = None,
                *, max_len: int | None = None) -> dict[str, Any]: ...


def model_identity(model_dir: Path) -> tuple[str, str]:
    """读取模型目录身份：(model_id, precision)。缺省文件时回退通用名。"""
    model_dir = Path(model_dir)
    precision = "unknown"
    p_file = model_dir / "PRECISION.txt"
    if p_file.is_file():
        precision = p_file.read_text(encoding="utf-8").strip() or "unknown"
    mid_file = model_dir / "MODEL_ID.txt"
    model_id = mid_file.read_text(encoding="utf-8").strip() if mid_file.is_file() else ""
    if not model_id:
        model_id = "laya" if precision == "fp8" else f"laya-{precision}"
    return model_id, precision


def _version_of(model_id: str) -> str | None:
    """'laya-v1.0-fp8' → 'v1.0'；不符合 laya-<版本>-<精度> 模式则 None。"""
    if not model_id.startswith("laya-"):
        return None
    rest = model_id[len("laya-"):]
    return rest.rsplit("-", 1)[0] if "-" in rest else None


def _accepted_ids(model_id: str, precision: str) -> list[str]:
    """单模型镜像可接受的 model id 集合（规范序：auto, laya, laya-<版本>, 完整 id）。
    laya/laya-<版本> 别名仅 fp8 默认镜像提供。"""
    ids = ["auto"]
    if precision == "fp8":
        version = _version_of(model_id)
        if version:
            ids += ["laya", f"laya-{version}"]
    ids.append(model_id)
    return ids


def _extract_text(state: Any) -> str:
    """把 Jev state 序列化为分类文本：消息正文优先，随后 sender/bio/bio_links、
    matched_rules、recent_context、known_ad_patterns/false_positives 也进入输入。"""
    if isinstance(state, str):
        return state
    if isinstance(state, list):
        parts = []
        for turn in state:
            if isinstance(turn, dict) and isinstance(turn.get("text"), str):
                parts.append(turn["text"])
        return "\n".join(parts) if parts else json.dumps(state, ensure_ascii=False)
    if not isinstance(state, dict):
        return str(state)

    msg = state.get("message")
    if isinstance(msg, dict):
        text = msg.get("text") or ""
    elif isinstance(msg, str):
        text = msg
    else:
        text = ""
    if not isinstance(text, str):
        text = str(text)
    if not text:
        for key in ("text", "body", "content"):
            if isinstance(state.get(key), str):
                text = state[key]
                break

    parts: list[str] = []
    if text:
        parts.append("［消息］" + text)

    chat = state.get("chat")
    if isinstance(chat, dict) and chat.get("title"):
        parts.append("［群］" + str(chat["title"]))

    sender = state.get("sender")
    if isinstance(sender, dict):
        bits = []
        name = " ".join(str(sender[k]) for k in ("first_name", "last_name")
                        if sender.get(k)).strip()
        if name:
            bits.append("昵称: " + name)
        if sender.get("username"):
            bits.append("用户名: @" + str(sender["username"]).lstrip("@"))
        if sender.get("bio"):
            bits.append("简介: " + str(sender["bio"]))
        if sender.get("age_hours") is not None:
            bits.append(f"入群时长(h): {sender['age_hours']}")
        if sender.get("msgs_in_group") is not None:
            bits.append(f"群内发言数: {sender['msgs_in_group']}")
        if sender.get("prior_ad_hits") is not None:
            bits.append(f"历史广告命中: {sender['prior_ad_hits']}")
        if bits:
            parts.append("［发送者］" + " | ".join(bits))
        links = sender.get("bio_links")
        if isinstance(links, list) and links:
            lb = []
            for x in links[:4]:
                if not isinstance(x, dict):
                    continue
                s = " / ".join(str(x[k]) for k in ("kind", "title", "handle", "about")
                               if x.get(k))
                if s:
                    lb.append(s)
            if lb:
                parts.append("［资料链接］" + " ; ".join(lb))

    rules = state.get("matched_rules")
    if isinstance(rules, list) and rules:
        rb = []
        for x in rules[:4]:
            if not isinstance(x, dict):
                continue
            rb.append(" / ".join(str(x[k]) for k in ("name", "category", "note")
                                 if x.get(k)))
        if rb:
            parts.append("［命中规则］" + " ; ".join(rb))

    ctx = state.get("recent_context")
    if isinstance(ctx, list) and ctx:
        cb = []
        for x in ctx[-5:]:
            if isinstance(x, dict):
                cb.append(f"{x.get('user', '')}: {x.get('text', '')}")
            elif isinstance(x, str):
                cb.append(x)
        if cb:
            parts.append("［近期消息］" + " / ".join(cb))

    for key, label in (("known_ad_patterns", "［已知广告模式］"),
                       ("known_false_positives", "［已知误报］")):
        val = state.get(key)
        if val:
            parts.append(label + str(val)[:400])

    if isinstance(msg, dict):
        flags = []
        for key, label in (("has_link", "含链接"), ("has_mention", "含@"),
                           ("has_media", "含媒体"), ("is_forwarded", "转发")):
            if msg.get(key):
                flags.append(label)
        if flags:
            parts.append("［消息特征］" + " ".join(flags))

    return "\n".join(parts) if parts else json.dumps(state, ensure_ascii=False)


def _labels(criteria: Any) -> list[str]:
    if isinstance(criteria, dict):
        return [str(k) for k in criteria]
    if isinstance(criteria, list):
        return [str(x) for x in criteria]
    return []


class _MicroBatcher:
    """请求合批器：把并发单条请求在窗口内合并成一次 padding 前向。"""

    def __init__(self, forward_fn: Callable[[list[str], int], list[float]],
                 max_batch: int, window_ms: int):
        self.forward_fn = forward_fn
        self.max_batch = max(1, max_batch)
        self.window_s = max(0, window_ms) / 1000.0
        self._queue: queue.Queue = queue.Queue()
        self._worker = threading.Thread(target=self._loop, daemon=True,
                                        name="ad-clf-microbatch")
        self._worker.start()

    def submit(self, text: str, seq_len: int) -> tuple[float, float]:
        item = {"text": text, "seq_len": seq_len, "event": threading.Event()}
        self._queue.put(item)
        item["event"].wait()
        if item.get("error") is not None:
            raise item["error"]
        return item["probs"]

    def _loop(self) -> None:
        while True:
            batch = [self._queue.get()]
            deadline = time.monotonic() + self.window_s
            while len(batch) < self.max_batch:
                remaining = deadline - time.monotonic()
                if remaining <= 0:
                    break
                try:
                    batch.append(self._queue.get(timeout=remaining))
                except queue.Empty:
                    break
            try:
                seq_len = min(512, max(b["seq_len"] for b in batch))
                probs = self.forward_fn([b["text"] for b in batch], seq_len)
                for b, p in zip(batch, probs):
                    b["probs"] = p
                    b["event"].set()
            except Exception as exc:  # 单批失败只影响本批请求
                for b in batch:
                    b["error"] = exc
                    b["event"].set()


class AdClassifierPredictor:
    """实现与 LayaAdapter 相同的 Predictor 协议，后端为 laya-ad 二分类模型。"""

    def __init__(self, model_dir: Path, device: str | None):
        threads = int(os.environ.get("LAYA_CPU_THREADS",
                                     str(max(1, (os.cpu_count() or 2) // 2))))
        if torch.cuda.is_available() and (device or "").startswith("cuda"):
            self.device = "cuda"
        else:
            self.device = device or "cpu"
        torch.set_num_threads(threads if self.device == "cpu" else max(4, threads // 2))
        self.model_dir = Path(model_dir)
        self.model, self.tokenizer = load_ad_model(self.model_dir, self.device)
        self.model_id, self.precision = model_identity(self.model_dir)
        self.accepted_ids = _accepted_ids(self.model_id, self.precision)
        id2label = {int(k): str(v).casefold() for k, v in self.model.config.id2label.items()}
        self.ad_id = next((i for i, name in id2label.items() if name == "ad"), None)
        self.clean_id = next((i for i, name in id2label.items() if name == "normal"), None)
        if self.ad_id is None or self.clean_id is None:
            raise RuntimeError(f"模型 id2label 缺少 ad/normal: {id2label}")
        self.kind_labels = [str(x) for x in getattr(self.model, "kind_labels", [])]
        self.kind_index = {label.casefold(): i for i, label in enumerate(self.kind_labels)}
        self.scope_labels = [str(x) for x in getattr(self.model, "scope_labels", [])]
        self.scope_index = {label.casefold(): i for i, label in enumerate(self.scope_labels)}
        self.severity_labels = [str(x) for x in getattr(self.model, "severity_labels", [])]
        self.severity_index = {label.casefold(): i for i, label in enumerate(self.severity_labels)}
        self._batcher = _MicroBatcher(self._forward_batch,
                                      int(os.environ.get("LAYA_AD_BATCH", "32")),
                                      int(os.environ.get("LAYA_AD_BATCH_WINDOW_MS", "15")))

    def served_ids(self) -> list[str]:
        return self.accepted_ids

    def _forward_batch(self, texts: list[str], seq_len: int):
        enc = self.tokenizer(texts, truncation=True, max_length=seq_len,
                             padding=True, return_tensors="pt")
        enc = {k: v.to(self.device) for k, v in enc.items()}
        with torch.inference_mode():
            out = self.model(**enc)
            logits = out.logits.float()
            probs = torch.softmax(logits, dim=-1)
            kind_logits = getattr(out, "kind_logits", None)
            kind_probs = (torch.softmax(kind_logits.float(), dim=-1)
                          if kind_logits is not None else None)
            scope_logits = getattr(out, "scope_logits", None)
            scope_probs = (torch.softmax(scope_logits.float(), dim=-1)
                           if scope_logits is not None else None)
            sev_logits = getattr(out, "severity_logits", None)
            sev_probs = (torch.softmax(sev_logits.float(), dim=-1)
                         if sev_logits is not None else None)
        return [(float(p[self.ad_id]), float(p[self.clean_id]),
                 kind_probs[i].cpu().tolist() if kind_probs is not None else None,
                 scope_probs[i].cpu().tolist() if scope_probs is not None else None,
                 sev_probs[i].cpu().tolist() if sev_probs is not None else None)
                for i, p in enumerate(probs)]

    def predict(self, state: Any, questions: dict[str, Any], model: str | None = None,
                *, max_len: int | None = None) -> dict[str, Any]:
        text = _extract_text(state)
        seq_len = min(int(max_len) if max_len else 512, 512)
        p_ad, p_clean, kind_probs, scope_probs, sev_probs = self._batcher.submit(text, seq_len)
        pred_label = "ad" if p_ad >= p_clean else "clean"

        answers: dict[str, Any] = {}
        for qid, q in questions.items():
            qtype = q.get("type")
            labels = _labels(q.get("criteria"))
            if qtype == "choice" and labels and {l.casefold() for l in labels} == _AD_LABELS:
                prob_map = {}
                for label in labels:
                    prob_map[label] = round(p_ad if label.casefold() == "ad" else p_clean, 4)
                answers[qid] = {
                    "type": "choice",
                    "choice": pred_label if pred_label in [l.casefold() for l in labels] else labels[0],
                    "probabilities": prob_map,
                    "confidence": round(max(p_ad, p_clean), 4),
                    "answer_confidence": round(max(p_ad, p_clean), 4),
                    "unsupported": False,
                }
            elif (qtype == "choice" and labels and kind_probs is not None
                  and all(l.casefold() in self.kind_index or l.casefold() == "none"
                          for l in labels)):
                raw = {}
                for label in labels:
                    key = label.casefold()
                    if key == "none":
                        raw[label] = p_clean
                    elif key in self.kind_index:
                        raw[label] = p_ad * kind_probs[self.kind_index[key]]
                    else:
                        raw = None
                        break
                if raw is None:
                    answers[qid] = self._unsupported(q)
                else:
                    total = sum(raw.values()) or 1.0
                    prob_map = {label: round(v / total, 4) for label, v in raw.items()}
                    best = max(raw.items(), key=lambda x: x[1])[0]
                    answers[qid] = {
                        "type": "choice",
                        "choice": best,
                        "probabilities": prob_map,
                        "confidence": round(max(raw.values()) / total, 4),
                        "answer_confidence": round(max(raw.values()) / total, 4),
                        "unsupported": False,
                    }
            elif (qtype == "choice" and labels and scope_probs is not None
                  and all(l.casefold() in self.scope_index for l in labels)):
                if p_ad >= p_clean:
                    vals = [scope_probs[self.scope_index[l.casefold()]] for l in labels]
                else:  # 非广告按 teacher 口径默认 message
                    vals = [1.0 if l.casefold() == "message" else 0.0 for l in labels]
                total = sum(vals) or 1.0
                prob_map = {label: round(v / total, 4) for label, v in zip(labels, vals)}
                best = max(zip(labels, vals), key=lambda x: x[1])[0]
                answers[qid] = {
                    "type": "choice",
                    "choice": best,
                    "probabilities": prob_map,
                    "confidence": round(max(vals) / total, 4),
                    "answer_confidence": round(max(vals) / total, 4),
                    "unsupported": False,
                }
            elif (qtype == "score" and sev_probs is not None and len(labels) == 4
                  and all(l.casefold() in self.severity_index for l in labels)):
                if p_ad >= p_clean:
                    vals = [sev_probs[self.severity_index[l.casefold()]] for l in labels]
                else:
                    vals = [1.0, 0.0, 0.0, 0.0]
                total = sum(vals) or 1.0
                probs = {str(i): round(v / total, 4) for i, v in enumerate(vals)}
                score = sum(i * (v / total) for i, v in enumerate(vals))
                answers[qid] = {
                    "type": "score",
                    "score": round(score, 4),
                    "legend": {str(i): str(label) for i, label in enumerate(labels)},
                    "probabilities": probs,
                    "confidence": round(max(vals) / total, 4),
                    "answer_confidence": round(max(vals) / total, 4),
                    "unsupported": False,
                }
            elif qtype == "noul" and _AD_RE.search(str(q.get("instructions", ""))):
                # noul 约定：true = 是广告（仅当题目语义与广告相关时才有意义）
                answers[qid] = {
                    "type": "noul",
                    "noul": round(p_ad, 4),
                    "confidence": round(max(p_ad, 1.0 - p_ad), 4),
                    "answer_confidence": round(max(p_ad, 1.0 - p_ad), 4),
                    "unsupported": False,
                }
            else:
                answers[qid] = self._unsupported(q)

        usage = {"input_tokens": self._count_tokens(text, seq_len) * len(questions),
                 "output_tokens": len(questions)}
        return {"answers": answers, "model": self.model_id, "usage": usage}

    def _count_tokens(self, text: str, seq_len: int) -> int:
        enc = self.tokenizer(text, truncation=True, max_length=seq_len)
        return int(len(enc["input_ids"]))

    @staticmethod
    def _unsupported(q: dict[str, Any]) -> dict[str, Any]:
        """无信息中性答案：均匀概率 + 低置信 + unsupported 标记，绝不冒充判断。"""
        qtype = q.get("type")
        if qtype == "choice":
            labels = _labels(q.get("criteria")) or ["unknown"]
            n = len(labels)
            return {"type": "choice", "choice": labels[0],
                    "probabilities": {l: round(1.0 / n, 4) for l in labels},
                    "confidence": round(1.0 / n, 4), "answer_confidence": round(1.0 / n, 4),
                    "unsupported": True}
        if qtype == "score":
            n = len(q.get("criteria") or [])
            uniform = 1.0 / max(1, n)
            return {"type": "score", "score": round(sum(i * uniform for i in range(n)), 4),
                    "legend": {str(i): str(c) for i, c in enumerate(q.get("criteria") or [])},
                    "probabilities": {str(i): round(uniform, 4) for i in range(n)},
                    "confidence": round(uniform, 4), "answer_confidence": round(uniform, 4),
                    "unsupported": True}
        # noul（与广告无关的是非题）：0.5 = 无信息
        return {"type": "noul", "noul": 0.5, "confidence": 0.5, "answer_confidence": 0.5,
                "unsupported": True}


class MultiPrecisionPredictor:
    """多精度镜像的 Predictor：LAYA_MODEL_DIR 下每个子目录一个精度变体
    （各自含 MODEL_ID.txt / PRECISION.txt / model.safetensors），
    按请求的 model id 懒加载切换；auto/laya/laya-<版本> 指向 fp8 变体。"""

    def __init__(self, model_dir: Path, device: str | None):
        self.root = Path(model_dir)
        self.variants: dict[str, tuple[Path, str]] = {}
        self.default_id: str | None = None
        for d in sorted(self.root.iterdir()):
            if not d.is_dir() or not (d / "model.safetensors").is_file():
                continue
            mid, prec = model_identity(d)
            self.variants[mid] = (d, prec)
            if prec == "fp8":
                self.default_id = mid
        if not self.variants:
            raise RuntimeError(f"多模型目录为空: {self.root}")
        if self.default_id is None:
            self.default_id = next(iter(self.variants))
        self.version = _version_of(self.default_id)
        self.accepted_ids = ["auto", "laya"] \
            + ([f"laya-{self.version}"] if self.version else []) \
            + sorted(self.variants)
        self.device = device or "cpu"
        self._cache: dict[str, AdClassifierPredictor] = {}

    def resolve(self, model: str | None) -> "AdClassifierPredictor":
        if model in (None, "auto", "laya") or (self.version and model == f"laya-{self.version}"):
            mid = self.default_id
        elif model in self.variants:
            mid = model
        else:
            raise ValueError(f"unknown model id: {model}")
        if mid not in self._cache:
            d, _ = self.variants[mid]
            self._cache[mid] = AdClassifierPredictor(d, self.device)
        return self._cache[mid]

    def refresh(self) -> None:
        """重新扫描模型目录（新下载的变体上线）。"""
        for d in sorted(self.root.iterdir()):
            if d.is_dir() and (d / "MODEL_ID.txt").is_file() and (d / "model.safetensors").is_file():
                mid, prec = model_identity(d)
                if mid not in self.variants:
                    self.variants[mid] = (d, prec)
                    if prec == "fp8" and self.default_id is None:
                        self.default_id = mid
        self.version = _version_of(self.default_id) or self.version
        self.accepted_ids = ["auto", "laya"] \
            + ([f"laya-{self.version}"] if self.version else []) \
            + sorted(self.variants)

    def served_ids(self) -> list[str]:
        self.refresh()
        return self.accepted_ids

    def predict(self, state: Any, questions: dict[str, Any], model: str | None = None,
                *, max_len: int | None = None) -> dict[str, Any]:
        predictor = self.resolve(model)
        return predictor.predict(state, questions, max_len=max_len)


def served_model_ids(model_dir: Path) -> list[str]:
    """create_app 启动时确定对外可用的 model id 列表（单模型 / 多模型目录皆可）。"""
    model_dir = Path(model_dir)
    if (model_dir / "MODEL_ID.txt").is_file():
        mid, prec = model_identity(model_dir)
        return _accepted_ids(mid, prec)
    ids = ["auto", "laya"]
    default_version = None
    seen: list[str] = []
    for d in sorted(model_dir.iterdir()):
        if d.is_dir() and (d / "MODEL_ID.txt").is_file():
            mid, prec = model_identity(d)
            if mid not in seen:
                seen.append(mid)
            if prec == "fp8":
                v = _version_of(mid)
                if v:
                    default_version = v
                    ids.append(f"laya-{v}")
    return ids + seen