package cn.iocoder.yudao.module.mes.framework.qrexport;

import jakarta.annotation.PostConstruct;
import jakarta.annotation.Resource;
import org.springframework.stereotype.Component;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Semaphore;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * 二维码导出任务运行态注册表
 * <p>
 * 保存「哪些任务正在生成」与「是否已请求取消」。取消标记放内存而非数据库，是因为生成线程需要
 * 每数百张就检查一次取消，走数据库会产生大量无谓的查询；本项目为单机部署，两者同进程。
 * <p>
 * 进程重启会让所有运行态丢失，此时由清理任务把残留的「生成中」任务回收为失败，不会永久卡住。
 */
@Component
public class MesProQrExportTaskRegistry {

    @Resource
    private MesProQrExportProperties properties;

    /**
     * 正在运行的任务：任务编号 → 取消标记
     */
    private final Map<Long, AtomicBoolean> runningTasks = new ConcurrentHashMap<>();

    /**
     * 并发许可：约束的是大批量后台任务数，小批量同步导出不经过此处
     */
    private Semaphore permits;

    @PostConstruct
    public void init() {
        permits = new Semaphore(Math.max(1, properties.getMaxConcurrentTasks()));
    }

    /**
     * 尝试占用一个运行名额
     *
     * @return true 表示占用成功，可开始生成
     */
    public boolean tryAcquire() {
        return permits.tryAcquire();
    }

    /**
     * 释放运行名额（生成结束、失败、取消均需调用）
     */
    public void release() {
        permits.release();
    }

    /**
     * 当前由本进程持有的任务数（含刚建好、尚未被线程接管的）
     */
    public int getRunningCount() {
        return runningTasks.size();
    }

    /**
     * 注册任务并返回其取消标记
     * <p>
     * 必须在任务落库后、投递线程池前调用：否则从「落库」到「线程开始运行」之间存在一个窗口，
     * 清理任务会把这条尚未注册的非终态记录误判为进程重启后的残留而回收掉。
     */
    public AtomicBoolean register(Long taskId) {
        AtomicBoolean flag = new AtomicBoolean(false);
        runningTasks.put(taskId, flag);
        return flag;
    }

    /**
     * 获取任务的取消标记
     *
     * @return 取消标记；任务不属于本进程时返回 null
     */
    public AtomicBoolean get(Long taskId) {
        return runningTasks.get(taskId);
    }

    /**
     * 任务是否由本进程持有（正在生成或等待线程接管）
     */
    public boolean isRunning(Long taskId) {
        return runningTasks.containsKey(taskId);
    }

    /**
     * 注销任务
     */
    public void unregister(Long taskId) {
        runningTasks.remove(taskId);
    }

    /**
     * 请求取消：任务在跑则置位取消标记，返回 true
     *
     * @return true 表示任务正在生成、取消请求已送达；false 表示任务不在本进程运行
     */
    public boolean requestCancel(Long taskId) {
        AtomicBoolean flag = runningTasks.get(taskId);
        if (flag == null) {
            return false;
        }
        flag.set(true);
        return true;
    }

}