# -*- coding: utf-8 -*-
"""WD Tagger v3 (eva02-large) 动漫色情检测服务：由 Go 侧启动，常驻提供 /health 和 /detect。

只依赖 Python 标准库 + onnxruntime + opencv-python + numpy。
onnxruntime 的 InferenceSession 不是线程安全的，同一个实例并发推理会崩，
所以这里用一个固定大小的推理线程池：每个 worker 线程各持一份独立模型实例，
N 个请求可以真·并行推理。

模型来源：SmilingWolf/wd-eva02-large-tagger-v3 (HuggingFace)
  - model.onnx：ONNX 模型（约 1.26GB）
  - selected_tags.csv：标签映射表，前 4 行是 rating 分类：
      index 0: general（安全）
      index 1: sensitive（轻微敏感）
      index 2: questionable（可疑）
      index 3: explicit（明确色情）

模型文件放在 <程序运行目录>/lib/anime_nsfw/ 下，首次启动自动从 HuggingFace 下载。

输入预处理：
  - 图片 pad 成正方形（白边填充）
  - bicubic 缩放到 448x448
  - 归一化：(pixel/255 - 0.5) / 0.5

NSFW 判定逻辑：
  - explicit 或 questionable 的 sigmoid 分数 >= 阈值 → 判定为色情
  - score 返回 max(explicit, questionable) 的 sigmoid 分数
"""
import concurrent.futures
import json
import os
import shutil
import sys
import threading
import urllib.request
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

import cv2
import numpy as np
import onnxruntime as ort

# 模型文件与 HuggingFace 下载地址
HF_REPO = "SmilingWolf/wd-eva02-large-tagger-v3"
HF_BASE_URL = "https://huggingface.co/%s/resolve/main" % HF_REPO
MODEL_FILENAME = "model.onnx"
TAGS_FILENAME = "selected_tags.csv"

# 模型输入参数（来自 config.json）
INPUT_SIZE = 448
MEAN = np.array([0.5, 0.5, 0.5], dtype=np.float32)
STD = np.array([0.5, 0.5, 0.5], dtype=np.float32)

# Rating 标签在 selected_tags.csv 中的位置（前 4 行）
RATING_GENERAL = 0
RATING_SENSITIVE = 1
RATING_QUESTIONABLE = 2
RATING_EXPLICIT = 3

# 模型目录由 Go 侧作为第二个参数传入（程序运行目录下的 lib）
MODEL_DIR = sys.argv[2] if len(sys.argv) > 2 else ""

# 推理并发数（每个 worker 持独立模型实例）
INFER_WORKERS = int(os.environ.get("KM_ANIME_NSFW_WORKERS", "2") or "2")
if INFER_WORKERS < 1:
    INFER_WORKERS = 1

_executor = None
_thread_local = threading.local()
_MODEL_PATH = None
_TAGS = None  # numpy array of tag names (loaded once)


def _looks_like_model(path):
    """粗略校验文件是不是真 ONNX 模型。"""
    try:
        with open(path, "rb") as f:
            head = f.read(4)
        if not head or head.startswith(b"<"):
            return False
        return os.path.getsize(path) > 100_000_000  # ONNX 模型约 1.26GB
    except OSError:
        return False


def _download_file(url, local):
    """带超时地把文件下载到 local。"""
    req = urllib.request.Request(url)
    with urllib.request.urlopen(req, timeout=120) as resp, open(local, "wb") as f:
        shutil.copyfileobj(resp, f)


