package com.agent.management.engine;

import com.agent.management.external.SseEventBus;
import com.agent.management.model.entity.WorkflowRunNode;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;
import java.util.function.Function;

import static org.assertj.core.api.Assertions.assertThat;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;

/**
 * {@link WorkflowLevelExecutor} dataflow 调度模型的核心场景测试。
 *
 * <p>覆盖 5 个关键场景：</p>
 * <ol>
 *   <li>菱形 DAG：dataflow 下 D 在 C 完成后立即调度，不等慢节点 B</li>
 *   <li>abort 策略：失败节点阻止下游调度</li>
 *   <li>skip 策略：失败节点仍允许下游执行</li>
 *   <li>条件分支未选中：未选分支目标节点标记 SKIPPED 并传播</li>
 *   <li>merge 节点入度归零：等所有前驱完成才调度</li>
 * </ol>
 *
 * <p>通过 mock NodeExecutor 按 nodeId 返回不同结果，真实 DagResolver / WorkflowContext，
 * 验证 sortOrder、status、是否调度等行为。</p>
 */
class WorkflowLevelExecutorTest {

    private SseEventBus eventBus;
    private NodeWorkspaceBuilder workspaceBuilder;
    private SseEmitter emitter;
    private Map<String, NodeExecutor> executorsMap;
    private Map<String, Integer> sortOrderCapture;
    private Map<String, String> statusCapture;

    @BeforeEach
    void setUp() {
        eventBus = mock(SseEventBus.class);
        workspaceBuilder = mock(NodeWorkspaceBuilder.class);
        when(workspaceBuilder.build(anyString(), any(), any(), any()))
                .thenReturn(new NodeWorkspace());
        emitter = mock(SseEmitter.class);
        executorsMap = new HashMap<>();
        sortOrderCapture = new ConcurrentHashMap<>();
        statusCapture = new ConcurrentHashMap<>();
    }

    /**
     * 辅助：构造 mock NodeExecutor，按 nodeId 返回不同结果。
     */
    private NodeExecutor mockExecutor(String type, Function<String, NodeExecutionResult> behavior) {
        NodeExecutor exec = mock(NodeExecutor.class);
        when(exec.getType()).thenReturn(type);
        when(exec.execute(anyString(), any(), any(), any())).thenAnswer(inv -> {
            String nodeId = inv.getArgument(0);
            return behavior.apply(nodeId);
        });
        return exec;
    }

    /**
     * 辅助：执行 DAG 并捕获每个节点的 sortOrder / status。
     */
    private WorkflowLevelExecutor.ExecutionOutcome runDag(String graphData) {
        DagResolver.ResolvedDag dag = DagResolver.resolve(graphData);
        WorkflowContext context = new WorkflowContext(1L, "run-1", null, new HashMap<>());
        WorkflowLevelExecutor executor = new WorkflowLevelExecutor(
                new LinkedHashMap<>(executorsMap), eventBus, workspaceBuilder);
        WorkflowLevelExecutor.ExecutionOutcome outcome = executor.executeByLevel(
                dag, context, "run-1", 1L, emitter);

        for (WorkflowRunNode node : outcome.nodeRecords) {
            sortOrderCapture.put(node.getNodeId(), node.getSortOrder());
            statusCapture.put(node.getNodeId(), node.getStatus());
        }
        return outcome;
    }

    /**
     * 辅助：构造 envelope 格式的输出（status=200 + data）。
     */
    private static Map<String, Object> envelope(String outputValue) {
        Map<String, Object> data = new LinkedHashMap<>();
        data.put("output", outputValue);
        Map<String, Object> env = new LinkedHashMap<>();
        env.put("status", 200);
        env.put("message", "ok");
        env.put("data", data);
        return env;
    }

    // ============================ 测试用例 ============================

