package com.easyagents.agent.runtime.hitl;

import com.easyagents.agent.runtime.AgentResumeRequest;
import com.easyagents.agent.runtime.AgentRuntimeException;
import org.junit.Assert;
import org.junit.Test;

import java.time.Instant;
import java.util.Map;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;

/**
 * 测试工具审批协调器。
 */
public class AgentToolApprovalCoordinatorTest {

    /**
     * 验证同批次全部调用批准后才签发逐调用执行授权。
     */
    @Test
    public void shouldAuthorizeBatchOnlyAfterAllCallsApproved() {
        AgentToolApprovalCoordinator coordinator = AgentToolApprovalCoordinator.enabled();
        AgentPendingState first = register(coordinator, "call-1", "search", Map.of("q", "first"), "batch-1");
        AgentPendingState second = register(coordinator, "call-2", "search", Map.of("q", "second"), "batch-1");

        AgentToolApprovalResolution firstResolution = coordinator.resolve(resume(first, true));

        Assert.assertEquals(AgentToolApprovalResolution.Status.WAITING, firstResolution.getStatus());
        Assert.assertEquals(1, firstResolution.getRemainingStates().size());
        assertAuthorizationRejected(coordinator, "call-1", "search", Map.of("q", "first"));

        AgentToolApprovalResolution secondResolution = coordinator.resolve(resume(second, true));

        Assert.assertEquals(AgentToolApprovalResolution.Status.READY, secondResolution.getStatus());
        coordinator.consumeExecutionAuthorization("call-1", "search", Map.of("q", "first"));
        coordinator.consumeExecutionAuthorization("call-2", "search", Map.of("q", "second"));
        assertAuthorizationRejected(coordinator, "call-1", "search", Map.of("q", "first"));
    }

    /**
     * 验证拒绝一个调用会关闭整个审批批次。
     */
    @Test
    public void shouldRejectWholeBatchWhenAnyCallRejected() {
        AgentToolApprovalCoordinator coordinator = AgentToolApprovalCoordinator.enabled();
        AgentPendingState first = register(coordinator, "call-1", "search", Map.of("q", "first"), "batch-1");
        AgentPendingState second = register(coordinator, "call-2", "search", Map.of("q", "second"), "batch-1");
        coordinator.resolve(resume(first, true));

        AgentResumeRequest rejection = resume(second, false);
        rejection.setRejectReason("not allowed");
        AgentToolApprovalResolution resolution = coordinator.resolve(rejection);

        Assert.assertEquals(AgentToolApprovalResolution.Status.REJECTED, resolution.getStatus());
        Assert.assertEquals("not allowed", resolution.getReason());
        assertAuthorizationRejected(coordinator, "call-1", "search", Map.of("q", "first"));
        try {
            coordinator.resolve(resume(first, true));
            Assert.fail("已消费的审批令牌不能重复使用");
        } catch (AgentRuntimeException expected) {
            Assert.assertTrue(expected.getMessage().contains("invalid"));
        }
    }

    /**
     * 验证过期令牌不能签发工具执行授权。
     */
    @Test
    public void shouldExpireApprovalBeforeResolution() {
        AgentToolApprovalCoordinator coordinator = AgentToolApprovalCoordinator.enabled();
        AgentPendingState expired = coordinator.register(
                "session-1",
                "agent-1",
                "call-expired",
                "search",
                "approve",
                Map.of("q", "expired"),
                Map.of(),
                Instant.now().minusSeconds(1),
                "batch-expired");

        AgentToolApprovalResolution resolution = coordinator.resolve(resume(expired, true));

        Assert.assertEquals(AgentToolApprovalResolution.Status.EXPIRED, resolution.getStatus());
        assertAuthorizationRejected(coordinator, "call-expired", "search", Map.of("q", "expired"));
    }

