package com.cmvr.test.flow.context;

import com.alibaba.fastjson2.JSONObject;
import com.cmvr.test.enums.TaskStatusEnum;
import com.cmvr.test.enums.TerminalStatusEnum;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 任务上下文管理器
 */
@Slf4j
@Component
public class TaskContextManager {

    /**
     * 实例 ID → 任务上下文
     */
    @Getter
    private final Map<String, TaskContext> instContextMap = new ConcurrentHashMap<>();

    /**
     * 终端 ID → 实例 ID（标识该终端被哪个任务占用）
     */
    private final Map<String, String> terminalInstMap = new ConcurrentHashMap<>();

    /**
     * 注册任务上下文（任务启动时调用）
     */
    public void register(TaskContext context) {
        instContextMap.put(context.getInstId(), context);
        if (context.getTerminalId() != null) {
            terminalInstMap.put(context.getTerminalId(), context.getInstId());
        }
        context.setTerminalStatus(TerminalStatusEnum.RUNNING);
        log.info("注册任务上下文：instId={}, terminalId={}", context.getInstId(), context.getTerminalId());
    }

    /**
     * 注销任务上下文（任务结束 / 终止后调用）
     */
    public void unregister(String instId) {
        TaskContext ctx = instContextMap.remove(instId);
        if (ctx != null) {
            ctx.setTerminalStatus(TerminalStatusEnum.READY);
            ctx.setTerminalLocked(false);
            if (ctx.getTerminalId() != null) {
                terminalInstMap.remove(ctx.getTerminalId());
            }
            log.info("注销任务上下文：instId={}, terminalId={}", instId, ctx.getTerminalId());
        }
    }

    /**
     * 获取任务上下文
     */
    public TaskContext get(String instId) {
        return instContextMap.get(instId);
    }

    /**
     * 获取任务状态
     */
    public TaskStatusEnum getStatus(String instId) {
        return instContextMap.get(instId).getStatus();
    }

    /**
     * 判断某个终端是否已被占用（用于任务启动前判断）
     */
    public boolean isTerminalLocked(String terminalId) {
        return terminalInstMap.containsKey(terminalId);
    }

    /**
     * 判断终端是否被其他任务占用（用于任务恢复时判断）
     */
    public boolean isTerminalOccupied(String terminalId, String excludingInstId) {
        return terminalInstMap.containsKey(terminalId)
                && !terminalInstMap.get(terminalId).equals(excludingInstId);
    }

    /**
     * 根据终端ID获取正在执行的实例ID
     */
    public String getInstIdByTerminal(String terminalId) {
        return terminalInstMap.get(terminalId);
    }

    /**
     * 清除所有上下文
     */
    public void clearAll() {
        instContextMap.clear();
        terminalInstMap.clear();
        log.info("任务上下文已清空");
    }

    /**
     * 存储某个节点的输出参数
     */
    public void setNodeOutParams(String instId, String nodeId, List<Integer> iterations, JSONObject params) {
        TaskContext context = instContextMap.get(instId);
        if (context != null) {
            context.setNodeOutput(nodeId, iterations, params);
        }
    }

    /**
     * 根据 nodeId 和 iterations 获取节点的输出参数
     */
    public JSONObject getNodeOutParams(String instId, String nodeId, List<Integer> iterations) {
        TaskContext context = instContextMap.get(instId);
        if (context != null) {
            return context.getNodeOutput(nodeId, iterations);
        }
        return null;
    }
}
