package tech.easyflow.ai.easyagentsflow.listener;

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.NodeStatus;
import com.easyagents.flow.core.chain.event.NodeEndEvent;
import com.easyagents.flow.core.chain.event.NodeStartEvent;
import com.easyagents.flow.core.chain.repository.InMemoryChainStateRepository;
import com.easyagents.flow.core.node.StartNode;
import org.junit.Assert;
import org.junit.Test;
import org.mockito.ArgumentCaptor;
import org.mockito.Mockito;
import tech.easyflow.ai.easyagentsflow.event.WorkflowExecutionAuditEvent;
import tech.easyflow.ai.easyagentsflow.event.WorkflowExecutionAuditProducer;

import java.lang.reflect.Field;
import java.util.List;
import java.util.Map;
import java.util.UUID;

/**
 * {@link ChainEventListenerForSave} 节点审计归属回归测试。
 */
public class ChainEventListenerForSaveTest {

    /**
     * 验证 parent-linked 节点开始与结束事件使用同一顶级实例顺序键。
     *
     * @throws Exception 测试依赖注入失败时抛出
     */
    @Test
    public void shouldUseSameRootInstanceForNodeStartAndEnd()
            throws Exception {
        String suffix =
                UUID.randomUUID().toString();
        String rootId =
                "root-" + suffix;
        String childId =
                "child-" + suffix;
        InMemoryChainStateRepository repository =
                new InMemoryChainStateRepository();
        ChainState root =
                repository.create(rootId);
        ChainState child =
                repository.create(childId);
        child.setParentInstanceId(rootId);
        child.setAuditInstanceId(null);

        ChainDefinition definition =
                new ChainDefinition();
        definition.setId("1");
        StartNode node =
                new StartNode();
        node.setId("node");
        definition.addNode(node);
        Chain chain =
                new Chain(definition, childId);
        chain.setChainStateRepository(repository);
        WorkflowExecutionAuditProducer producer =
                Mockito.mock(
                        WorkflowExecutionAuditProducer.class);
        ChainEventListenerForSave listener =
                new ChainEventListenerForSave();
        Field producerField =
                ChainEventListenerForSave.class
                        .getDeclaredField("auditProducer");
        producerField.setAccessible(true);
        producerField.set(listener, producer);

        listener.onEvent(
                new NodeStartEvent(
                        chain,
                        node,
                        "attempt",
                        NodeStatus.RUNNING,
                        chain.getAuditInstanceId()),
                chain);
        listener.onEvent(
                new NodeEndEvent(
                        chain,
                        node,
                        Map.of("value", "ok"),
                        null,
                        NodeStatus.SUCCEEDED,
                        "attempt"),
                chain);

        ArgumentCaptor<WorkflowExecutionAuditEvent> captor =
                ArgumentCaptor.forClass(
                        WorkflowExecutionAuditEvent.class);
        Mockito.verify(
                        producer,
                        Mockito.times(2))
                .send(captor.capture());
        List<WorkflowExecutionAuditEvent> events =
                captor.getAllValues();
        Assert.assertEquals(
                WorkflowExecutionAuditEvent.Type.NODE_STARTED,
                events.get(0).getType());
        Assert.assertEquals(
                WorkflowExecutionAuditEvent.Type.NODE_ENDED,
                events.get(1).getType());
        Assert.assertEquals(
                rootId,
                events.get(0).getInstanceId());
        Assert.assertEquals(
                rootId,
                events.get(1).getInstanceId());
    }
}