    /**
     * 验证并发重复点击同一令牌时最多一个请求可以成功消费。
     *
     * @throws Exception 并发任务执行失败时抛出
     */
    @Test
    public void shouldConsumeConcurrentDuplicateApprovalOnlyOnce() throws Exception {
        AgentToolApprovalCoordinator coordinator = AgentToolApprovalCoordinator.enabled();
        AgentPendingState pending = register(
                coordinator, "call-1", "search", Map.of("q", "easyflow"), "batch-1");
        CountDownLatch start = new CountDownLatch(1);
        ExecutorService executor = Executors.newFixedThreadPool(2);
        try {
            Future<Boolean> first = executor.submit(() -> resolveAfter(start, coordinator, pending));
            Future<Boolean> second = executor.submit(() -> resolveAfter(start, coordinator, pending));
            start.countDown();

            int successCount = (first.get() ? 1 : 0) + (second.get() ? 1 : 0);

            Assert.assertEquals(1, successCount);
            coordinator.consumeExecutionAuthorization(
                    "call-1", "search", Map.of("q", "easyflow"));
            assertAuthorizationRejected(
                    coordinator, "call-1", "search", Map.of("q", "easyflow"));
        } finally {
            executor.shutdownNow();
        }
    }

    /**
     * 验证工具名称或入参变化时批准凭证立即失效。
     */
    @Test
    public void shouldRejectExecutionWhenApprovedCallIsModified() {
        AgentToolApprovalCoordinator coordinator = AgentToolApprovalCoordinator.enabled();
        AgentPendingState pending = register(
                coordinator, "call-1", "search", Map.of("q", "easyflow"), "batch-1");
        coordinator.resolve(resume(pending, true));

        assertAuthorizationRejected(
                coordinator, "call-1", "search", Map.of("q", "modified"));
        assertAuthorizationRejected(
                coordinator, "call-1", "search", Map.of("q", "easyflow"));
    }

    /**
     * 验证跨节点受信任恢复仍需绑定明确的工具调用信息。
     */
    @Test
    public void shouldAuthorizeTrustedExecutionByToolCallIdentity() {
        AgentToolApprovalCoordinator coordinator = AgentToolApprovalCoordinator.enabled();
        AgentResumeRequest request = new AgentResumeRequest();
        AgentResumeToken token = new AgentResumeToken();
        token.setValue("persisted-token");
        request.setResumeToken(token);
        request.setApproved(true);
        request.setTrusted(true);
        request.setMetadata(Map.of(
                "toolCallId", "call-1",
                "toolName", "search",
                "toolInput", Map.of("q", "easyflow")));

        coordinator.authorizeTrustedExecution(request);

        coordinator.consumeExecutionAuthorization(
                "call-1", "search", Map.of("q", "easyflow"));
        assertAuthorizationRejected(
                coordinator, "call-1", "search", Map.of("q", "easyflow"));
    }

    /**
     * 验证受信任恢复可以签发当前 Turn 内可复用的 MCP 审批作用域。
     */
    @Test
    public void shouldAuthorizeTrustedMcpScope() {
        AgentToolApprovalCoordinator coordinator = AgentToolApprovalCoordinator.enabled();
        AgentResumeRequest request = new AgentResumeRequest();
        AgentResumeToken token = new AgentResumeToken();
        token.setValue("persisted-token");
        request.setResumeToken(token);
        request.setApproved(true);
        request.setTrusted(true);
        request.setMetadata(Map.of(
                "toolCallId", "call-1",
                "toolName", "mcp_101_search",
                "toolInput", Map.of("q", "easyflow"),
                "toolType", "MCP",
                "mcpId", "101"));

        coordinator.authorizeTrustedExecution(request);

        Assert.assertTrue(coordinator.isReusableApprovalGranted(
                Map.of("toolType", "MCP", "mcpId", "101")));
    }

