package tech.easyflow.ai.easyagentsflow.event;

import com.alibaba.fastjson2.JSON;
import com.easyagents.flow.core.chain.repository.InMemoryLoopResultRepository;
import com.easyagents.flow.core.chain.repository.LoopInputReference;
import com.easyagents.flow.core.chain.repository.LoopResultReference;
import org.junit.Assert;
import org.junit.Test;
import org.mockito.Mockito;
import org.springframework.context.annotation.AnnotationConfigApplicationContext;
import tech.easyflow.ai.entity.WorkflowExecResult;
import tech.easyflow.ai.entity.WorkflowExecStep;
import tech.easyflow.ai.service.WorkflowExecResultService;
import tech.easyflow.ai.service.WorkflowExecStepService;
import tech.easyflow.common.mq.config.MQProperties;
import tech.easyflow.common.mq.core.MQDeadLetterService;
import tech.easyflow.common.mq.core.MQMessage;
import tech.easyflow.common.mq.core.MQProducer;

import java.lang.reflect.Field;
import java.math.BigInteger;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Date;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 工作流执行审计异步持久化测试。
 */
public class WorkflowExecutionAuditConsumerTest {

    /**
     * 验证存在测试构造器时 Spring 仍能选择生产构造器创建 Bean。
     */
    @Test
    public void shouldCreateProducerThroughSpringContext() {
        MQProducer mqProducer =
                Mockito.mock(MQProducer.class);
        MQDeadLetterService deadLetterService =
                Mockito.mock(
                        MQDeadLetterService.class);
        try (AnnotationConfigApplicationContext context =
                     new AnnotationConfigApplicationContext()) {
            context.registerBean(
                    MQProducer.class,
                    () -> mqProducer);
            context.registerBean(
                    MQDeadLetterService.class,
                    () -> deadLetterService);
            context.registerBean(
                    WorkflowExecutionAuditProducer.class);
            context.refresh();

            Assert.assertNotNull(
                    context.getBean(
                            WorkflowExecutionAuditProducer.class));
        }
    }

    /**
     * 验证生产者固定投递到单一有序分片。
     */
    @Test
    public void shouldPublishAuditEventToOrderedShard() {
        MQProducer mqProducer = Mockito.mock(MQProducer.class);
        MQDeadLetterService deadLetterService =
                Mockito.mock(MQDeadLetterService.class);
        WorkflowExecutionAuditProducer producer =
                new WorkflowExecutionAuditProducer(
                        mqProducer, deadLetterService);
        try {
            WorkflowExecutionAuditEvent event = event(
                    WorkflowExecutionAuditEvent.Type.CHAIN_STARTED,
                    "instance-1:chain-started",
                    "instance-1",
                    new WorkflowExecResult(),
                    null);

            producer.send(event);

            Mockito.verify(mqProducer).send(Mockito.argThat(message ->
                    WorkflowExecutionAuditMqConstants.TOPIC.equals(
                            message.getTopic())
                            && "instance-1:chain-started".equals(
                            message.getMessageId())
                            && "instance-1".equals(message.getKey())
                            && message.getBody().contains(
                            "CHAIN_STARTED")));
        } finally {
            producer.close();
        }
    }

