package tech.easyflow.agent.runtime;

import com.easyagents.agent.runtime.AgentInitRequest;
import com.easyagents.agent.runtime.AgentResumeRequest;
import com.easyagents.agent.runtime.AgentRuntime;
import com.easyagents.agent.runtime.event.AgentRuntimeEvent;
import org.junit.Assert;
import org.junit.Test;
import reactor.core.publisher.Flux;
import tech.easyflow.common.web.exceptions.BusinessException;
import tech.easyflow.core.runtime.ChatAssistantAccumulator;
import tech.easyflow.core.runtime.ChatRuntimeContext;

import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;

/**
 * Agent 运行态注册表测试。
 */
public class AgentRunRegistryTest {

    /**
     * 验证批准请求会恢复当前运行时。
     */
    @Test
    public void approveShouldResumeRuntime() {
        AgentRunRegistry registry = new AgentRunRegistry();
        CapturingRuntime runtime = new CapturingRuntime();

        registry.register(context("request-1", "session-1", "user-1", runtime));
        registry.registerResumeToken("request-1", "token-1");
        registry.approve("request-1", "token-1", "user-1");

        AgentResumeRequest request = runtime.resumeRequest.get();
        Assert.assertNotNull(request);
        Assert.assertTrue(request.isApproved());
        Assert.assertEquals("token-1", request.getResumeToken().getValue());
    }

    /**
     * 验证拒绝请求支持通过恢复令牌反查运行时。
     */
    @Test
    public void rejectShouldResolveRequestByResumeToken() {
        AgentRunRegistry registry = new AgentRunRegistry();
        CapturingRuntime runtime = new CapturingRuntime();

        registry.register(context("request-2", "session-2", "user-1", runtime));
        registry.registerResumeToken("request-2", "token-2");
        registry.reject(null, "token-2", "user-1", "denied");

        AgentResumeRequest request = runtime.resumeRequest.get();
        Assert.assertNotNull(request);
        Assert.assertFalse(request.isApproved());
        Assert.assertEquals("denied", request.getRejectReason());
    }

    /**
     * 验证运行结束后清理运行态与恢复令牌索引。
     */
    @Test
    public void removeShouldClearRuntimeAndResumeTokens() {
        AgentRunRegistry registry = new AgentRunRegistry();

        registry.register(context("request-3", "session-3", "user-1", new CapturingRuntime()));
        registry.registerResumeToken("request-3", "token-3");
        registry.remove("request-3");

        Assert.assertThrows(BusinessException.class, () -> registry.approve(null, "token-3", "user-1"));
    }

    /**
     * 验证运行审批只能由运行发起人处理。
     */
    @Test
    public void approveShouldRejectDifferentOwner() {
        AgentRunRegistry registry = new AgentRunRegistry();
        CapturingRuntime runtime = new CapturingRuntime();

        registry.register(context("request-4", "session-4", "user-1", runtime));
        registry.registerResumeToken("request-4", "token-4");

        Assert.assertThrows(BusinessException.class, () -> registry.approve("request-4", "token-4", "user-2"));
        Assert.assertNull(runtime.resumeRequest.get());

        registry.approve("request-4", "token-4", "user-1");
        Assert.assertTrue(runtime.resumeRequest.get().isApproved());
    }

    /**
     * 验证传入 requestId 时仍会校验恢复令牌归属，避免错误令牌打断挂起运行。
     */
    @Test
    public void approveShouldRejectTokenNotBelongingToRequest() {
        AgentRunRegistry registry = new AgentRunRegistry();
        CapturingRuntime runtime = new CapturingRuntime();

        registry.register(context("request-5", "session-5", "user-1", runtime));

        Assert.assertThrows(BusinessException.class,
                () -> registry.approve("request-5", "wrong-token", "user-1"));
        Assert.assertNull(runtime.resumeRequest.get());
    }

    /**
     * 验证同一会话同一时刻只允许一个运行态。
     */
    @Test
    public void registerShouldRejectActiveRunInSameSession() {
        AgentRunRegistry registry = new AgentRunRegistry();

        registry.register(context("request-6", "session-6", "user-1", new CapturingRuntime()));

        Assert.assertThrows(BusinessException.class,
                () -> registry.register(context("request-7", "session-6", "user-1", new CapturingRuntime())));
    }

    /**
     * 验证按会话清理会取消当前运行并释放同会话运行锁。
     */
    @Test
    public void cancelSessionShouldRemoveActiveRun() {
        AgentRunRegistry registry = new AgentRunRegistry();
        AgentRunRegistry.AgentRunContext context = context("request-8", "session-8", "user-1", new CapturingRuntime());

        registry.register(context);
        registry.cancelSession("session-8");

        Assert.assertNull(registry.get("request-8"));
        registry.register(context("request-9", "session-8", "user-1", new CapturingRuntime()));
        Assert.assertNotNull(registry.get("request-9"));
    }

    /**
     * 验证按会话清理时会校验运行归属。
     */
    @Test
    public void cancelSessionShouldRejectDifferentOwner() {
        AgentRunRegistry registry = new AgentRunRegistry();

        registry.register(context("request-10", "session-10", "user-1", new CapturingRuntime()));

        Assert.assertThrows(BusinessException.class, () -> registry.cancelSession("session-10", "user-2"));
        Assert.assertNotNull(registry.get("request-10"));
    }

    private AgentRunRegistry.AgentRunContext context(String requestId,
                                                     String sessionId,
                                                     String userId,
                                                     AgentRuntime runtime) {
        return new AgentRunRegistry.AgentRunContext(
                requestId,
                sessionId,
                runtime,
                null,
                new ChatRuntimeContext(),
                new StringBuilder(),
                new ChatAssistantAccumulator(),
                new AtomicBoolean(false),
                false,
                new AgentRunRegistry.RunOwner("agent-1", sessionId, userId),
                null,
                event -> {
                },
                error -> {
                },
                () -> {
                }
        );
    }

    private static final class CapturingRuntime implements AgentRuntime {

        private final AtomicReference<AgentResumeRequest> resumeRequest = new AtomicReference<>();

        @Override
        public void init(AgentInitRequest request) {
            // 测试桩无需初始化。
        }

        @Override
        public Flux<AgentRuntimeEvent> stream(com.easyagents.agent.runtime.message.AgentMessage userMessage) {
            return Flux.empty();
        }

        @Override
        public Flux<AgentRuntimeEvent> resume(AgentResumeRequest request) {
            resumeRequest.set(request);
            return Flux.empty();
        }
    }
}
