package com.easyagents.flow.core.test;

import com.alibaba.fastjson.JSONObject;
import com.easyagents.flow.core.chain.Chain;
import com.easyagents.flow.core.chain.ChainDefinition;
import com.easyagents.flow.core.chain.ChainState;
import com.easyagents.flow.core.chain.ChainStatus;
import com.easyagents.flow.core.chain.Edge;
import com.easyagents.flow.core.chain.Node;
import com.easyagents.flow.core.chain.NodeCondition;
import com.easyagents.flow.core.chain.NodeJoinMode;
import com.easyagents.flow.core.chain.NodeState;
import com.easyagents.flow.core.chain.Parameter;
import com.easyagents.flow.core.chain.RefType;
import com.easyagents.flow.core.chain.repository.InMemoryChainStateRepository;
import com.easyagents.flow.core.chain.repository.InMemoryNodeStateRepository;
import com.easyagents.flow.core.chain.runtime.ChainExecutor;
import com.easyagents.flow.core.chain.runtime.InMemoryTriggerStore;
import com.easyagents.flow.core.chain.runtime.TriggerScheduler;
import com.easyagents.flow.core.node.BaseNode;
import com.easyagents.flow.core.node.EndNode;
import com.easyagents.flow.core.node.StartNode;
import com.easyagents.flow.core.parser.ChainParser;
import com.easyagents.flow.core.parser.impl.EndNodeParser;
import org.junit.Assert;
import org.junit.Test;

import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.BooleanSupplier;

/**
 * 验证普通节点的直接入边汇聚模式。
 */
public class NodeJoinModeTest {

    @Test
    public void shouldParseJoinModeAndRejectInvalidValues() {
        ChainParser parser = ChainParser.builder()
                .withDefaultParsers(true)
                .build();

        Assert.assertEquals(
                NodeJoinMode.ANY,
                parseEndNode(parser, null, null).getJoinMode());
        Assert.assertEquals(
                NodeJoinMode.ALL,
                parseEndNode(parser, "all", null).getJoinMode());
        Assert.assertEquals(
                NodeJoinMode.ANY,
                parseEndNode(parser, "ANY", null).getJoinMode());

        assertInvalidJoinMode(() -> parseEndNode(parser, "first", null));
        assertInvalidJoinMode(() -> parseEndNode(parser, "", null));
        assertInvalidJoinMode(() -> parseEndNode(parser, "all", "loop"));

        ProbeJoinNode node = new ProbeJoinNode(new AtomicInteger());
        node.setJoinMode(NodeJoinMode.ALL);
        assertInvalidJoinMode(() -> node.setParentId("loop"));
    }

    @Test
    public void shouldWaitForEveryInboundEdgeBeforeExecutingAndCheckingCondition()
            throws Exception {
        JoinFixture fixture = createJoinFixture(NodeJoinMode.ALL, true);
        String instanceId = null;
        try {
            instanceId = fixture.executor.executeAsync(
                    fixture.definition.getId(), Collections.emptyMap());
            Assert.assertTrue(fixture.branchACompleted.await(2, TimeUnit.SECONDS));
            Assert.assertTrue(fixture.branchBStarted.await(2, TimeUnit.SECONDS));

            String currentInstanceId = instanceId;
            await(() -> hasTriggerEdge(
                    fixture.nodeStateRepository,
                    currentInstanceId,
                    "join",
                    "a-join"));
            Assert.assertEquals(0, fixture.joinExecutions.get());
            Assert.assertEquals(0, fixture.conditionChecks.get());

            fixture.releaseBranchB.countDown();
            ChainState finalState = awaitTerminal(
                    fixture.chainStateRepository, instanceId);

            Assert.assertEquals(1, fixture.joinExecutions.get());
            Assert.assertEquals(1, fixture.conditionChecks.get());
            Assert.assertEquals("A+B", finalState.getExecuteResult().get("combined"));
            Assert.assertEquals(Boolean.TRUE, finalState.getExecuteResult().get("sawBoth"));
        } finally {
            fixture.releaseBranchB.countDown();
            fixture.scheduler.shutdown();
        }
    }

    @Test
    public void shouldKeepAnyModeFirstArrivalBehavior() throws Exception {
        JoinFixture fixture = createJoinFixture(NodeJoinMode.ANY, false);
        try {
            String instanceId = fixture.executor.executeAsync(
                    fixture.definition.getId(), Collections.emptyMap());
            Assert.assertTrue(fixture.branchACompleted.await(2, TimeUnit.SECONDS));
            Assert.assertTrue(fixture.branchBStarted.await(2, TimeUnit.SECONDS));

            await(() -> fixture.joinExecutions.get() > 0);
            ChainState finalState = awaitTerminal(
                    fixture.chainStateRepository, instanceId);

            Assert.assertEquals(1, fixture.joinExecutions.get());
            Assert.assertEquals(Boolean.FALSE, finalState.getExecuteResult().get("sawBoth"));
        } finally {
            fixture.releaseBranchB.countDown();
            fixture.scheduler.shutdown();
        }
    }