    /**
     * 验证结束事件仅携带结束时间时仍可安全序列化投递。
     */
    @Test
    public void shouldPublishEndEventsWithoutStartTime() {
        MQProducer mqProducer =
                Mockito.mock(MQProducer.class);
        MQDeadLetterService deadLetterService =
                Mockito.mock(
                        MQDeadLetterService.class);
        WorkflowExecutionAuditProducer producer =
                new WorkflowExecutionAuditProducer(
                        mqProducer,
                        deadLetterService);
        Mockito.when(mqProducer.send(
                        Mockito.any()))
                .thenReturn("message-id");
        try {
            WorkflowExecStep step =
                    new WorkflowExecStep();
            step.setExecKey("step-ended");
            step.setEndTime(new Date());
            WorkflowExecResult result =
                    new WorkflowExecResult();
            result.setExecKey("instance-ended");
            result.setEndTime(new Date());

            producer.send(event(
                    WorkflowExecutionAuditEvent.Type.NODE_ENDED,
                    "step-ended:event",
                    "instance-ended",
                    null,
                    step));
            producer.send(event(
                    WorkflowExecutionAuditEvent.Type.CHAIN_ENDED,
                    "instance-ended:event",
                    "instance-ended",
                    result,
                    null));

            Mockito.verify(
                    mqProducer,
                    Mockito.times(2))
                    .send(Mockito.argThat(message ->
                            message.getBody() != null
                                    && (message.getBody()
                                    .contains("NODE_ENDED")
                                    || message.getBody()
                                    .contains("CHAIN_ENDED"))));
        } finally {
            producer.close();
        }
    }

    /**
     * 验证一个实例的毒消息退避不会阻塞其他发送 lane。
     *
     * @throws Exception 等待健康实例发送失败时抛出
     */
    @Test
    public void shouldIsolateRetryHeadBlockingAcrossLanes()
            throws Exception {
        MQProducer mqProducer =
                Mockito.mock(MQProducer.class);
        MQDeadLetterService deadLetterService =
                Mockito.mock(
                        MQDeadLetterService.class);
        CountDownLatch healthySent =
                new CountDownLatch(1);
        Mockito.when(mqProducer.send(
                        Mockito.any()))
                .thenAnswer(invocation -> {
                    MQMessage message =
                            invocation.getArgument(0);
                    if ("instance-0".equals(
                            message.getKey())) {
                        throw new IllegalStateException(
                                "poison");
                    }
                    healthySent.countDown();
                    return "sent";
                });
        WorkflowExecutionAuditProducer producer =
                new WorkflowExecutionAuditProducer(
                        mqProducer,
                        deadLetterService,
                        8,
                        100,
                        1024L * 1024L,
                        8L * 1024L * 1024L,
                        0L);
        try {
            producer.send(event(
                    WorkflowExecutionAuditEvent.Type.NODE_STARTED,
                    "poison",
                    "instance-0",
                    null,
                    new WorkflowExecStep()));
            producer.send(event(
                    WorkflowExecutionAuditEvent.Type.NODE_STARTED,
                    "healthy",
                    "instance-1",
                    null,
                    new WorkflowExecStep()));

            Assert.assertTrue(
                    healthySent.await(
                            1L,
                            TimeUnit.SECONDS));
        } finally {
            producer.close();
        }
        Mockito.verify(deadLetterService)
                .deadLetter(
                        Mockito.argThat(message ->
                                "poison".equals(
                                        message.getMessageId())),
                        Mockito.contains("shutdown"));
    }

    /**
     * 验证同一实例失败恢复后仍按原事件顺序发送。
     *
     * @throws Exception 等待重试发送失败时抛出
     */
    @Test
    public void shouldPreserveOrderWithinAuditLane()
            throws Exception {
        MQProducer mqProducer =
                Mockito.mock(MQProducer.class);
        MQDeadLetterService deadLetterService =
                Mockito.mock(
                        MQDeadLetterService.class);
        AtomicInteger firstAttempts =
                new AtomicInteger();
        CountDownLatch sent =
                new CountDownLatch(2);
        List<String> order =
                Collections.synchronizedList(
                        new ArrayList<>());
        Mockito.when(mqProducer.send(
                        Mockito.any()))
                .thenAnswer(invocation -> {
                    MQMessage message =
                            invocation.getArgument(0);
                    if ("first".equals(
                            message.getMessageId())
                            && firstAttempts
                            .getAndIncrement() == 0) {
                        throw new IllegalStateException(
                                "temporary");
                    }
                    order.add(
                            message.getMessageId());
                    sent.countDown();
                    return "sent";
                });
        WorkflowExecutionAuditProducer producer =
                new WorkflowExecutionAuditProducer(
                        mqProducer,
                        deadLetterService,
                        2,
                        100,
                        1024L * 1024L,
                        8L * 1024L * 1024L,
                        1000L);
        try {
            producer.send(event(
                    WorkflowExecutionAuditEvent.Type.NODE_STARTED,
                    "first",
                    "same-instance",
                    null,
                    new WorkflowExecStep()));
            producer.send(event(
                    WorkflowExecutionAuditEvent.Type.NODE_ENDED,
                    "second",
                    "same-instance",
                    null,
                    new WorkflowExecStep()));

            Assert.assertTrue(
                    sent.await(
                            2L,
                            TimeUnit.SECONDS));
            Assert.assertEquals(
                    List.of("first", "second"),
                    order);
        } finally {
            producer.close();
        }
    }