    /**
     * 场景 1：fork chain（A→B→D 和 A→C→E），B 固定 300ms，C 立即完成。
     *
     * <p>dataflow 下 C 完成立即触发 E（不等 B 完成），因此 E.sortOrder 应小于 B.sortOrder。
     * 若为 BSP 模型，B、C 同层，E 必须等 B 完成才能开始，E.sortOrder > B.sortOrder。</p>
     */
    @Test
    void forkChain_fastBranchDownstreamTriggersImmediately() throws Exception {
        String graphData = "{"
                + "\"nodes\":["
                + "  {\"id\":\"A\",\"type\":\"llm\",\"data\":{\"label\":\"A\"}},"
                + "  {\"id\":\"B\",\"type\":\"llm\",\"data\":{\"label\":\"B\"}},"
                + "  {\"id\":\"C\",\"type\":\"llm\",\"data\":{\"label\":\"C\"}},"
                + "  {\"id\":\"D\",\"type\":\"llm\",\"data\":{\"label\":\"D\"}},"
                + "  {\"id\":\"E\",\"type\":\"llm\",\"data\":{\"label\":\"E\"}}"
                + "],\"edges\":["
                + "  {\"source\":\"A\",\"target\":\"B\"},"
                + "  {\"source\":\"A\",\"target\":\"C\"},"
                + "  {\"source\":\"B\",\"target\":\"D\"},"
                + "  {\"source\":\"C\",\"target\":\"E\"}"
                + "]}";

        NodeExecutor llm = mockExecutor("llm", nodeId -> {
            try {
                switch (nodeId) {
                    case "A":
                        return NodeExecutionResult.success("A", envelope("A-done"));
                    case "B":
                        // B 固定 300ms 延迟（模拟"慢节点"）
                        Thread.sleep(300);
                        return NodeExecutionResult.success("B", envelope("B-done"));
                    case "C":
                        return NodeExecutionResult.success("C", envelope("C-done"));
                    case "D":
                        return NodeExecutionResult.success("D", envelope("D-done"));
                    case "E":
                        return NodeExecutionResult.success("E", envelope("E-done"));
                    default:
                        return NodeExecutionResult.failed(nodeId, "未知节点");
                }
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
                return NodeExecutionResult.failed(nodeId, "中断");
            }
        });
        executorsMap.put("llm", llm);

        WorkflowLevelExecutor.ExecutionOutcome outcome = runDag(graphData);

        assertThat(outcome.failed).isFalse();
        assertThat(statusCapture).containsEntry("A", "SUCCESS")
                .containsEntry("B", "SUCCESS")
                .containsEntry("C", "SUCCESS")
                .containsEntry("D", "SUCCESS")
                .containsEntry("E", "SUCCESS");
        // 关键：E（C 的下游）在 B 之前完成（dataflow 下 C 完成立即触发 E）
        assertThat(sortOrderCapture.get("E"))
                .as("E 应在 B 之前完成（dataflow 下 C 完成立即触发 E）")
                .isLessThan(sortOrderCapture.get("B"));
        // A 是起点，sortOrder=0
        assertThat(sortOrderCapture.get("A")).isEqualTo(0);
        // 全部 5 个节点都有 sortOrder
        assertThat(sortOrderCapture).hasSize(5);
    }

