package tech.easyflow.agent.runtime.asynctool;

import com.easyagents.agent.runtime.tool.AgentToolContext;
import com.easyagents.agent.runtime.tool.asynctool.AsyncToolObserveRequest;
import com.easyagents.agent.runtime.tool.asynctool.AsyncToolResultRequest;
import com.easyagents.agent.runtime.tool.asynctool.AsyncToolSubmitResult;
import com.easyagents.agent.runtime.tool.asynctool.AsyncToolTaskStatus;
import com.easyagents.agent.runtime.tool.asynctool.AsyncToolTaskView;
import org.junit.Assert;
import org.junit.Test;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import tech.easyflow.agent.runtime.tool.AgentToolExecutionResult;
import tech.easyflow.agent.runtime.tool.PluginToolExecutor;
import tech.easyflow.agent.runtime.tool.WorkflowToolExecutor;
import tech.easyflow.ai.entity.PluginItem;
import tech.easyflow.ai.entity.Workflow;

import java.math.BigInteger;
import java.util.ArrayList;
import java.util.Comparator;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.UnaryOperator;

/**
 * Workflow 与 Plugin 异步子工具测试。
 */
public class WorkflowPluginAsyncSubToolsTest {

    /**
     * 验证 Workflow 异步子工具会把业务执行结果保留到任务视图。
     *
     * @throws Exception 等待后台执行超时时抛出
     */
    @Test
    public void workflowAsyncSubToolsShouldKeepBusinessResultInTaskView() throws Exception {
        ThreadPoolTaskExecutor executor = executor();
        try {
            Map<String, Object> businessResult = Map.of("workflowOutput", "ok");
            WorkflowAsyncSubTools subTools = new WorkflowAsyncSubTools(workflow(),
                    "workflow_demo",
                    "测试工作流",
                    new StubWorkflowToolExecutor(businessResult),
                    new InMemoryTaskStore(),
                    executor);

            AsyncToolTaskView view = submitAndResult(subTools);

            Assert.assertEquals(AsyncToolTaskStatus.SUCCEEDED, view.getStatus());
            Assert.assertEquals(businessResult, view.getResult());
            Assert.assertEquals("workflow-run-1", view.getPayload().get("businessExecutionId"));
        } finally {
            executor.shutdown();
        }
    }

    /**
     * 验证 Plugin 异步子工具会把业务执行结果保留到任务视图。
     *
     * @throws Exception 等待后台执行超时时抛出
     */
    @Test
    public void pluginAsyncSubToolsShouldKeepBusinessResultInTaskView() throws Exception {
        ThreadPoolTaskExecutor executor = executor();
        try {
            Map<String, Object> businessResult = Map.of("pluginOutput", List.of("a", "b"));
            PluginAsyncSubTools subTools = new PluginAsyncSubTools(pluginItem(),
                    "plugin_demo",
                    "测试插件",
                    new StubPluginToolExecutor(businessResult),
                    new InMemoryTaskStore(),
                    executor);

            AsyncToolTaskView view = submitAndResult(subTools);

            Assert.assertEquals(AsyncToolTaskStatus.SUCCEEDED, view.getStatus());
            Assert.assertEquals(businessResult, view.getResult());
        } finally {
            executor.shutdown();
        }
    }

    private AsyncToolTaskView submitAndResult(AbstractAgentAsyncSubTools subTools) throws Exception {
        AgentToolContext context = context();
        AsyncToolSubmitResult submitted = subTools.submit(Map.of("keyword", "hello"), context);
        waitTerminal(subTools, submitted.getTaskId(), context);
        AsyncToolResultRequest request = new AsyncToolResultRequest();
        request.setTaskId(submitted.getTaskId());
        return subTools.result(request, context);
    }

    private void waitTerminal(AbstractAgentAsyncSubTools subTools, String taskId, AgentToolContext context) throws Exception {
        long deadline = System.currentTimeMillis() + 3000L;
        AsyncToolTaskView view = observe(subTools, taskId, context);
        while (!Boolean.TRUE.equals(view.getTerminal()) && System.currentTimeMillis() < deadline) {
            Thread.sleep(20L);
            view = observe(subTools, taskId, context);
        }
        Assert.assertTrue("异步任务应在测试超时前完成", Boolean.TRUE.equals(view.getTerminal()));
    }

