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);
    }

    /**
     * 验证审批过期命令走专用运行时入口，不重复消费持久化 pending。
     *
     * @throws Exception 消息序列化异常
     */
    @Test
    public void consumerShouldHandleExpireCommand() throws Exception {
        AgentRuntimeProperties properties = new AgentRuntimeProperties();
        properties.setInstanceId("node-a");
        MQProperties mqProperties = new MQProperties();
        RecordingAgentRunService service = new RecordingAgentRunService();
        RecordingCommandResultRegistry resultRegistry = new RecordingCommandResultRegistry();
        AgentRuntimeCommandConsumer consumer =
                new AgentRuntimeCommandConsumer(new ObjectMapper(), properties, mqProperties, service, resultRegistry);
        AgentRuntimeCommandMessage command = command("cmd-expire", "node-a");
        command.setAction(AgentRuntimeCommandAction.EXPIRE);
        command.setReason("expired");

        consumer.handle(List.of(message(command)));

        Assert.assertEquals(1, service.expireCount);
        Assert.assertEquals("request-cmd-expire", service.lastRequestId);
        Assert.assertEquals("expired", service.lastReason);
        Assert.assertEquals("cmd-expire", resultRegistry.lastSuccessCommandId);
    }

    /**
     * 验证跨节点 AG-UI 审批命令在 owner 节点走脱敏决议入口。
     *
     * @throws Exception 消息序列化异常
     */
    @Test
    public void consumerShouldRouteAguiApprovalToResolvedEventPath() throws Exception {
        AgentRuntimeProperties properties = new AgentRuntimeProperties();
        properties.setInstanceId("node-a");
        RecordingAgentRunService service = new RecordingAgentRunService();
        RecordingCommandResultRegistry resultRegistry = new RecordingCommandResultRegistry();
        AgentRuntimeCommandConsumer consumer = new AgentRuntimeCommandConsumer(
                new ObjectMapper(), properties, new MQProperties(), service, resultRegistry);
        AgentRuntimeCommandMessage command = command("cmd-agui", "node-a");
        command.setApprovalId("approval-public");

        consumer.handle(List.of(message(command)));

        Assert.assertEquals(1, service.aguiApproveCount);
        Assert.assertEquals("approval-public", service.lastApprovalId);
        Assert.assertEquals(0, service.approveCount);
    }

    /**
     * 验证 Agent 集群取消命令只取消目标节点的对应 Agent 运行。
     *
     * @throws Exception 消息序列化异常
     */
    @Test
    public void consumerShouldHandleCancelAgentCommand() throws Exception {
        AgentRuntimeProperties properties = new AgentRuntimeProperties();
        properties.setInstanceId("node-a");
        RecordingAgentRunService service = new RecordingAgentRunService();
        RecordingCommandResultRegistry resultRegistry = new RecordingCommandResultRegistry();
        AgentRuntimeCommandConsumer consumer = new AgentRuntimeCommandConsumer(
                new ObjectMapper(),
                properties,
                new MQProperties(),
                service,
                resultRegistry
        );
        AgentRuntimeCommandMessage command = command("cmd-cancel", "node-a");
        command.setAction(AgentRuntimeCommandAction.CANCEL_AGENT);
        command.setAgentId("1001");

        consumer.handle(List.of(message(command)));

        Assert.assertEquals("1001", service.lastCancelledAgentId);
        Assert.assertEquals("cmd-cancel", resultRegistry.lastSuccessCommandId);
    }

    /**
     * 验证跨节点精确取消命令保留请求、用户和原因。
     *
     * @throws Exception 消息序列化异常
     */
    @Test
    public void consumerShouldHandleCancelRunCommand() throws Exception {
        AgentRuntimeProperties properties = new AgentRuntimeProperties();
        properties.setInstanceId("node-a");
        RecordingAgentRunService service = new RecordingAgentRunService();
        RecordingCommandResultRegistry resultRegistry = new RecordingCommandResultRegistry();
        AgentRuntimeCommandConsumer consumer = new AgentRuntimeCommandConsumer(
                new ObjectMapper(), properties, new MQProperties(), service, resultRegistry);
        AgentRuntimeCommandMessage command = command("cmd-cancel-run", "node-a");
        command.setAction(AgentRuntimeCommandAction.CANCEL_RUN);
        command.setReason("用户停止");

        consumer.handle(List.of(message(command)));

        Assert.assertEquals("request-cmd-cancel-run", service.lastCancelledRequestId);
        Assert.assertEquals("1", service.lastCancelledUserId);
        Assert.assertEquals("用户停止", service.lastReason);
        Assert.assertEquals("cmd-cancel-run", resultRegistry.lastSuccessCommandId);
    }

    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 int aguiApproveCount;
        private int expireCount;
        private String lastRequestId;
        private String lastReason;
        private String lastCancelledAgentId;
        private String lastCancelledRequestId;
        private String lastCancelledUserId;
        private String lastApprovalId;

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

        @Override
        public void approveAguiRuntimeLocal(String requestId,
                                            String resumeToken,
                                            String approvalId,
                                            BigInteger operatorId,
                                            String userId) {
            aguiApproveCount++;
            lastRequestId = requestId;
            lastApprovalId = approvalId;
        }

        @Override
        public void expireApprovalLocal(String requestId, String resumeToken, String reason) {
            expireCount++;
            lastRequestId = requestId;
            lastReason = reason;
        }

        @Override
        public void cancelAgentLocal(String agentId) {
            lastCancelledAgentId = agentId;
        }

        @Override
        public void cancelRunLocal(String requestId, String userId, String reason) {
            lastCancelledRequestId = requestId;
            lastCancelledUserId = userId;
            lastReason = reason;
        }
    }

    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");
        }
    }
}
