package tech.easyflow.agent.runtime.asynctool;

import com.easyagents.agent.runtime.tool.AgentToolContext;
import com.easyagents.agent.runtime.tool.asynctool.AsyncToolCancelRequest;
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 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;

/**
 * EasyFlow 异步业务工具基类测试。
 */
public class AbstractAgentAsyncSubToolsTest {

    /**
     * 验证 submit、observe、result 与 list 的基础任务生命周期。
     *
     * @throws Exception 等待后台执行超时时抛出
     */
    @Test
    public void asyncSubToolsShouldSubmitObserveResultAndListCurrentSessionTasks() throws Exception {
        ThreadPoolTaskExecutor executor = executor();
        try {
            InMemoryTaskStore store = new InMemoryTaskStore();
            TestAsyncSubTools subTools = new TestAsyncSubTools(store, executor);
            AgentToolContext context = context("session-a");

            AsyncToolSubmitResult submitted = subTools.submit(Map.of("keyword", "hello"), context);
            Assert.assertEquals(AsyncToolTaskStatus.PENDING, submitted.getStatus());
            Assert.assertTrue(submitted.getTaskId().startsWith("async_"));

            AsyncToolTaskView completed = waitTerminal(subTools, submitted.getTaskId(), context);
            Assert.assertEquals(AsyncToolTaskStatus.SUCCEEDED, completed.getStatus());
            Assert.assertEquals(Map.of("echo", "hello"), completed.getResult());
            Assert.assertTrue(completed.getNextCursor() >= 2L);

            AsyncToolResultRequest resultRequest = new AsyncToolResultRequest();
            resultRequest.setTaskId(submitted.getTaskId());
            resultRequest.setCursor(1L);
            AsyncToolTaskView result = subTools.result(resultRequest, context);
            Assert.assertEquals(AsyncToolTaskStatus.SUCCEEDED, result.getStatus());
            Assert.assertEquals(Map.of("echo", "hello"), result.getResult());
            Assert.assertFalse(result.getEvents().isEmpty());

            Assert.assertEquals(1, subTools.list(null, context).getTasks().size());
            Assert.assertTrue(subTools.list(null, context("session-b")).getTasks().isEmpty());

            AsyncToolTaskView crossedSessionView = observe(subTools, submitted.getTaskId(), context("session-b"));
            Assert.assertEquals(AsyncToolTaskStatus.FAILED, crossedSessionView.getStatus());
            Assert.assertEquals("TASK_NOT_FOUND", crossedSessionView.getErrorType());
        } finally {
            executor.shutdown();
        }
    }

    /**
     * 验证首版取消语义返回明确失败结果。
     */
    @Test
    public void cancelShouldReturnUnsupportedFailure() {
        TestAsyncSubTools subTools = new TestAsyncSubTools(new InMemoryTaskStore(), executor());
        AsyncToolCancelRequest request = new AsyncToolCancelRequest();
        request.setTaskId("task-1");

        var result = subTools.cancel(request, context("session-a"));

        Assert.assertEquals(AsyncToolTaskStatus.FAILED, result.getStatus());
        Assert.assertEquals("不支持取消", result.getMessage());
    }

    private AsyncToolTaskView waitTerminal(TestAsyncSubTools 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()));
        return view;
    }

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

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

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

    private static final class TestAsyncSubTools extends AbstractAgentAsyncSubTools {

        private TestAsyncSubTools(AgentAsyncToolTaskStore taskStore, ThreadPoolTaskExecutor taskExecutor) {
            super(taskStore, taskExecutor);
        }

        @Override
        protected String toolType() {
            return "PLUGIN";
        }

        @Override
        protected String toolName() {
            return "test_tool";
        }

        @Override
        protected String displayName() {
            return "测试工具";
        }

        @Override
        protected String businessId() {
            return "business-1";
        }

        @Override
        protected AgentToolExecutionResult executeBusiness(Map<String, Object> arguments) {
            return new AgentToolExecutionResult(Map.of("echo", arguments.get("keyword")), "business-run-1");
        }
    }

    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) {
            String key = key(sessionId, taskId);
            AgentAsyncToolTaskRecord updated = records.computeIfPresent(key,
                    (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;
        }
    }
}