def _ensure_model(model_dir):
    """确保模型和标签文件都存在，返回 (model_path, tags_csv_path)。"""
    if not model_dir:
        raise RuntimeError("未指定模型目录")

    anime_dir = os.path.join(model_dir, "anime_nsfw")
    os.makedirs(anime_dir, exist_ok=True)

    model_path = os.path.join(anime_dir, MODEL_FILENAME)
    tags_path = os.path.join(anime_dir, TAGS_FILENAME)

    # 下载模型（如果本地没有或无效）
    if not os.path.isfile(model_path) or not _looks_like_model(model_path):
        url = "%s/%s" % (HF_BASE_URL, MODEL_FILENAME)
        print("=== [AnimeNSFW] 正在下载模型: %s" % url, flush=True)
        _download_file(url, model_path)
        if not _looks_like_model(model_path):
            raise RuntimeError("下载的模型文件无效")
        print("=== [AnimeNSFW] 模型下载完成: %s" % model_path, flush=True)
    else:
        print("=== [AnimeNSFW] 使用本地模型: %s" % model_path, flush=True)

    # 下载标签 CSV（如果本地没有）
    if not os.path.isfile(tags_path):
        url = "%s/%s" % (HF_BASE_URL, TAGS_FILENAME)
        print("=== [AnimeNSFW] 正在下载标签: %s" % url, flush=True)
        _download_file(url, tags_path)
        print("=== [AnimeNSFW] 标签下载完成: %s" % tags_path, flush=True)

    return model_path, tags_path


def _load_tags(tags_path):
    """加载 selected_tags.csv，返回 tag name 的 numpy array。"""
    import csv
    tags = []
    with open(tags_path, "r", encoding="utf-8") as f:
        reader = csv.DictReader(f)
        for row in reader:
            tags.append(row["name"])
    return np.array(tags)


def _new_session(model_path):
    """创建一份独立 ONNX 推理会话（每个 worker 各持一份）。"""
    sess_opts = ort.SessionOptions()
    sess_opts.inter_op_num_threads = 1
    return ort.InferenceSession(model_path, sess_options=sess_opts, providers=["CPUExecutionProvider"])


