// Package nsfw 负责"色情图片识别"：本机 Python + NudeNet 起一个本地 HTTP 服务，
// Go 侧把图片路径发给它，拿到"是否色情"和最高敏感分。
//
// 大白话理解：NudeNet 是开源里口碑不错的离线鉴黄模型（python 库，自带模型）。
// 我们不想让 Go 里掺一堆深度学习依赖，所以让 python 当"外援"：
//   - 扫描开始时拉起一个常驻的 python 小服务（模型只加载一次）；
//   - 探测图片时逐张把路径 POST 给它；
//   - 扫描结束把服务停掉、清理临时脚本。
//
// 机器上没有 python 或没装 nudenet 时，Start 会报错，调用方跳过识别、
// 走默认整理流程——这个功能是纯附加的，不影响其它步骤。
package nsfw

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"time"

	_ "embed"
)

// 内嵌的 python 服务脚本：用标准库起 HTTP 服务，负责加载模型、接收路径、返回结果。
// 不依赖 flask/fastapi，省去额外装包。
//
//go:embed server.py
var serverScript []byte

// 敏感类集合：检测到这些部位"暴露"且分数达标才算色情。
// 默认只取最明确的四类（生殖器/肛/女性乳房），避免把泳装、露背照误判。
var sensitiveClasses = map[string]bool{
	"FEMALE_GENITALIA_EXPOSED": true,
	"MALE_GENITALIA_EXPOSED":   true,
	"ANUS_EXPOSED":             true,
	"FEMALE_BREAST_EXPOSED":    true,
}

// detectorStartTimeout 等待 python 服务就绪（模型加载）的超时时间。
// 首次加载模型可能要十几秒，给足余量。
const detectorStartTimeout = 60 * time.Second

// requestTimeout 单张图片检测的超时时间。检测很快（毫秒级），
// 万一卡住（模型没加载好/图太大）不至于拖死整个扫描。
const requestTimeout = 30 * time.Second

// Detector 封装一个本地 python 服务的客户端。
// 多个探测协程会并发调 Detect，底层 HTTP 客户端线程安全；
// 服务端单线程处理，请求自己排队，不用额外加锁。
type Detector struct {
	url       string  // 本地服务地址，如 http://127.0.0.1:54321
	threshold float64 // 判定阈值 0~1
	client    *http.Client
	cmd       *exec.Cmd  // python 进程句柄（Stop 时杀掉）
	scriptDir string     // 临时脚本所在目录（Stop 时清理）
	mu        sync.Mutex // 保护下面的关闭状态
	stopped   bool       // 已经 Stop 过了
	cancel    context.CancelFunc
	done      chan struct{} // 进程退出时关闭（Wait 协程负责）
}

// Start 拉起本地 python 检测服务并等它就绪。
//   - pythonPath：python 可执行文件路径；为空时在 PATH 里找 python/python3。
//   - threshold：判定阈值 0~1。
//
// 找不到 python、装不了 nudenet、或服务起不来时返回错误（调用方跳过识别即可）。
func Start(pythonPath string, threshold float64) (*Detector, error) {
	if threshold <= 0 || threshold >= 1 {
		threshold = 0.5 // 阈值传得不合理就用默认
	}
	// 找 python
	py := resolvePython(pythonPath)
	if py == "" {
		return nil, fmt.Errorf("未找到 python，敏感识别不可用")
	}

	// 先把内嵌的 server.py 写到临时目录
	tmpDir, err := os.MkdirTemp("", "kmp-nsfw-*")
	if err != nil {
		return nil, fmt.Errorf("创建临时目录失败: %v", err)
	}
	scriptPath := filepath.Join(tmpDir, "server.py")
	if err := os.WriteFile(scriptPath, serverScript, 0644); err != nil {
		_ = os.RemoveAll(tmpDir)
		return nil, fmt.Errorf("写入检测脚本失败: %v", err)
	}

	// 找一个空闲端口
	port, err := freePort()
	if err != nil {
		_ = os.RemoveAll(tmpDir)
		return nil, fmt.Errorf("找空闲端口失败: %v", err)
	}

	// 拉起服务。标准输出/错误接到程序自身的输出，方便看 python 报错。
	ctx, cancel := context.WithCancel(context.Background())
	cmd := exec.CommandContext(ctx, py, scriptPath, fmt.Sprintf("%d", port))
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Start(); err != nil {
		cancel()
		_ = os.RemoveAll(tmpDir)
		return nil, fmt.Errorf("启动检测服务失败: %v", err)
	}

	done := make(chan struct{})
	d := &Detector{
		url:       fmt.Sprintf("http://127.0.0.1:%d", port),
		threshold: threshold,
		client:    &http.Client{Timeout: requestTimeout},
		cmd:       cmd,
		scriptDir: tmpDir,
		cancel:    cancel,
		done:      done,
	}
	// 后台等进程退出：ProcessState 要 Wait 之后才有值，这里用 done 通知，
	// 这样 waitReady 能及时发现"python 没装 nudenet 直接退掉"的情况。
	go func() {
		_ = cmd.Wait()
		close(done)
	}()

	// 等服务就绪（模型加载完成）。就绪前先探测 /health。
	if err := d.waitReady(); err != nil {
		d.Stop()
		return nil, err
	}
	return d, nil
}