    /**
     * 验证 MCP 批准可在当前 Turn 按稳定 mcpId 复用。
     */
    @Test
    public void shouldReuseApprovedMcpScopeWithinCurrentTurn() {
        AgentToolApprovalCoordinator coordinator = AgentToolApprovalCoordinator.enabled();
        Map<String, Object> mcpMetadata = Map.of("toolType", "MCP", "mcpId", "101");
        AgentPendingState pending = coordinator.register(
                "session-1",
                "agent-1",
                "call-resolve",
                "mcp_101_resolve_library_id",
                "approve",
                Map.of("libraryName", "AG-UI"),
                mcpMetadata,
                Instant.now().plusSeconds(60),
                "batch-mcp");

        coordinator.resolve(resume(pending, true));

        Assert.assertTrue(coordinator.isReusableApprovalGranted(mcpMetadata));
        Assert.assertFalse(coordinator.isReusableApprovalGranted(
                Map.of("toolType", "MCP", "mcpId", "102")));
        Assert.assertFalse(coordinator.isReusableApprovalGranted(
                Map.of("toolType", "MCP", "mcpName", "context7")));

        coordinator.clearReusableApprovalScopes();

        Assert.assertFalse(coordinator.isReusableApprovalGranted(mcpMetadata));
    }

    /**
     * 验证受控 Shell 脚本只能按受信任内容摘要在当前 Turn 复用审批。
     */
    @Test
    public void shouldReuseApprovedShellScriptScopeWithinCurrentTurn() {
        AgentToolApprovalCoordinator coordinator = AgentToolApprovalCoordinator.enabled();
        Map<String, Object> scriptMetadata = Map.of(
                "operateTool", true,
                "operateToolType", "SHELL",
                "approvalScope", "SHELL_SCRIPT:abc123");
        AgentPendingState pending = coordinator.register(
                "session-1", "agent-1", "call-script", "execute_shell_command", "approve",
                Map.of("command", "python3 report.py"), scriptMetadata,
                Instant.now().plusSeconds(60), "batch-script");

        coordinator.resolve(resume(pending, true));

        Assert.assertTrue(coordinator.isReusableApprovalGranted(scriptMetadata));
        Assert.assertFalse(coordinator.isReusableApprovalGranted(Map.of(
                "operateTool", true,
                "operateToolType", "SHELL",
                "approvalScope", "SHELL_SCRIPT:changed")));
        Assert.assertNull(coordinator.reusableApprovalScope(Map.of(
                "approvalScope", "SHELL_SCRIPT:abc123")));
    }

    /**
     * 验证跨节点受信任恢复可恢复 Shell 脚本内容摘要作用域。
     */
    @Test
    public void shouldAuthorizeTrustedShellScriptScope() {
        AgentToolApprovalCoordinator coordinator = AgentToolApprovalCoordinator.enabled();
        AgentResumeRequest request = new AgentResumeRequest();
        AgentResumeToken token = new AgentResumeToken();
        token.setValue("persisted-token");
        request.setResumeToken(token);
        request.setApproved(true);
        request.setTrusted(true);
        request.setMetadata(Map.of(
                "toolCallId", "call-script",
                "toolName", "execute_shell_command",
                "toolInput", Map.of("command", "node report.mjs"),
                "operateTool", true,
                "operateToolType", "SHELL",
                "approvalScope", "SHELL_SCRIPT:def456"));

        coordinator.authorizeTrustedExecution(request);

        Assert.assertTrue(coordinator.isReusableApprovalGranted(Map.of(
                "operateTool", true,
                "operateToolType", "SHELL",
                "approvalScope", "SHELL_SCRIPT:def456")));
    }