    /**
     * 场景 2：A→B→C，A 失败 + failStrategy=abort。
     *
     * <p>预期：B 和 C 不被调度（executor.execute 不被调用），workflow 标记为 failed，
     * finalOutputs 为空。</p>
     */
    @Test
    void failedNodeWithAbort_blocksDownstream() {
        String graphData = "{"
                + "\"nodes\":["
                + "  {\"id\":\"A\",\"type\":\"llm\",\"data\":{\"label\":\"A\",\"failStrategy\":\"abort\"}},"
                + "  {\"id\":\"B\",\"type\":\"llm\",\"data\":{\"label\":\"B\"}},"
                + "  {\"id\":\"C\",\"type\":\"llm\",\"data\":{\"label\":\"C\"}}"
                + "],\"edges\":["
                + "  {\"source\":\"A\",\"target\":\"B\"},"
                + "  {\"source\":\"B\",\"target\":\"C\"}"
                + "]}";

        AtomicInteger bExecutions = new AtomicInteger();
        AtomicInteger cExecutions = new AtomicInteger();
        NodeExecutor llm = mockExecutor("llm", nodeId -> {
            switch (nodeId) {
                case "A":
                    return NodeExecutionResult.failed("A", "故意失败");
                case "B":
                    bExecutions.incrementAndGet();
                    return NodeExecutionResult.success("B", envelope("B-done"));
                case "C":
                    cExecutions.incrementAndGet();
                    return NodeExecutionResult.success("C", envelope("C-done"));
                default:
                    return NodeExecutionResult.failed(nodeId, "未知");
            }
        });
        executorsMap.put("llm", llm);

        WorkflowLevelExecutor.ExecutionOutcome outcome = runDag(graphData);

        assertThat(outcome.failed).isTrue();
        assertThat(outcome.errorMsg).contains("A");
        assertThat(outcome.finalOutputs).isEmpty();
        assertThat(bExecutions.get()).as("B 不应被调度").isEqualTo(0);
        assertThat(cExecutions.get()).as("C 不应被调度").isEqualTo(0);
        // B、C 应被记录为 SKIPPED
        assertThat(statusCapture).containsEntry("A", "FAILED")
                .containsEntry("B", "SKIPPED")
                .containsEntry("C", "SKIPPED");
    }

    /**
     * 场景 3：A→B→C，A 失败 + failStrategy=skip。
     *
     * <p>预期：B 和 C 仍被调度执行，workflow 不标记为 failed。</p>
     */
    @Test
    void failedNodeWithSkip_continuesWorkflow() {
        String graphData = "{"
                + "\"nodes\":["
                + "  {\"id\":\"A\",\"type\":\"llm\",\"data\":{\"label\":\"A\",\"failStrategy\":\"skip\"}},"
                + "  {\"id\":\"B\",\"type\":\"llm\",\"data\":{\"label\":\"B\"}},"
                + "  {\"id\":\"C\",\"type\":\"llm\",\"data\":{\"label\":\"C\"}}"
                + "],\"edges\":["
                + "  {\"source\":\"A\",\"target\":\"B\"},"
                + "  {\"source\":\"B\",\"target\":\"C\"}"
                + "]}";

        NodeExecutor llm = mockExecutor("llm", nodeId -> {
            switch (nodeId) {
                case "A":
                    return NodeExecutionResult.failed("A", "故意失败");
                case "B":
                    return NodeExecutionResult.success("B", envelope("B-done"));
                case "C":
                    return NodeExecutionResult.success("C", envelope("C-done"));
                default:
                    return NodeExecutionResult.failed(nodeId, "未知");
            }
        });
        executorsMap.put("llm", llm);

        WorkflowLevelExecutor.ExecutionOutcome outcome = runDag(graphData);

        assertThat(outcome.failed).as("skip 策略不算失败").isFalse();
        assertThat(statusCapture).containsEntry("A", "FAILED")
                .containsEntry("B", "SUCCESS")
                .containsEntry("C", "SUCCESS");
    }