def _preprocess(data):
    """图片字节 → 模型输入 tensor。

    步骤：
      1. 解码图片（从字节）
      2. RGBA → BGR
      3. Pad 成正方形（白边 255 填充）
      4. Bicubic 缩放到 INPUT_SIZE x INPUT_SIZE
      5. 归一化：(pixel/255 - mean) / std
      6. HWC → NCHW
    """
    mat = cv2.imdecode(np.frombuffer(data, np.uint8), -1)
    if mat is None:
        raise OSError("图片解码失败")

    # RGBA → BGR
    if mat.ndim == 3 and mat.shape[2] == 4:
        mat = cv2.cvtColor(mat, cv2.COLOR_RGBA2BGR)
    elif mat.ndim == 2:
        mat = cv2.cvtColor(mat, cv2.COLOR_GRAY2BGR)

    # Pad 成正方形
    h, w = mat.shape[:2]
    max_size = max(w, h)
    x_pad = max_size - w
    y_pad = max_size - h
    mat_pad = cv2.copyMakeBorder(mat, 0, y_pad, 0, x_pad, cv2.BORDER_CONSTANT, value=(255, 255, 255))

    # Bicubic 缩放到模型输入尺寸
    resized = cv2.resize(mat_pad, (INPUT_SIZE, INPUT_SIZE), interpolation=cv2.INTER_CUBIC)

    # 归一化：BGR → RGB, uint8 → float32, (pixel/255 - mean) / std
    img = cv2.cvtColor(resized, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
    img = (img - MEAN) / STD

    # 模型输入格式 NHWC → NCHW 不需要，直接加 batch 维度
    img = np.expand_dims(img, axis=0)
    return img


def _sigmoid(x):
    """数值稳定的 sigmoid。"""
    return np.where(x >= 0, 1.0 / (1.0 + np.exp(-x)), np.exp(x) / (1.0 + np.exp(x)))


def _infer(path):
    """整个推理任务（读文件 + 预处理 + 推理 + 后处理）。"""
    # 读文件
    lp = path
    if os.name == "nt":
        if not lp.startswith("\\\\?\\"):
            if lp.startswith("\\\\"):
                lp = "\\\\?\\UNC\\" + lp.lstrip("\\")
            elif os.path.isabs(lp):
                lp = "\\\\?\\" + os.path.normpath(lp)
    if not os.path.exists(lp):
        raise OSError("文件不存在: %s" % path)
    with open(lp, "rb") as f:
        data = f.read()
    if not data:
        raise OSError("空文件: %s" % path)

    # 预处理
    input_tensor = _preprocess(data)

    # 推理（每个 worker 线程持独立 session）
    sess = getattr(_thread_local, "session", None)
    if sess is None:
        sess = _new_session(_MODEL_PATH)
        _thread_local.session = sess

    input_name = sess.get_inputs()[0].name
    outputs = sess.run(None, {input_name: input_tensor})
    logits = outputs[0][0]  # shape: (num_classes,)

    # Sigmoid 转概率
    probs = _sigmoid(logits)

    # 提取 rating 分数
    rating_general = float(probs[RATING_GENERAL])
    rating_sensitive = float(probs[RATING_SENSITIVE])
    rating_questionable = float(probs[RATING_QUESTIONABLE])
    rating_explicit = float(probs[RATING_EXPLICIT])

    # 找 top 10 一般标签（跳过前 4 个 rating）
    tag_probs = probs[4:]
    top_indices = np.argsort(tag_probs)[::-1][:10]
    top_tags = []
    for idx in top_indices:
        tag_name = str(_TAGS[idx + 4])
        tag_score = float(tag_probs[idx])
        if tag_score > 0.05:
            top_tags.append({"tag": tag_name, "score": round(tag_score, 4)})

    return {
        "ratings": {
            "general": round(rating_general, 4),
            "sensitive": round(rating_sensitive, 4),
            "questionable": round(rating_questionable, 4),
            "explicit": round(rating_explicit, 4),
        },
        "top_tags": top_tags,
    }


class Handler(BaseHTTPRequestHandler):
    def _send(self, obj, status=200):
        data = json.dumps(obj, ensure_ascii=False).encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type", "application/json; charset=utf-8")
        self.send_header("Content-Length", str(len(data)))
        self.end_headers()
        self.wfile.write(data)

    def do_GET(self):
        if self.path == "/health":
            self._send({"ok": True})
        else:
            self._send({"ok": False, "error": "not found"}, 404)

    def do_POST(self):
        if self.path != "/detect":
            return self._send({"ok": False, "error": "not found"}, 404)
        length = int(self.headers.get("Content-Length", 0))
        try:
            body = json.loads(self.rfile.read(length) or b"{}")
            path = body.get("path", "")
            threshold = float(body.get("threshold", 0.5))
            if not path:
                return self._send({"ok": False, "error": "缺少图片路径"})
            result = _executor.submit(_infer, path).result()
        except Exception as exc:
            return self._send({"ok": False, "error": str(exc)})

        # NSFW 判定：explicit 或 questionable 分数 >= 阈值，且必须高于 general。
        # WD Tagger 的 rating 是独立 sigmoid（不归一化），非动漫图四个分值往往
        # 都很接近；若 general 最高说明模型认为是安全图，只看绝对值会大量误判。
        ratings = result["ratings"]
        nsfw_raw = max(ratings["explicit"], ratings["questionable"])
        if nsfw_raw > ratings["general"]:
            nsfw_score = nsfw_raw
            nsfw = nsfw_score >= threshold
        else:
            nsfw_score = 0.0
            nsfw = False

        self._send({
            "ok": True,
            "nsfw": nsfw,
            "score": nsfw_score,
            "ratings": ratings,
            "top_tags": result["top_tags"],
        })

    def log_message(self, *args):
        pass


class Server(ThreadingHTTPServer):
    request_queue_size = 128


if __name__ == "__main__":
    port = int(sys.argv[1]) if len(sys.argv) > 1 else 0

    # 加载模型和标签
    _MODEL_PATH, tags_path = _ensure_model(MODEL_DIR)
    _TAGS = _load_tags(tags_path)
    print("=== [AnimeNSFW] 标签加载完成, 共 %d 个" % len(_TAGS), flush=True)

    # 预创建第一个 session 验证模型可用
    sess = _new_session(_MODEL_PATH)
    print("=== [AnimeNSFW] 模型加载成功, 输入: %s" % str(sess.get_inputs()[0].shape), flush=True)
    del sess

    # 启动推理线程池
    _executor = concurrent.futures.ThreadPoolExecutor(max_workers=INFER_WORKERS)
    print("=== [AnimeNSFW] 服务启动, 端口=%d, workers=%d" % (port, INFER_WORKERS), flush=True)
    Server(("127.0.0.1", port), Handler).serve_forever()