    /**
     * 验证拒绝和过期不会产生可复用 MCP 批准。
     */
    @Test
    public void shouldNotReuseRejectedOrExpiredMcpApproval() {
        AgentToolApprovalCoordinator coordinator = AgentToolApprovalCoordinator.enabled();
        Map<String, Object> rejectedMetadata = Map.of("toolType", "MCP", "mcpId", "201");
        AgentPendingState rejected = coordinator.register(
                "session-1", "agent-1", "call-rejected", "mcp_rejected", "approve",
                Map.of(), rejectedMetadata, Instant.now().plusSeconds(60), "batch-rejected");
        coordinator.resolve(resume(rejected, false));

        Map<String, Object> expiredMetadata = Map.of("toolType", "MCP", "mcpId", "202");
        AgentPendingState expired = coordinator.register(
                "session-1", "agent-1", "call-expired-mcp", "mcp_expired", "approve",
                Map.of(), expiredMetadata, Instant.now().minusSeconds(1), "batch-expired-mcp");
        coordinator.resolve(resume(expired, true));

        Assert.assertFalse(coordinator.isReusableApprovalGranted(rejectedMetadata));
        Assert.assertFalse(coordinator.isReusableApprovalGranted(expiredMetadata));
    }

    /**
     * 验证同一 toolCallId 不能被重新绑定到不同工具内容。
     */
    @Test
    public void shouldRejectDuplicateToolCallIdWithDifferentInput() {
        AgentToolApprovalCoordinator coordinator = AgentToolApprovalCoordinator.enabled();
        register(coordinator, "call-1", "search", Map.of("q", "easyflow"), "batch-1");

        try {
            register(coordinator, "call-1", "search", Map.of("q", "modified"), "batch-2");
            Assert.fail("重复 toolCallId 不能绑定不同入参");
        } catch (AgentRuntimeException expected) {
            Assert.assertTrue(expected.getMessage().contains("Duplicate toolCallId"));
        }
    }

    /**
     * 注册测试用审批状态。
     *
     * @param coordinator 审批协调器
     * @param toolCallId 工具调用ID
     * @param toolName 工具名称
     * @param toolInput 工具入参
     * @param batchId 审批批次ID
     * @return 待审批状态
     */
    private AgentPendingState register(AgentToolApprovalCoordinator coordinator,
                                       String toolCallId,
                                       String toolName,
                                       Map<String, Object> toolInput,
                                       String batchId) {
        return coordinator.register(
                "session-1",
                "agent-1",
                toolCallId,
                toolName,
                "approve",
                toolInput,
                Map.of(),
                Instant.now().plusSeconds(60),
                batchId);
    }

    /**
     * 创建测试用恢复请求。
     *
     * @param state 待审批状态
     * @param approved 是否批准
     * @return 恢复请求
     */
    private AgentResumeRequest resume(AgentPendingState state, boolean approved) {
        AgentResumeRequest request = new AgentResumeRequest();
        request.setResumeToken(state.getResumeToken());
        request.setApproved(approved);
        return request;
    }

    /**
     * 等待并发起跑信号后消费审批令牌。
     *
     * @param start 起跑信号
     * @param coordinator 审批协调器
     * @param pending 待审批状态
     * @return 成功消费时为 true
     * @throws InterruptedException 等待被中断时抛出
     */
    private boolean resolveAfter(CountDownLatch start,
                                 AgentToolApprovalCoordinator coordinator,
                                 AgentPendingState pending) throws InterruptedException {
        start.await();
        try {
            coordinator.resolve(resume(pending, true));
            return true;
        } catch (AgentRuntimeException expected) {
            return false;
        }
    }

    /**
     * 断言工具执行授权不可用。
     *
     * @param coordinator 审批协调器
     * @param toolCallId 工具调用ID
     * @param toolName 工具名称
     * @param toolInput 工具入参
     */
    private void assertAuthorizationRejected(AgentToolApprovalCoordinator coordinator,
                                             String toolCallId,
                                             String toolName,
                                             Map<String, Object> toolInput) {
        try {
            coordinator.consumeExecutionAuthorization(toolCallId, toolName, toolInput);
            Assert.fail("未授权或已消费的工具调用必须被拒绝");
        } catch (AgentToolApprovalRejectedException expected) {
            Assert.assertNotNull(expected.getMessage());
        }
    }
}
