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

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 MesProBagCodeGenerateTaskRegistry {

    @Resource
    private MesProBagCodeGenerateProperties properties;

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

    /**
     * 并发许可：约束后台生成任务数
     */
    private Semaphore permits;

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

    /**
     * 尝试占用一个运行名额
     */
    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;
    }

    /**
     * 获取任务的取消标记；任务不属于本进程时返回 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);
    }

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

}