package tech.easyflow.agent.distributed;

import com.fasterxml.jackson.databind.ObjectMapper;
import org.junit.Assert;
import org.junit.Test;
import tech.easyflow.agent.config.AgentRuntimeProperties;
import tech.easyflow.agent.distributed.AgentRuntimeCommandAction;
import tech.easyflow.agent.distributed.AgentRuntimeCommandConsumer;
import tech.easyflow.agent.distributed.AgentRuntimeCommandMessage;
import tech.easyflow.agent.distributed.AgentRuntimeCommandResultRegistry;
import tech.easyflow.agent.runtime.AgentRunService;
import tech.easyflow.common.mq.config.MQProperties;
import tech.easyflow.common.mq.core.MQMessage;

import java.math.BigInteger;
import java.util.List;

/**
 * {@link AgentRuntimeCommandConsumer} 回归测试。
 */
public class AgentRuntimeCommandConsumerTest {

    /**
     * 验证消费者只处理发给当前节点的命令。
     *
     * @throws Exception 消息序列化异常
     */
    @Test
    public void consumerShouldHandleOnlyCurrentNodeCommand() throws Exception {
        AgentRuntimeProperties properties = new AgentRuntimeProperties();
        properties.setInstanceId("node-a");
        MQProperties mqProperties = new MQProperties();
        mqProperties.getRedis().setChatPersistShardCount(4);
        RecordingAgentRunService service = new RecordingAgentRunService();
        RecordingCommandResultRegistry resultRegistry = new RecordingCommandResultRegistry();
        AgentRuntimeCommandConsumer consumer =
                new AgentRuntimeCommandConsumer(new ObjectMapper(), properties, mqProperties, service, resultRegistry);

        consumer.handle(List.of(message(command("cmd-1", "node-b")), message(command("cmd-2", "node-a"))));

        Assert.assertEquals(1, service.approveCount);
        Assert.assertEquals("request-cmd-2", service.lastRequestId);
        Assert.assertEquals(4, consumer.subscription().getShardCount());
        Assert.assertFalse(consumer.subscription().isBatchEnabled());
        Assert.assertEquals("cmd-2", resultRegistry.lastSuccessCommandId);
    }

    /**
     * 验证 owner 本机执行失败时写入失败结果，避免 MQ 重试重复消费一次性 token。
     *
     * @throws Exception 消息序列化异常
     */
    @Test
    public void consumerShouldMarkFailureWhenLocalRuntimeFails() throws Exception {
        AgentRuntimeProperties properties = new AgentRuntimeProperties();
        properties.setInstanceId("node-a");
        MQProperties mqProperties = new MQProperties();
        FailingAgentRunService service = new FailingAgentRunService();
        RecordingCommandResultRegistry resultRegistry = new RecordingCommandResultRegistry();
        AgentRuntimeCommandConsumer consumer =
                new AgentRuntimeCommandConsumer(new ObjectMapper(), properties, mqProperties, service, resultRegistry);

        consumer.handle(List.of(message(command("cmd-1", "node-a"))));

        Assert.assertEquals("cmd-1", resultRegistry.lastFailureCommandId);
        Assert.assertEquals("runtime missing", resultRegistry.lastFailureMessage);
    }

    /**
     * 验证成功结果写入失败不会再次执行或改写为失败结果。
     *
     * @throws Exception 消息序列化异常
     */
    @Test
    public void consumerShouldNotMarkFailureWhenSuccessResultWriteFails() throws Exception {
        AgentRuntimeProperties properties = new AgentRuntimeProperties();
        properties.setInstanceId("node-a");
        MQProperties mqProperties = new MQProperties();
        RecordingAgentRunService service = new RecordingAgentRunService();
        FailingSuccessResultRegistry resultRegistry = new FailingSuccessResultRegistry();
        AgentRuntimeCommandConsumer consumer =
                new AgentRuntimeCommandConsumer(new ObjectMapper(), properties, mqProperties, service, resultRegistry);

        consumer.handle(List.of(message(command("cmd-1", "node-a"))));

        Assert.assertEquals(1, service.approveCount);
        Assert.assertNull(resultRegistry.lastFailureCommandId);
    }

    private AgentRuntimeCommandMessage command(String commandId, String targetNodeId) {
        AgentRuntimeCommandMessage command = new AgentRuntimeCommandMessage();
        command.setCommandId(commandId);
        command.setRequestId("request-" + commandId);
        command.setResumeToken("token-" + commandId);
        command.setAction(AgentRuntimeCommandAction.APPROVE);
        command.setOperatorId(BigInteger.ONE);
        command.setUserId("1");
        command.setTargetNodeId(targetNodeId);
        return command;
    }

    private MQMessage message(AgentRuntimeCommandMessage command) throws Exception {
        MQMessage message = new MQMessage();
        message.setMessageId(command.getCommandId());
        message.setBody(new ObjectMapper().writeValueAsString(command));
        return message;
    }

    private static final class RecordingAgentRunService extends AgentRunService {

        private int approveCount;
        private String lastRequestId;

        @Override
        public void approveRuntimeLocal(String requestId, String resumeToken, BigInteger operatorId, String userId) {
            approveCount++;
            lastRequestId = requestId;
        }
    }

    private static class RecordingCommandResultRegistry extends AgentRuntimeCommandResultRegistry {

        private String lastSuccessCommandId;
        String lastFailureCommandId;
        private String lastFailureMessage;

        private RecordingCommandResultRegistry() {
            super(null, null, null);
        }

        @Override
        public void markSuccess(String commandId) {
            lastSuccessCommandId = commandId;
        }

        @Override
        public void markFailure(String commandId, String message) {
            lastFailureCommandId = commandId;
            lastFailureMessage = message;
        }
    }

    private static final class FailingAgentRunService extends AgentRunService {

        @Override
        public void approveRuntimeLocal(String requestId, String resumeToken, BigInteger operatorId, String userId) {
            throw new RuntimeException("runtime missing");
        }
    }

    private static final class FailingSuccessResultRegistry extends RecordingCommandResultRegistry {

        @Override
        public void markSuccess(String commandId) {
            super.markSuccess(commandId);
            throw new RuntimeException("redis down");
        }
    }
}