    private AsyncToolTaskView observe(AbstractAgentAsyncSubTools subTools, String taskId, AgentToolContext context) {
        AsyncToolObserveRequest request = new AsyncToolObserveRequest();
        request.setTaskId(taskId);
        return subTools.observe(request, context);
    }

    private AgentToolContext context() {
        AgentToolContext context = new AgentToolContext();
        context.setRequestId("request-1");
        context.setTraceId("trace-1");
        context.setSessionId("session-1");
        context.setAgentId("agent-1");
        context.setToolCallId("tool-call-1");
        return context;
    }

    private Workflow workflow() {
        Workflow workflow = new Workflow();
        workflow.setId(BigInteger.valueOf(101L));
        workflow.setTitle("测试工作流");
        return workflow;
    }

    private PluginItem pluginItem() {
        PluginItem pluginItem = new PluginItem();
        pluginItem.setId(BigInteger.valueOf(102L));
        pluginItem.setName("测试插件");
        return pluginItem;
    }

    private ThreadPoolTaskExecutor executor() {
        ThreadPoolTaskExecutor executor = new ThreadPoolTaskExecutor();
        executor.setCorePoolSize(1);
        executor.setMaxPoolSize(1);
        executor.setQueueCapacity(4);
        executor.setThreadNamePrefix("workflow-plugin-async-test-");
        executor.initialize();
        return executor;
    }

    private static final class StubWorkflowToolExecutor extends WorkflowToolExecutor {

        private final Map<String, Object> businessResult;

        private StubWorkflowToolExecutor(Map<String, Object> businessResult) {
            super(null);
            this.businessResult = businessResult;
        }

        @Override
        public AgentToolExecutionResult execute(Workflow workflow, Map<String, Object> arguments) {
            return new AgentToolExecutionResult(businessResult, "workflow-run-1");
        }
    }

    private static final class StubPluginToolExecutor extends PluginToolExecutor {

        private final Map<String, Object> businessResult;

        private StubPluginToolExecutor(Map<String, Object> businessResult) {
            this.businessResult = businessResult;
        }

        @Override
        public AgentToolExecutionResult execute(PluginItem pluginItem, Map<String, Object> arguments) {
            return new AgentToolExecutionResult(businessResult, null);
        }
    }

    private static final class InMemoryTaskStore implements AgentAsyncToolTaskStore {

        private final Map<String, AgentAsyncToolTaskRecord> records = new ConcurrentHashMap<>();

        @Override
        public void create(AgentAsyncToolTaskRecord record) {
            record.setSessionScopedKey(key(record.getSessionId(), record.getTaskId()));
            records.put(record.getSessionScopedKey(), record);
        }

        @Override
        public Optional<AgentAsyncToolTaskRecord> get(String sessionId, String taskId) {
            return Optional.ofNullable(records.get(key(sessionId, taskId)));
        }

        @Override
        public Optional<AgentAsyncToolTaskRecord> update(String sessionId,
                                                         String taskId,
                                                         UnaryOperator<AgentAsyncToolTaskRecord> updater) {
            AgentAsyncToolTaskRecord updated = records.computeIfPresent(key(sessionId, taskId),
                    (ignored, existing) -> updater == null ? existing : updater.apply(existing));
            return Optional.ofNullable(updated);
        }

        @Override
        public List<AgentAsyncToolTaskRecord> list(String sessionId, AsyncToolTaskStatus status) {
            List<AgentAsyncToolTaskRecord> result = new ArrayList<>();
            for (AgentAsyncToolTaskRecord record : records.values()) {
                if (sessionId.equals(record.getSessionId()) && (status == null || status == record.getStatus())) {
                    result.add(record);
                }
            }
            result.sort(Comparator.comparing(AgentAsyncToolTaskRecord::getCreatedAt).reversed());
            return result;
        }

        private String key(String sessionId, String taskId) {
            return sessionId + ":" + taskId;
        }
    }
}