    @Test
    public void shouldKeepDefaultAnyRetryAndLoopBehavior() throws Exception {
        ScheduledExecutorService schedulerPool =
                Executors.newSingleThreadScheduledExecutor();
        ExecutorService workerPool = Executors.newFixedThreadPool(3);
        TriggerScheduler scheduler = new TriggerScheduler(
                new InMemoryTriggerStore(), schedulerPool, workerPool, 1000L);
        AtomicInteger executions = new AtomicInteger();
        ChainDefinition definition = new ChainDefinition();
        definition.setId("join-mode-retry-loop");
        StartNode start = new StartNode();
        start.setId("start");
        RetryLoopNode worker = new RetryLoopNode(executions);
        worker.setId("worker");
        worker.setRetryEnable(true);
        worker.setMaxRetryCount(1);
        worker.setRetryIntervalMs(0L);
        worker.setLoopEnable(true);
        worker.setMaxLoopCount(2);
        worker.setLoopIntervalMs(0L);
        EndNode end = endNode("worker", "count", "count");
        definition.addNode(start);
        definition.addNode(worker);
        definition.addNode(end);
        definition.addEdge(edge("start-worker", "start", "worker"));
        definition.addEdge(edge("worker-end", "worker", "end"));

        ChainExecutor executor = new ChainExecutor(
                ignored -> definition,
                new InMemoryChainStateRepository(),
                new InMemoryNodeStateRepository(),
                scheduler);
        try {
            Map<String, Object> result = executor.execute(
                    definition.getId(), Collections.emptyMap(), 5L, TimeUnit.SECONDS);

            Assert.assertEquals(NodeJoinMode.ANY, worker.getJoinMode());
            Assert.assertEquals(3, executions.get());
            Assert.assertEquals(3, result.get("count"));
        } finally {
            scheduler.shutdown();
        }
    }

    private JoinFixture createJoinFixture(
            NodeJoinMode joinMode, boolean withCondition) {
        JoinFixture fixture = new JoinFixture();
        fixture.schedulerPool = Executors.newScheduledThreadPool(2);
        fixture.workerPool = Executors.newFixedThreadPool(4);
        fixture.scheduler = new TriggerScheduler(
                new InMemoryTriggerStore(),
                fixture.schedulerPool,
                fixture.workerPool,
                1000L);
        fixture.chainStateRepository = new InMemoryChainStateRepository();
        fixture.nodeStateRepository = new InMemoryNodeStateRepository();
        fixture.definition = new ChainDefinition();
        fixture.definition.setId("join-mode-" + joinMode.getValue());

        StartNode start = new StartNode();
        start.setId("start");
        BranchNode branchA = new BranchNode(
                "A", fixture.branchACompleted, null, null);
        branchA.setId("a");
        BranchNode branchB = new BranchNode(
                "B", null, fixture.branchBStarted, fixture.releaseBranchB);
        branchB.setId("b");
        ProbeJoinNode join = new ProbeJoinNode(fixture.joinExecutions);
        join.setId("join");
        join.setJoinMode(joinMode);
        if (withCondition) {
            join.setCondition(new BothOutputsCondition(fixture.conditionChecks));
        }
        EndNode end = endNode("join", "combined", "combined");
        end.addOutputDef(outputRef("sawBoth", "join.sawBoth"));

        fixture.definition.addNode(start);
        fixture.definition.addNode(branchA);
        fixture.definition.addNode(branchB);
        fixture.definition.addNode(join);
        fixture.definition.addNode(end);
        fixture.definition.addEdge(edge("start-a", "start", "a"));
        fixture.definition.addEdge(edge("start-b", "start", "b"));
        fixture.definition.addEdge(edge("a-join", "a", "join"));
        fixture.definition.addEdge(edge("b-join", "b", "join"));
        fixture.definition.addEdge(edge("join-end", "join", "end"));

        fixture.executor = new ChainExecutor(
                ignored -> fixture.definition,
                fixture.chainStateRepository,
                fixture.nodeStateRepository,
                fixture.scheduler);
        return fixture;
    }

    private Node parseEndNode(
            ChainParser parser, String joinMode, String parentId) {
        JSONObject data = new JSONObject();
        if (joinMode != null) {
            data.put("joinMode", joinMode);
        }
        JSONObject nodeJson = new JSONObject();
        nodeJson.put("id", "end");
        nodeJson.put("type", "endNode");
        nodeJson.put("parentId", parentId);
        nodeJson.put("data", data);
        return new EndNodeParser().parse(
                nodeJson, new JSONObject(), parser);
    }

    private void assertInvalidJoinMode(Runnable action) {
        try {
            action.run();
            Assert.fail("Expected invalid join mode");
        } catch (IllegalArgumentException expected) {
            Assert.assertTrue(expected.getMessage().contains("joinMode"));
        }
    }

