# -*- coding: utf-8 -*-
"""wd-tagger 本地检测服务：由 Go 侧启动，常驻提供 /health 和 /detect。

基于 SmilingWolf/wd-vit-large-tagger-v3 模型，专门用于二次元（anime/hentai）图片分级。
输出 rating:safe / rating:questionable / rating:explicit 三个分级分数，
以及若干高置信度标签（用于辅助判断内容类型）。

只依赖 Python 标准库 + onnxruntime + Pillow。
模型与标签文件放在 <程序运行目录>/lib/wd-tagger/ 下，首次启动自动从 HuggingFace 下载。

推理并发策略与 nsfw/server.py 相同：每个 worker 线程各持一份独立 session 实例，
避免 onnxruntime 同一实例并发推理崩溃的问题。
"""
import concurrent.futures
import csv
import io
import json
import os
import shutil
import sys
import tempfile
import threading
import urllib.request
import zipfile
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

try:
    import numpy as np
except ImportError:
    print("=== 错误: 缺少 numpy，请执行 pip install numpy", flush=True)
    sys.exit(1)

try:
    import onnxruntime as ort
except ImportError:
    print("=== 错误: 缺少 onnxruntime，请执行 pip install onnxruntime", flush=True)
    sys.exit(1)

try:
    from PIL import Image
except ImportError:
    print("=== 错误: 缺少 Pillow，请执行 pip install Pillow", flush=True)
    sys.exit(1)

# 模型相关配置
MODEL_DIR_NAME = "wd-tagger"
MODEL_FILENAME = "model.onnx"
TAGS_FILENAME = "selected_tags.csv"

# HuggingFace 下载地址（wd-vit-tagger-v3，轻量版，约 340MB）
HF_MODEL_URL = "https://huggingface.co/SmilingWolf/wd-vit-tagger-v3/resolve/main/model.onnx"
HF_TAGS_URL = "https://huggingface.co/SmilingWolf/wd-vit-tagger-v3/resolve/main/selected_tags.csv"

# 模型输入尺寸
INPUT_SIZE = 448

# ImageNet 归一化参数
IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)

# 我们关注的分级标签
RATING_TAGS = {
    "rating:safe": "safe",
    "rating:questionable": "questionable",
    "rating:explicit": "explicit",
}

# 模型目录由 Go 侧作为第二个参数传入
MODEL_DIR = sys.argv[2] if len(sys.argv) > 2 else ""

# 推理并发数（ViT-Large 模型较大，默认 2 避免 CPU 争抢）
INFER_WORKERS = int(os.environ.get("KM_HENTAI_WORKERS", "2") or "2")
if INFER_WORKERS < 1:
    INFER_WORKERS = 1

_executor = None
_thread_local = threading.local()

# 标签表：启动时从 CSV 加载
_tag_names = []
_rating_indices = {}  # tag名 -> 在输出数组中的索引


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) > 1_000_000
    except OSError:
        return False


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


def _ensure_model(model_dir):
    """确保模型和标签文件就绪，返回 (model_path, tags_path) 或 (None, None) 表示用包自带。"""
    if not model_dir:
        return None, None

    wd_dir = os.path.join(model_dir, MODEL_DIR_NAME)
    model_path = os.path.join(wd_dir, MODEL_FILENAME)
    tags_path = os.path.join(wd_dir, TAGS_FILENAME)

    # 检查模型文件
    if os.path.isfile(model_path) and _looks_like_model(model_path):
        print("=== 使用本地 wd-tagger 模型: %s" % model_path, flush=True)
    else:
        # 下载模型
        try:
            os.makedirs(wd_dir, exist_ok=True)
            print("=== 本地未找到 wd-tagger 模型, 正在下载: %s" % model_path, flush=True)
            _download_file(HF_MODEL_URL, model_path)
            if _looks_like_model(model_path):
                print("=== wd-tagger 模型下载完成: %s" % model_path, flush=True)
            else:
                print("=== 下载的模型文件无效, 删除", flush=True)
                try:
                    os.remove(model_path)
                except OSError:
                    pass
                return None, None
        except Exception as exc:
            print("=== wd-tagger 模型下载失败: %s" % exc, flush=True)
            return None, None

    # 检查标签文件
    if os.path.isfile(tags_path):
        print("=== 使用本地标签文件: %s" % tags_path, flush=True)
    else:
        try:
            print("=== 下载标签文件: %s" % tags_path, flush=True)
            _download_file(HF_TAGS_URL, tags_path)
            print("=== 标签文件下载完成", flush=True)
        except Exception as exc:
            # 标签文件下载失败不影响模型运行，只是无法解析输出标签
            print("=== 警告: 标签文件下载失败，将无法解析输出标签: %s" % exc, flush=True)
            tags_path = ""  # 标记标签文件不可用，但仍返回模型路径

    return model_path, tags_path


def _load_tags(tags_path):
    """加载标签 CSV，建立 tag名 -> 索引 的映射。"""
    global _tag_names, _rating_indices
    _tag_names = []
    _rating_indices = {}

    if not tags_path or not os.path.isfile(tags_path):
        print("=== 警告: 标签文件不存在，无法解析输出", flush=True)
        return

    with open(tags_path, "r", encoding="utf-8") as f:
        reader = csv.DictReader(f)
        for row in reader:
            name = row.get("name", "")
            _tag_names.append(name)
            if name in RATING_TAGS:
                _rating_indices[name] = len(_tag_names) - 1

    print("=== 已加载 %d 个标签, 其中分级标签: %s" % (len(_tag_names), list(_rating_indices.keys())), flush=True)