// waitReady 轮询本地服务，直到能正常应答。
func (d *Detector) waitReady() error {
	deadline := time.Now().Add(detectorStartTimeout)
	for time.Now().Before(deadline) {
		// 服务进程提前退了（多半是 python 没装 nudenet）直接报错
		select {
		case <-d.done:
			return fmt.Errorf("检测服务提前退出，请确认已安装 nudenet (pip install nudenet)")
		default:
		}
		req, err := http.NewRequest(http.MethodGet, d.url+"/health", nil)
		if err != nil {
			return err
		}
		resp, err := d.client.Do(req)
		if err == nil {
			_, _ = io.Copy(io.Discard, resp.Body)
			resp.Body.Close()
			if resp.StatusCode == http.StatusOK {
				return nil
			}
		}
		time.Sleep(500 * time.Millisecond)
	}
	return fmt.Errorf("检测服务在 %v 内未就绪", detectorStartTimeout)
}

// Detect 检测一张图片是否色情。返回：
//   - score：命中的敏感类里最高分数（没有命中为 0）
//   - nsfw：是否判定为色情（有敏感类且分数 >= 阈值）
//   - err：检测失败（服务挂了/图读不出等），由调用方兜底按未命中处理
func (d *Detector) Detect(path string) (score float64, nsfw bool, err error) {
	body, _ := json.Marshal(map[string]interface{}{
		"path":      path,
		"threshold": d.threshold,
	})
	resp, err := d.client.Post(d.url+"/detect", "application/json", bytes.NewReader(body))
	if err != nil {
		return 0, false, err
	}
	defer resp.Body.Close()
	data, err := io.ReadAll(resp.Body)
	if err != nil {
		return 0, false, err
	}
	var out struct {
		Ok    bool    `json:"ok"`
		Error string  `json:"error"`
		Nsfw  bool    `json:"nsfw"`
		Score float64 `json:"score"`
	}
	if err := json.Unmarshal(data, &out); err != nil {
		return 0, false, fmt.Errorf("解析检测结果失败: %v", err)
	}
	if !out.Ok {
		return 0, false, fmt.Errorf("检测服务返回错误: %s", out.Error)
	}
	return out.Score, out.Nsfw, nil
}

// Stop 停掉检测服务并清理临时文件。可以安全地重复调用。
func (d *Detector) Stop() {
	d.mu.Lock()
	if d.stopped {
		d.mu.Unlock()
		return
	}
	d.stopped = true
	d.mu.Unlock()

	if d.cancel != nil {
		d.cancel()
	}
	if d.cmd != nil && d.cmd.Process != nil {
		// 进程活着才 kill；kill 失败也无所谓，服务很快就没了
		_ = d.cmd.Process.Kill()
	}
	if d.done != nil {
		// 等 Wait 协程回收进程，避免僵尸进程；进程已死则立刻返回
		<-d.done
	}
	if d.scriptDir != "" {
		_ = os.RemoveAll(d.scriptDir)
	}
}

// resolvePython 在配置值 / PATH 里找可用的 python 解释器。
// Windows 上常见 python / python3 / py；Linux/mac 常见 python3 / python。
func resolvePython(configured string) string {
	candidates := []string{}
	if strings.TrimSpace(configured) != "" {
		candidates = append(candidates, strings.TrimSpace(configured))
	}
	if os.Getenv("KM_NSFW_PYTHON") != "" {
		candidates = append(candidates, os.Getenv("KM_NSFW_PYTHON"))
	}
	candidates = append(candidates, "python", "python3", "py")
	for _, c := range candidates {
		if p, err := exec.LookPath(c); err == nil {
			return p
		}
		// 配置给的是绝对路径：直接看文件在不在
		if strings.ContainsAny(c, `\/`) {
			if _, err := os.Stat(c); err == nil {
				return c
			}
		}
	}
	return ""
}

// freePort 找一个当前空闲的 TCP 端口（挑到立刻释放，竞态可接受）。
func freePort() (int, error) {
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		return 0, err
	}
	defer ln.Close()
	return ln.Addr().(*net.TCPAddr).Port, nil
}