    private static boolean hasTriggerEdge(
            InMemoryNodeStateRepository repository,
            String instanceId,
            String nodeId,
            String edgeId) {
        NodeState state = repository.load(instanceId, nodeId);
        return state != null && state.getTriggerEdgeIds().contains(edgeId);
    }

    private static ChainState awaitTerminal(
            InMemoryChainStateRepository repository,
            String instanceId) throws Exception {
        await(() -> {
            ChainState state = repository.load(instanceId);
            return state != null
                    && state.getStatus() != null
                    && state.getStatus().isTerminal();
        });
        ChainState state = repository.load(instanceId);
        Assert.assertEquals(ChainStatus.SUCCEEDED, state.getStatus());
        return state;
    }

    private static void await(BooleanSupplier condition) throws Exception {
        long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(3L);
        while (!condition.getAsBoolean() && System.nanoTime() < deadline) {
            Thread.sleep(10L);
        }
        Assert.assertTrue("condition was not met before timeout", condition.getAsBoolean());
    }

    private static EndNode endNode(
            String sourceNodeId, String sourceName, String outputName) {
        EndNode end = new EndNode();
        end.setId("end");
        end.addOutputDef(outputRef(
                outputName, sourceNodeId + "." + sourceName));
        return end;
    }

    private static Parameter outputRef(String name, String ref) {
        Parameter parameter = new Parameter();
        parameter.setName(name);
        parameter.setRef(ref);
        parameter.setRefType(RefType.REF);
        return parameter;
    }

    private static Edge edge(String id, String source, String target) {
        Edge edge = new Edge();
        edge.setId(id);
        edge.setSource(source);
        edge.setTarget(target);
        return edge;
    }

    private static final class BranchNode extends BaseNode {
        private final String value;
        private final CountDownLatch completed;
        private final CountDownLatch started;
        private final CountDownLatch release;

        private BranchNode(
                String value,
                CountDownLatch completed,
                CountDownLatch started,
                CountDownLatch release) {
            this.value = value;
            this.completed = completed;
            this.started = started;
            this.release = release;
        }

        @Override
        public Map<String, Object> execute(Chain chain) {
            if (started != null) {
                started.countDown();
            }
            if (release != null) {
                try {
                    if (!release.await(3L, TimeUnit.SECONDS)) {
                        throw new IllegalStateException("branch release timed out");
                    }
                } catch (InterruptedException error) {
                    Thread.currentThread().interrupt();
                    throw new IllegalStateException("branch interrupted", error);
                }
            }
            if (completed != null) {
                completed.countDown();
            }
            return Collections.singletonMap("value", value);
        }
    }

    private static final class ProbeJoinNode extends BaseNode {
        private final AtomicInteger executions;

        private ProbeJoinNode(AtomicInteger executions) {
            this.executions = executions;
        }

        @Override
        public Map<String, Object> execute(Chain chain) {
            executions.incrementAndGet();
            Object a = chain.getExecutionState().getMemory().get("a.value");
            Object b = chain.getExecutionState().getMemory().get("b.value");
            Map<String, Object> result = new HashMap<>();
            result.put("combined", String.valueOf(a) + "+" + String.valueOf(b));
            result.put("sawBoth", a != null && b != null);
            return result;
        }
    }

    private static final class BothOutputsCondition implements NodeCondition {
        private static final long serialVersionUID = 1L;
        private final AtomicInteger checks;

        private BothOutputsCondition(AtomicInteger checks) {
            this.checks = checks;
        }

        @Override
        public boolean check(
                Chain chain,
                NodeState context,
                Map<String, Object> executeResult) {
            checks.incrementAndGet();
            Map<String, Object> memory = chain.getExecutionState().getMemory();
            return memory.containsKey("a.value") && memory.containsKey("b.value");
        }
    }

    private static final class RetryLoopNode extends BaseNode {
        private final AtomicInteger executions;

        private RetryLoopNode(AtomicInteger executions) {
            this.executions = executions;
        }

        @Override
        public Map<String, Object> execute(Chain chain) {
            int count = executions.incrementAndGet();
            if (count == 1) {
                throw new IllegalStateException("retry once");
            }
            return Collections.singletonMap("count", count);
        }
    }

    private static final class JoinFixture {
        private final CountDownLatch branchACompleted = new CountDownLatch(1);
        private final CountDownLatch branchBStarted = new CountDownLatch(1);
        private final CountDownLatch releaseBranchB = new CountDownLatch(1);
        private final AtomicInteger joinExecutions = new AtomicInteger();
        private final AtomicInteger conditionChecks = new AtomicInteger();
        private ScheduledExecutorService schedulerPool;
        private ExecutorService workerPool;
        private TriggerScheduler scheduler;
        private InMemoryChainStateRepository chainStateRepository;
        private InMemoryNodeStateRepository nodeStateRepository;
        private ChainDefinition definition;
        private ChainExecutor executor;
    }
}