def _new_session(model_path):
    """创建一份独立的 ONNX 推理 session。"""
    if model_path:
        return ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])
    # 没有本地模型，无法推理
    raise RuntimeError("wd-tagger 模型不可用")


def get_session():
    """确保模型已就绪，返回一个推理 session（启动时调用）。"""
    global _session
    if "_session" not in dir() or _session is None:
        model_path, tags_path = _ensure_model(MODEL_DIR)
        if not model_path:
            raise RuntimeError("wd-tagger 模型不可用")
        if not tags_path:
            raise RuntimeError("wd-tagger 标签文件不可用，无法解析输出。请手动下载 %s 到 %s" % (TAGS_FILENAME, os.path.join(MODEL_DIR, MODEL_DIR_NAME)))
        _load_tags(tags_path)
        _session = _new_session(model_path)
    return _session


_session = None


def _preprocess(img_bytes):
    """把图片字节预处理成模型输入的 numpy 数组。

    步骤：
    1. 解码为 PIL Image
    2. 转为 RGB（如果有 alpha 通道）
    3. 直接 resize 到 INPUT_SIZE x INPUT_SIZE（官方推荐方式）
    4. 归一化（ImageNet 参数）
    5. 转为 NHWC 格式的 float32 数组（模型期望 [batch, H, W, C]）
    """
    img = Image.open(io.BytesIO(img_bytes))
    if img.mode != "RGB":
        img = img.convert("RGB")

    # 直接 resize 到目标尺寸（wd-tagger 官方推荐）
    img = img.resize((INPUT_SIZE, INPUT_SIZE), Image.LANCZOS)

    # 转 numpy 数组并归一化
    arr = np.array(img, dtype=np.float32) / 255.0
    arr = (arr - IMAGENET_MEAN) / IMAGENET_STD

    # HWC -> NHWC (模型期望 [batch, height, width, channels])
    arr = np.expand_dims(arr, axis=0)
    return arr


def _infer_bytes(session, img_bytes):
    """用给定的 session 推理一张图片（字节），返回 (ratings_dict, top_tags_list)。"""
    inp = _preprocess(img_bytes)
    input_name = session.get_inputs()[0].name
    outputs = session.run(None, {input_name: inp})

    # 输出是 [1, num_tags] 的分数数组
    scores = outputs[0][0] if len(outputs[0].shape) == 2 else outputs[0]

    # 提取分级分数
    ratings = {}
    for tag_name, idx in _rating_indices.items():
        if idx < len(scores):
            rating_key = RATING_TAGS[tag_name]
            ratings[rating_key] = float(scores[idx])

    # 提取 top 标签（排除分级标签，取分数最高的 10 个）
    top_tags = []
    non_rating_indices = [i for i in range(len(scores)) if i not in _rating_indices.values()]
    if non_rating_indices:
        top_indices = sorted(non_rating_indices, key=lambda i: scores[i], reverse=True)[:10]
        for idx in top_indices:
            if idx < len(_tag_names) and scores[idx] > 0.5:
                top_tags.append({"tag": _tag_names[idx], "score": float(scores[idx])})

    return ratings, top_tags


def _read_bytes(path):
    """用 Python 读文件字节。"""
    if os.name == "nt":
        if path.startswith("\\\\?\\"):
            lp = path
        elif path.startswith("\\\\"):
            lp = "\\\\?\\UNC\\" + path.lstrip("\\")
        elif os.path.isabs(path):
            lp = "\\\\?\\" + os.path.normpath(path)
        else:
            lp = path
    else:
        lp = path

    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)
    return data


def _infer(path):
    """整个推理任务：读文件 + 预处理 + 推理。"""
    data = _read_bytes(path)
    session = getattr(_thread_local, "session", None)
    if session is None:
        # 每个 worker 线程各持一份独立 session
        model_path = os.path.join(MODEL_DIR, MODEL_DIR_NAME, MODEL_FILENAME) if MODEL_DIR else ""
        session = _new_session(model_path)
        _thread_local.session = session
    return _infer_bytes(session, data)


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": "缺少图片路径"})
            # 丢给推理线程池
            ratings, top_tags = _executor.submit(_infer, path).result()
        except Exception as exc:
            return self._send({"ok": False, "error": str(exc)})

        # 判断是否为 hentai：explicit 分数 >= 阈值 或 questionable 分数 >= 阈值
        explicit_score = ratings.get("explicit", 0.0)
        questionable_score = ratings.get("questionable", 0.0)
        safe_score = ratings.get("safe", 0.0)

        # hentai 判定：explicit 或 questionable 分数超过阈值
        is_hentai = (explicit_score >= threshold) or (questionable_score >= threshold)

        # 取最高的非 safe 分数作为"敏感分"
        max_score = max(explicit_score, questionable_score)

        # 确定分级标签
        rating = "safe"
        if explicit_score >= questionable_score and explicit_score >= threshold:
            rating = "explicit"
        elif questionable_score >= threshold:
            rating = "questionable"

        self._send({
            "ok": True,
            "hentai": is_hentai,
            "score": max_score,
            "rating": rating,
            "ratings": ratings,
            "topTags": 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
    # 启动前先把模型准备好
    get_session()
    # 固定大小推理线程池
    _executor = concurrent.futures.ThreadPoolExecutor(max_workers=INFER_WORKERS)
    Server(("127.0.0.1", port), Handler).serve_forever()