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

只依赖 Python 标准库 + nudenet 包。
onnxruntime 的 InferenceSession 不是线程安全的，同一个实例并发推理会崩，
所以这里用一个固定大小的推理线程池：每个 worker 线程各持一份独立模型实例，
N 个请求可以真·并行推理。收到 /detect 就把图片喂给 NudeDetector，返回结果 JSON。

模型加载优先级（模型放在 <程序运行目录>/lib/nwfs/320n.onnx）：
  1. 本地已有该文件 → 直接加载（离线可用，适合随程序分发）；
  2. 本地没有 → 自动下载到 lib/nwfs/320n.onnx 再加载（仅第一次需要网络）；
  3. 下载失败 → 回退用 nudenet 包自带的模型。
下载地址是 NudeNet 官方 release 的 320n.onnx（约 11.6MB）。

Windows 超长路径说明：OpenCV 的 imread 走老式 C API，路径超过 260 字符就
打不开（报 can't open/read file）。所以这里改用 Python 的 open() 读字节，
再把字节直接交给 NudeNet（v3 的 detect 支持 bytes）——从内存解码，彻底
绕开路径限制，同时兼容老版 v2（只收路径时落到临时短路径文件再试）。
"""
import concurrent.futures
import json
import os
import shutil
import sys
import tempfile
import threading
import urllib.request
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

from nudenet import NudeDetector

# 只报 OpenCV 致命错误，别让 imread 的 WARN 刷屏（也避免误导排查）
try:
    import cv2

    cv2.setLogLevel(8)  # LOG_LEVEL_ERROR
except Exception:
    pass

# 只有这些"暴露部位"才算色情；其它如肚子/腋窝/脚踝暴露不判。
SENSITIVE = {
    "FEMALE_GENITALIA_EXPOSED",
    "MALE_GENITALIA_EXPOSED",
    "ANUS_EXPOSED",
    "FEMALE_BREAST_EXPOSED",
}

# 模型文件与官方下载地址
MODEL_FILENAME = "nwfs/320n.onnx"
MODEL_URL = "https://github.com/notAI-tech/NudeNet/releases/download/v3.4-weights/320n.onnx"

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

_detector = None

# 推理并发数：onnxruntime 的 InferenceSession.run() 不是线程安全的，
# 同一个实例并发推理会崩，所以每个 worker 线程各持有一个独立实例（NudeDetector），
# N 个 worker 即可真·并行推理。默认 4，可用环境变量 KM_NSFW_WORKERS 调整。
INFER_WORKERS = int(os.environ.get("KM_NSFW_WORKERS", "4") or "4")
if INFER_WORKERS < 1:
    INFER_WORKERS = 1

_executor = None                 # 推理线程池（固定 N 个 worker，模型各持一份）
_thread_local = threading.local()  # 每个 worker 线程自己的检测器
_MODEL_PATH = None               # 启动时解析好的模型路径（None 表示用包自带模型）


def _looks_like_model(path):
    """粗略校验文件是不是真 ONNX 模型：开头不是 HTML、大小不小于 1MB（320n 模型约 11.6MB）。

    防止 GitHub 被墙时把"登录网页/占位页"当成模型存下来，下次加载直接报错。
    """
    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_model(local):
    """带超时地把模型下载到 local。urlretrieve 没有超时，网络挂起会卡死整个服务。"""
    req = urllib.request.Request(MODEL_URL)
    with urllib.request.urlopen(req, timeout=30) as resp, open(local, "wb") as f:
        shutil.copyfileobj(resp, f)


def _ensure_model(model_dir):
    """返回要加载的模型路径；返回 None 表示用 nudenet 包自带的模型。"""
    if not model_dir:
        return None
    local = os.path.join(model_dir, MODEL_FILENAME)
    if os.path.isfile(local) and _looks_like_model(local):
        print("=== 使用本地模型: %s" % local, flush=True)
        return local
    # 本地没有（或文件无效）→ 下载到 lib 目录（仅第一次需要网络，带 30 秒超时）
    try:
        os.makedirs(model_dir, exist_ok=True)
        print("=== 本地未找到有效模型, 正在下载: %s" % local, flush=True)
        _download_model(local)
        if _looks_like_model(local):
            print("=== 模型下载完成: %s" % local, flush=True)
            return local
        # 下载回来的是网页/占位文件，不是真模型：删掉免得下次继续踩坑
        print("=== 下载的文件无效(可能被墙返回了网页), 删除并回退用包自带模型", flush=True)
        try:
            os.remove(local)
        except OSError:
            pass
    except Exception as exc:  # noqa: BLE001 - 下载失败就回退，别让服务起不来
        print("=== 模型下载失败, 回退用包自带模型: %s" % exc, flush=True)
    return None


def get_detector():
    """确保模型已就绪（下载/加载），返回一个检测器实例。

    启动时调用一次：把模型下载/解析好，/health 就绪即代表模型可用。
    真正推理用的独立实例由 _new_detector() 按线程各建一份，共享这份模型路径。
    """
    global _detector, _MODEL_PATH
    if _detector is None:
        _MODEL_PATH = _ensure_model(MODEL_DIR)
        _detector = _new_detector()
    return _detector


def _new_detector():
    """创建一份独立检测器实例（每个推理 worker 各持一份，可并行推理）。"""
    mp = _MODEL_PATH
    try:
        return NudeDetector(model_path=mp) if mp else NudeDetector()
    except Exception as exc:  # noqa: BLE001 - 指定路径加载失败就回退包自带模型，别崩
        print("=== 本地模型加载失败(%s), 回退用包自带模型" % exc, flush=True)
        if mp:
            try:
                os.remove(mp)
            except OSError:
                pass
        return NudeDetector()


def _longpath(path):
    """Windows 下给绝对路径加 \\\\?\\ 前缀，让 Python 原生 open 也能处理超长路径。"""
    if os.name != "nt":
        return path
    if path.startswith("\\\\?\\"):
        return path
    if path.startswith("\\\\"):
        # UNC 网络路径：\\server\share → \\?\UNC\server\share
        return "\\\\?\\UNC\\" + path.lstrip("\\")
    if os.path.isabs(path):
        return "\\\\?\\" + os.path.normpath(path)
    return path


def _read_bytes(path):
    """用 Python 读文件字节（不经过 OpenCV），读不了就抛异常由上层兜底。"""
    lp = _longpath(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 _detect_bytes(det, data, path):
    """把图片字节喂给 NudeNet。v3 直接吃 bytes；老版 v2 只收路径时落到临时短路径重试。

    用线程自己的检测器实例推理，各实例互不干扰，可以真·并行。
    """
    try:
        return det.detect(data)
    except (TypeError, AttributeError):
        tmp = None
        try:
            fd, tmp = tempfile.mkstemp(suffix=os.path.splitext(path)[1] or ".jpg")
            with os.fdopen(fd, "wb") as f:
                f.write(data)
            return det.detect(tmp)
        finally:
            if tmp:
                try:
                    os.unlink(tmp)
                except OSError:
                    pass


def _infer(path):
    """整个推理任务（读文件 + 检测），提交到固定线程池执行。

    读文件在池外并行进行；推理用线程各自的独立检测器，N 个 worker 真·并行。
    """
    data = _read_bytes(path)
    det = getattr(_thread_local, "detector", None)
    if det is None:
        det = _new_detector()
        _thread_local.detector = det
    return _detect_bytes(det, data, path)


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": "缺少图片路径"})
            # 丢给推理线程池：读文件并行、推理由 N 个 worker 并行处理
            detections = _executor.submit(_infer, path).result()
        except Exception as exc:  # noqa: BLE001 - 服务端任何异常都要回给 Go，别让进程崩
            return self._send({"ok": False, "error": str(exc)})

        max_score = 0.0
        nsfw = False
        for det in detections or []:
            cls = det.get("class", "")
            sc = float(det.get("score", 0) or 0)
            if cls in SENSITIVE:
                if sc > max_score:
                    max_score = sc
                if sc >= threshold:
                    nsfw = True
        self._send({"ok": True, "nsfw": nsfw, "score": max_score, "detections": detections})

    def log_message(self, *args):  # 不刷屏
        pass


# 多线程接收连接 + 大监听队列：Go 侧会并发发请求（worker 数可到 8），
# 单线程 HTTPServer 默认 backlog 只有 5，并发一高连接就被内核拒绝（报 10061 连接拒绝）。
# ThreadingHTTPServer 让连接即时被 accept；推理再交给固定线程池并行执行
# （每个 worker 一份独立模型实例，onnxruntime 同一实例并发推理会崩）。
class Server(ThreadingHTTPServer):
    request_queue_size = 128


if __name__ == "__main__":
    port = int(sys.argv[1]) if len(sys.argv) > 1 else 0
    # 启动前先把模型准备好（本地加载 / 下载 / 回退），/health 就绪即代表模型可用，
    # 这样 Go 侧 waitReady 等到的就是"模型已就绪"的服务。
    get_detector()
    # 固定大小推理线程池：每个 worker 线程持一份独立模型实例，可并行推理
    _executor = concurrent.futures.ThreadPoolExecutor(max_workers=INFER_WORKERS)
    Server(("127.0.0.1", port), Handler).serve_forever()