    /**
     * 验证超出单条字节预算的审计消息直接进入死信并显式失败。
     */
    @Test
    public void shouldRejectOversizedAuditMessage() {
        MQProducer mqProducer =
                Mockito.mock(MQProducer.class);
        MQDeadLetterService deadLetterService =
                Mockito.mock(
                        MQDeadLetterService.class);
        WorkflowExecutionAuditProducer producer =
                new WorkflowExecutionAuditProducer(
                        mqProducer,
                        deadLetterService,
                        1,
                        10,
                        512L,
                        2048L,
                        0L);
        WorkflowExecResult result =
                new WorkflowExecResult();
        result.setOutput(
                "x".repeat(1024));
        try {
            producer.send(event(
                    WorkflowExecutionAuditEvent.Type.CHAIN_ENDED,
                    "oversized",
                    "instance",
                    result,
                    null));
            Assert.fail(
                    "oversized audit message should fail");
        } catch (IllegalArgumentException expected) {
            Assert.assertTrue(
                    expected.getMessage()
                            .contains("byte limit"));
        } finally {
            producer.close();
        }
        Mockito.verify(
                        deadLetterService)
                .deadLetter(
                        Mockito.argThat(message ->
                                "oversized".equals(
                                        message.getMessageId())),
                        Mockito.contains("byte limit"));
        Mockito.verifyNoInteractions(
                mqProducer);
    }

    /**
     * 验证关闭期间仍在直发的失败消息会转入死信且不会重新形成孤儿积压。
     *
     * @throws Exception 并发关闭、等待或反射读取失败时抛出
     */
    @Test
    public void shouldNotEnqueueAfterConcurrentClose()
            throws Exception {
        MQProducer mqProducer =
                Mockito.mock(MQProducer.class);
        MQDeadLetterService deadLetterService =
                Mockito.mock(
                        MQDeadLetterService.class);
        CountDownLatch sending =
                new CountDownLatch(1);
        CountDownLatch releaseSend =
                new CountDownLatch(1);
        Mockito.when(mqProducer.send(
                        Mockito.any()))
                .thenAnswer(invocation -> {
                    sending.countDown();
                    releaseSend.await(
                            2L,
                            TimeUnit.SECONDS);
                    throw new IllegalStateException(
                            "send failed during close");
                });
        WorkflowExecutionAuditProducer producer =
                new WorkflowExecutionAuditProducer(
                        mqProducer,
                        deadLetterService,
                        1,
                        100,
                        1024L * 1024L,
                        8L * 1024L * 1024L,
                        0L);
        ExecutorService callers =
                Executors.newFixedThreadPool(2);
        try {
            Future<?> sender =
                    callers.submit(() -> {
                        try {
                            producer.send(event(
                                    WorkflowExecutionAuditEvent.Type.NODE_STARTED,
                                    "closing",
                                    "instance",
                                    null,
                                    new WorkflowExecStep()));
                            Assert.fail(
                                    "send should report concurrent close");
                        } catch (IllegalStateException expected) {
                            Assert.assertTrue(
                                    expected.getMessage()
                                            .contains("closed"));
                        }
                    });
            Assert.assertTrue(
                    sending.await(
                            1L,
                            TimeUnit.SECONDS));
            Future<?> closer =
                    callers.submit(
                            producer::close);
            closer.get(
                    1L,
                    TimeUnit.SECONDS);
            releaseSend.countDown();
            sender.get(
                    2L,
                    TimeUnit.SECONDS);

            Field backlogCountField =
                    WorkflowExecutionAuditProducer.class
                            .getDeclaredField(
                                    "backlogCount");
            backlogCountField.setAccessible(true);
            Assert.assertEquals(
                    0,
                    backlogCountField.getInt(
                            producer));
            Mockito.verify(deadLetterService)
                    .deadLetter(
                            Mockito.argThat(message ->
                                    "closing".equals(
                                            message.getMessageId())),
                            Mockito.contains(
                                    "closed during send"));
        } finally {
            releaseSend.countDown();
            producer.close();
            callers.shutdownNow();
        }
    }