    /**
     * 场景 4：A→X→B（branch-1），X→C（branch-2），X 选择 branch-1。
     *
     * <p>预期：B 正常执行，C 标记 SKIPPED，C 的 executor 不被调用。</p>
     */
    @Test
    void conditionBranchUnselected_targetMarkedAsSkipped() {
        String graphData = "{"
                + "\"nodes\":["
                + "  {\"id\":\"A\",\"type\":\"llm\",\"data\":{\"label\":\"A\"}},"
                + "  {\"id\":\"X\",\"type\":\"condition\",\"data\":{\"label\":\"X\"}},"
                + "  {\"id\":\"B\",\"type\":\"llm\",\"data\":{\"label\":\"B\"}},"
                + "  {\"id\":\"C\",\"type\":\"llm\",\"data\":{\"label\":\"C\"}}"
                + "],\"edges\":["
                + "  {\"source\":\"A\",\"target\":\"X\"},"
                + "  {\"source\":\"X\",\"target\":\"B\",\"sourceHandle\":\"branch-1\"},"
                + "  {\"source\":\"X\",\"target\":\"C\",\"sourceHandle\":\"branch-2\"}"
                + "]}";

        AtomicInteger cExecutions = new AtomicInteger();
        NodeExecutor llm = mockExecutor("llm", nodeId -> {
            switch (nodeId) {
                case "A":
                    return NodeExecutionResult.success("A", envelope("A-done"));
                case "B":
                    return NodeExecutionResult.success("B", envelope("B-done"));
                case "C":
                    cExecutions.incrementAndGet();
                    return NodeExecutionResult.success("C", envelope("C-done"));
                default:
                    return NodeExecutionResult.failed(nodeId, "未知");
            }
        });
        NodeExecutor cond = mockExecutor("condition", nodeId ->
                NodeExecutionResult.success(nodeId, envelope("X-done"), "branch-1"));
        executorsMap.put("llm", llm);
        executorsMap.put("condition", cond);

        WorkflowLevelExecutor.ExecutionOutcome outcome = runDag(graphData);

        assertThat(outcome.failed).isFalse();
        assertThat(statusCapture).containsEntry("A", "SUCCESS")
                .containsEntry("X", "SUCCESS")
                .containsEntry("B", "SUCCESS")
                .containsEntry("C", "SKIPPED");
        assertThat(cExecutions.get()).as("C 的 executor 不应被调用").isEqualTo(0);
    }

    /**
     * 场景 5：A→M, B→M（M 有两个前驱）。A 慢，B 快。
     *
     * <p>预期：M 必须在 A 和 B 都完成后才开始执行。
     * 通过测量 M.startTime >= max(A.endTime, B.endTime) 验证。</p>
     */
    @Test
    void mergeNode_scheduledOnlyAfterAllPredecessorsComplete() throws Exception {
        String graphData = "{"
                + "\"nodes\":["
                + "  {\"id\":\"A\",\"type\":\"llm\",\"data\":{\"label\":\"A\"}},"
                + "  {\"id\":\"B\",\"type\":\"llm\",\"data\":{\"label\":\"B\"}},"
                + "  {\"id\":\"M\",\"type\":\"llm\",\"data\":{\"label\":\"M\"}}"
                + "],\"edges\":["
                + "  {\"source\":\"A\",\"target\":\"M\"},"
                + "  {\"source\":\"B\",\"target\":\"M\"}"
                + "]}";

        CountDownLatch latchA = new CountDownLatch(1);
        AtomicLong mStartTime = new AtomicLong();
        AtomicLong aEndTime = new AtomicLong();
        AtomicLong bEndTime = new AtomicLong();
        NodeExecutor llm = mockExecutor("llm", nodeId -> {
            try {
                switch (nodeId) {
                    case "A":
                        Thread.sleep(150);
                        aEndTime.set(System.currentTimeMillis());
                        latchA.countDown();
                        return NodeExecutionResult.success("A", envelope("A-done"));
                    case "B":
                        // B 等 A 完成信号后再记录结束时间，确保 A、B 完成时间可比较
                        latchA.await(2, TimeUnit.SECONDS);
                        bEndTime.set(System.currentTimeMillis());
                        return NodeExecutionResult.success("B", envelope("B-done"));
                    case "M":
                        mStartTime.set(System.currentTimeMillis());
                        return NodeExecutionResult.success("M", envelope("M-done"));
                    default:
                        return NodeExecutionResult.failed(nodeId, "未知");
                }
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
                return NodeExecutionResult.failed(nodeId, "中断");
            }
        });
        executorsMap.put("llm", llm);

        WorkflowLevelExecutor.ExecutionOutcome outcome = runDag(graphData);

        assertThat(outcome.failed).isFalse();
        assertThat(statusCapture).containsEntry("A", "SUCCESS")
                .containsEntry("B", "SUCCESS")
                .containsEntry("M", "SUCCESS");
        // M 必须在 A 和 B 都完成后才开始（>= 较晚完成者的结束时间）
        assertThat(mStartTime.get())
                .as("M 必须在 A 和 B 都完成后才调度")
                .isGreaterThanOrEqualTo(Math.max(aEndTime.get(), bEndTime.get()) - 1);
    }
}
