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

只依赖 Python 标准库 + nudenet 包，模型在进程里只加载一次。
收到 /detect 就把图片喂给 NudeDetector，返回结果 JSON。

模型加载优先级（模型放在 <程序运行目录>/lib/320n.onnx）：
  1. 本地已有该文件 → 直接加载（离线可用，适合随程序分发）；
  2. 本地没有 → 自动下载到 lib/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 json
import os
import sys
import tempfile
import urllib.request
from http.server import BaseHTTPRequestHandler, HTTPServer

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 = "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


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 _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 目录（仅第一次需要网络）
    try:
        os.makedirs(model_dir, exist_ok=True)
        print("=== 本地未找到有效模型, 正在下载: %s" % local, flush=True)
        urllib.request.urlretrieve(MODEL_URL, 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():
    global _detector
    if _detector is None:
        mp = _ensure_model(MODEL_DIR)
        try:
            _detector = NudeDetector(model_path=mp)
        except Exception as exc:  # noqa: BLE001 - 指定路径加载失败就回退包自带模型，别崩
            print("=== 本地模型加载失败(%s), 回退用包自带模型" % exc, flush=True)
            if mp:
                try:
                    os.remove(mp)
                except OSError:
                    pass
            _detector = NudeDetector()
    return _detector


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


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": "缺少图片路径"})
            data = _read_bytes(path)
            detections = _detect_bytes(get_detector(), data, path)
        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


if __name__ == "__main__":
    port = int(sys.argv[1]) if len(sys.argv) > 1 else 0
    # 启动前先把模型准备好（本地加载 / 下载 / 回退），/health 就绪即代表模型可用，
    # 这样 Go 侧 waitReady 等到的就是"模型已就绪"的服务。
    get_detector()
    HTTPServer(("127.0.0.1", port), Handler).serve_forever()