    /**
     * 验证启动、节点开始、节点结束和流程结束事件按顺序幂等落库。
     */
    @Test
    public void shouldApplyOrderedExecutionAuditEvents() {
        WorkflowExecResultService resultService =
                Mockito.mock(WorkflowExecResultService.class);
        WorkflowExecStepService stepService =
                Mockito.mock(WorkflowExecStepService.class);
        InMemoryLoopResultRepository loopRepository =
                new InMemoryLoopResultRepository();
        WorkflowExecutionAuditConsumer consumer =
                new WorkflowExecutionAuditConsumer(
                        resultService,
                        stepService,
                        new MQProperties(),
                        loopRepository);

        WorkflowExecResult persistedResult = new WorkflowExecResult();
        persistedResult.setId(BigInteger.ONE);
        persistedResult.setExecKey("instance-1");
        Mockito.when(resultService.getByExecKey("instance-1"))
                .thenReturn(persistedResult);
        Mockito.when(resultService.updateByExecKey(Mockito.any()))
                .thenReturn(1);
        Mockito.when(stepService.updateByExecKey(Mockito.any()))
                .thenReturn(1);

        WorkflowExecResult startRecord = new WorkflowExecResult();
        startRecord.setExecKey("instance-1");
        startRecord.setStatus(1);
        WorkflowExecStep startStep = new WorkflowExecStep();
        startStep.setExecKey("step-1");
        startStep.setNodeId("node-1");
        startStep.setNodeName("node");
        startStep.setStatus(1);
        WorkflowExecStep endStep = new WorkflowExecStep();
        endStep.setExecKey("step-1");
        endStep.setStatus(2);
        endStep.setOutput("{\"value\":1}");
        WorkflowExecResult endRecord = new WorkflowExecResult();
        endRecord.setExecKey("instance-1");
        endRecord.setStatus(2);
        endRecord.setOutput("{\"value\":1}");

        consumer.handle(List.of(
                message(event(WorkflowExecutionAuditEvent.Type.CHAIN_STARTED,
                        "start", "instance-1", startRecord, null)),
                message(event(WorkflowExecutionAuditEvent.Type.NODE_STARTED,
                        "node-start", "instance-1", null, startStep)),
                message(event(WorkflowExecutionAuditEvent.Type.NODE_ENDED,
                        "node-end", "instance-1", null, endStep)),
                message(event(WorkflowExecutionAuditEvent.Type.CHAIN_ENDED,
                        "end", "instance-1", endRecord, null))
        ));

        Mockito.verify(resultService).save(Mockito.argThat(record ->
                "instance-1".equals(record.getExecKey())
                        && Integer.valueOf(1).equals(record.getStatus())));
        Mockito.verify(stepService).save(Mockito.argThat(step ->
                "step-1".equals(step.getExecKey())
                        && BigInteger.ONE.equals(step.getRecordId())));
        Mockito.verify(stepService).updateByExecKey(
                Mockito.argThat(step ->
                        "step-1".equals(step.getExecKey())
                                && "{\"value\":1}".equals(
                                step.getOutput())));
        Mockito.verify(resultService).updateByExecKey(
                Mockito.argThat(record ->
                        "instance-1".equals(record.getExecKey())
                                && "{\"value\":1}".equals(
                                record.getOutput())));
        Mockito.verify(stepService, Mockito.never())
                .getByExecKey(Mockito.anyString());
    }

