package com.cmvr.test.flow.context;

import com.cmvr.test.enums.NodeTypeEnum;
import com.cmvr.test.flow.builder.FlowNodeWrapper;
import com.cmvr.test.flow.builder.TaskKeyBuilder;
import com.cmvr.test.flow.runtime.message.TaskNodeExecuteContext;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CountDownLatch;
import java.util.stream.Collectors;

/**
 * 线程注册中心 —— 仅保存“仍在执行”的 TaskExecutionContext
 * 支持：注册、注销、精确中断、批量中断、恢复查询
 */
@Slf4j
@Component
public class TaskThreadRegistry {

    // 结构为 instId_itemId_nodeId -> TaskNodeExecuteContext
    private final Map<String, TaskNodeExecuteContext> nodeContextMap = new ConcurrentHashMap<>();

    // 结构为 instId_loopNodeId -> CountDownLatch
    private final Map<String, CountDownLatch> loopLatchMap = new ConcurrentHashMap<>();

    /**
     * 注册节点执行上下文
     */
    public void registerNodeContext(String instId, String itemId, String nodeId, TaskNodeExecuteContext context) {
        String key = TaskKeyBuilder.buildKey(instId, itemId, nodeId);
        nodeContextMap.put(key, context);
        log.info("注册节点上下文: key={}, thread={}", key, context.getThread().getName());
    }

    /**
     * 获取上下文
     */
    public TaskNodeExecuteContext getNodeContext(String instId, String itemId, String nodeId) {
        return nodeContextMap.get(TaskKeyBuilder.buildKey(instId, itemId, nodeId));
    }

    /**
     * 获取线程
     */
    public List<TaskNodeExecuteContext> getPendingNodes(String instId) {
        List<TaskNodeExecuteContext> pending = new ArrayList<>();
        for (Map.Entry<String, TaskNodeExecuteContext> entry : nodeContextMap.entrySet()) {
            if (entry.getKey().startsWith(instId)) {
                pending.add(entry.getValue());
            }
        }
        return pending;
    }

    /**
     * 中断指定节点线程
     */
    public void interruptNode(String instId, String itemId, String nodeId) {
        TaskNodeExecuteContext ctx = getNodeContext(instId, itemId, nodeId);
        if (ctx != null && ctx.getThread() != null) {
            log.info("中断节点线程: instId={}, itemId={}, nodeId={}, thread={}", instId, itemId, nodeId, ctx.getThread().getName());
            ctx.getThread().interrupt();
        }
    }

    /**
     * 中断所有节点线程
     */
    public void interruptAll(String instId) {
        for (Map.Entry<String, TaskNodeExecuteContext> entry : nodeContextMap.entrySet()) {
            if (entry.getKey().startsWith(instId)) {
                FlowNodeWrapper node = entry.getValue().getNode();
                if (node.getNodeType().equals(NodeTypeEnum.LOOP)) {
                    log.info("LOOP节点不直接中断");
                    continue;
                }
                Thread thread = entry.getValue().getThread();
                log.info("中断线程:{}", thread.getName());
                thread.interrupt();
            }
        }
//        releaseAllLatches(instId);
        log.info("-------------------------------------");
        log.info("pending:\n{}", getPendingNodes(instId).stream().map(x -> x.getNode().getNodeId()).collect(Collectors.toList()));
        log.info("-------------------------------------");
    }

    /**
     * 移除节点上下文（节点执行成功时调用）
     */
    public void unregisterNodeContext(String instId, String itemId, String nodeId) {
        String key = TaskKeyBuilder.buildKey(instId, itemId, nodeId);
        nodeContextMap.remove(key);
        log.info("移除节点上下文: key={}", key);
    }

    /**
     * 获取所有未完成节点上下文（可用于恢复）
     */
    public Map<String, TaskNodeExecuteContext> getAllContexts() {
        return nodeContextMap;
    }

    /**
     * 注册 loop 节点 latch（每次只允许一个 latch）
     */
    public void registerLatch(String instId, String loopNodeId, CountDownLatch latch) {
        String key = TaskKeyBuilder.buildLoopKey(instId, loopNodeId);
        loopLatchMap.put(key, latch);
    }

    /**
     * 获取 latch（可用于测试或回调判断）
     */
    public CountDownLatch getLatch(String instId, String loopNodeId) {
        String key = TaskKeyBuilder.buildLoopKey(instId, loopNodeId);
        boolean flag = loopLatchMap.containsKey(key);
        System.out.println("这个latch是存在的吗："+flag);
        return loopLatchMap.get(key);
    }

    /**
     * 主动释放 latch（用于暂停任务时中断 loop await）
     */
    public void releaseLatch(String instId, String loopNodeId) {
        String key = TaskKeyBuilder.buildLoopKey(instId, loopNodeId);
        CountDownLatch latch = loopLatchMap.remove(key);
        if (latch != null) {
            try {
                latch.countDown();
                log.info("释放 latch: {} -> {}", key, latch);
            } catch (Exception ignored) {
            }
        }
    }

    /**
     * 暂停任务时释放该实例下所有 loop latch
     */
    public void releaseAllLatches(String instId) {
        List<String> keysToRelease = loopLatchMap.keySet().stream()
                .filter(key -> key.startsWith(instId + "_"))
                .collect(Collectors.toList());

        for (String key : keysToRelease) {
            CountDownLatch latch = loopLatchMap.remove(key);
            if (latch != null) {
                try {
                    latch.countDown();
                    log.info("释放 latch: {}", key);
                } catch (Exception ignored) {
                }
            }
        }
    }
}