    /**
     * 验证大型结果引用在审计消费线程还原，持久记录仍保持完整 JSON。
     */
    @Test
    public void shouldResolveLargeReferenceInAuditConsumer() {
        WorkflowExecResultService resultService =
                Mockito.mock(
                        WorkflowExecResultService.class);
        WorkflowExecStepService stepService =
                Mockito.mock(
                        WorkflowExecStepService.class);
        InMemoryLoopResultRepository loopRepository =
                new InMemoryLoopResultRepository();
        String resultId = "instance:dataset:rows";
        loopRepository.storeInput(
                resultId, List.of(1, 2, 3));
        WorkflowExecutionAuditConsumer consumer =
                new WorkflowExecutionAuditConsumer(
                        resultService,
                        stepService,
                        new MQProperties(),
                        loopRepository);
        Mockito.when(stepService.updateByExecKey(
                        Mockito.any()))
                .thenReturn(1);
        WorkflowExecStep incoming =
                new WorkflowExecStep();
        incoming.setExecKey("step-reference");
        incoming.setOutput(JSON.toJSONString(
                Map.of(
                        "data",
                        new LoopInputReference(
                                resultId, 3))));

        consumer.handle(List.of(message(event(
                WorkflowExecutionAuditEvent.Type.NODE_ENDED,
                "node-reference",
                "instance",
                null,
                incoming))));

        Mockito.verify(stepService).updateByExecKey(
                Mockito.argThat(step ->
                        "{\"data\":[1,2,3]}"
                                .equals(step.getOutput())));
    }

    /**
     * 验证节点启动输入引用在审计消费者中还原后再保存。
     */
    @Test
    public void shouldResolveLargeInputReferenceWhenCreatingStep() {
        WorkflowExecResultService resultService =
                Mockito.mock(
                        WorkflowExecResultService.class);
        WorkflowExecStepService stepService =
                Mockito.mock(
                        WorkflowExecStepService.class);
        InMemoryLoopResultRepository loopRepository =
                new InMemoryLoopResultRepository();
        String resultId = "instance:dataset:input";
        loopRepository.storeInput(
                resultId, List.of(1, 2, 3));
        WorkflowExecutionAuditConsumer consumer =
                new WorkflowExecutionAuditConsumer(
                        resultService,
                        stepService,
                        new MQProperties(),
                        loopRepository);
        WorkflowExecResult record =
                new WorkflowExecResult();
        record.setId(BigInteger.ONE);
        Mockito.when(resultService.getByExecKey(
                        "instance"))
                .thenReturn(record);
        WorkflowExecStep incoming =
                new WorkflowExecStep();
        incoming.setExecKey("step-input");
        incoming.setInput(JSON.toJSONString(
                Map.of(
                        "items",
                        new LoopInputReference(
                                resultId, 3))));

        consumer.handle(List.of(message(event(
                WorkflowExecutionAuditEvent.Type.NODE_STARTED,
                "node-input",
                "instance",
                null,
                incoming))));

        Mockito.verify(stepService).save(
                Mockito.argThat(step ->
                        BigInteger.ONE.equals(
                                step.getRecordId())
                                && "{\"items\":[1,2,3]}"
                                .equals(step.getInput())));
    }

    /**
     * 验证节点结束与流程结束审计均还原循环累计输出。
     */
    @Test
    @SuppressWarnings("unchecked")
    public void shouldResolveLoopResultReferenceForEndedAudits() {
        WorkflowExecResultService resultService =
                Mockito.mock(
                        WorkflowExecResultService.class);
        WorkflowExecStepService stepService =
                Mockito.mock(
                        WorkflowExecStepService.class);
        InMemoryLoopResultRepository loopRepository =
                new InMemoryLoopResultRepository();
        String resultId = "instance:loop:result";
        loopRepository.append(
                resultId,
                0,
                Map.of("answer", List.of("first", "second")));
        loopRepository.append(
                resultId,
                1,
                Map.of("answer", List.of("third")));
        WorkflowExecutionAuditConsumer consumer =
                new WorkflowExecutionAuditConsumer(
                        resultService,
                        stepService,
                        new MQProperties(),
                        loopRepository);
        Mockito.when(stepService.updateByExecKey(
                        Mockito.any()))
                .thenReturn(1);
        Mockito.when(resultService.updateByExecKey(
                        Mockito.any()))
                .thenReturn(1);
        Map<String, Object> referenceOutput =
                Map.of(
                        "answers",
                        new LoopResultReference(
                                resultId,
                                2,
                                "answer",
                                true));
        WorkflowExecStep incomingStep =
                new WorkflowExecStep();
        incomingStep.setExecKey("step-loop");
        incomingStep.setOutput(
                JSON.toJSONString(
                        referenceOutput));
        WorkflowExecResult incomingResult =
                new WorkflowExecResult();
        incomingResult.setExecKey(
                "instance-loop");
        incomingResult.setOutput(
                JSON.toJSONString(
                        referenceOutput));

        consumer.handle(List.of(
                message(event(
                        WorkflowExecutionAuditEvent.Type.NODE_ENDED,
                        "node-loop-ended",
                        "instance-loop",
                        null,
                        incomingStep)),
                message(event(
                        WorkflowExecutionAuditEvent.Type.CHAIN_ENDED,
                        "chain-loop-ended",
                        "instance-loop",
                        incomingResult,
                        null))));

        org.mockito.ArgumentCaptor<WorkflowExecStep>
                stepCaptor =
                org.mockito.ArgumentCaptor.forClass(
                        WorkflowExecStep.class);
        org.mockito.ArgumentCaptor<WorkflowExecResult>
                resultCaptor =
                org.mockito.ArgumentCaptor.forClass(
                        WorkflowExecResult.class);
        Mockito.verify(stepService)
                .updateByExecKey(stepCaptor.capture());
        Mockito.verify(resultService)
                .updateByExecKey(resultCaptor.capture());
        Map<String, Object> stepOutput =
                JSON.parseObject(
                        stepCaptor.getValue().getOutput(),
                        Map.class);
        Map<String, Object> resultOutput =
                JSON.parseObject(
                        resultCaptor.getValue().getOutput(),
                        Map.class);
        Assert.assertEquals(
                List.of("first", "second", "third"),
                stepOutput.get("answers"));
        Assert.assertEquals(
                List.of("first", "second", "third"),
                resultOutput.get("answers"));
    }

    /**
     * 构造审计事件。
     *
     * @param type 事件类型
     * @param eventId 事件 ID
     * @param instanceId 实例 ID
     * @param result 工作流记录
     * @param step 节点步骤
     * @return 审计事件
     */
    private WorkflowExecutionAuditEvent event(WorkflowExecutionAuditEvent.Type type,
                                              String eventId,
                                              String instanceId,
                                              WorkflowExecResult result,
                                              WorkflowExecStep step) {
        WorkflowExecutionAuditEvent event = new WorkflowExecutionAuditEvent();
        event.setType(type);
        event.setEventId(eventId);
        event.setInstanceId(instanceId);
        event.setOccurredAt(new Date());
        event.setResult(result);
        event.setStep(step);
        return event;
    }

    /**
     * 将审计事件包装为通用 MQ 消息。
     *
     * @param event 审计事件
     * @return MQ 消息
     */
    private MQMessage message(WorkflowExecutionAuditEvent event) {
        MQMessage message = new MQMessage();
        message.setMessageId(event.getEventId());
        message.setBody(JSON.toJSONString(event));
        return message;
    }
}
