package com.easyagents.agent.runtime.tool.asynctool;

import com.easyagents.agent.runtime.event.AgentRuntimeEvent;
import com.easyagents.agent.runtime.event.AgentRuntimeEventType;
import com.easyagents.agent.runtime.hitl.AgentToolApprovalRequest;
import com.easyagents.agent.runtime.tool.AgentToolContext;
import com.easyagents.agent.runtime.tool.AgentToolInvoker;
import com.easyagents.agent.runtime.tool.AgentToolResult;
import com.easyagents.agent.runtime.tool.AgentToolSpec;
import org.junit.Assert;
import org.junit.Test;

import java.time.Duration;
import java.time.Instant;
import java.util.*;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.atomic.AtomicReference;

/**
 * 测试异步工具声明展开器。
 */
public class AsyncToolSpecExpanderTest {

    @Test
    public void shouldExpandFiveToolSpecsAndInvokers() {
        AsyncToolSpec spec = spec(new StubSubTools());
        AsyncToolSpecExpander expander = new AsyncToolSpecExpander();

        List<AgentToolSpec> specs = expander.expandSpecs(spec);
        Map<String, AgentToolInvoker> invokers = expander.expandInvokers(spec);

        Assert.assertEquals(List.of("demo_task_submit", "demo_task_observe", "demo_task_result",
                "demo_task_cancel", "demo_task_list"), specs.stream().map(AgentToolSpec::getName).toList());
        Assert.assertEquals(5, invokers.size());
        Assert.assertEquals("object", specs.get(1).getParametersSchema().get("type"));
        Assert.assertEquals("AsyncToolTaskView", specs.get(1).getOutputSchema().get("title"));
        Assert.assertTrue(specs.get(0).getDescription().contains("default entry point when the user asks to run this tool"));
        Assert.assertTrue(specs.get(1).getDescription().contains("Use immediately after submit"));
        Assert.assertTrue(specs.get(1).getDescription().contains("Do not ask the user for task_id immediately after submit"));
        Assert.assertTrue(specs.get(2).getDescription().contains("final result"));
        Assert.assertTrue(specs.get(3).getDescription().contains("only when the user explicitly asks to cancel"));
        Assert.assertTrue(specs.get(4).getDescription().contains("only when the user explicitly asks to list"));
    }

    @Test
    public void shouldKeepRuntimeMetadataWhenUserMetadataUsesReservedKeys() {
        AsyncToolSpec spec = spec(new StubSubTools());
        spec.setMetadata(Map.of("asyncTool", false, "asyncToolName", "user_name", "custom", "value"));
        AsyncToolSpecExpander expander = new AsyncToolSpecExpander();

        AgentToolSpec toolSpec = expander.expandSpecs(spec).get(0);
        AgentToolResult result = expander.expandInvokers(spec).get("demo_task_submit")
                .invoke(Map.of(), context(new ArrayList<>()));

        Assert.assertEquals(true, toolSpec.getMetadata().get("asyncTool"));
        Assert.assertEquals("demo_task", toolSpec.getMetadata().get("asyncToolName"));
        Assert.assertEquals("value", toolSpec.getMetadata().get("custom"));
        Assert.assertEquals(true, result.getMetadata().get("asyncTool"));
        Assert.assertEquals("demo_task", result.getMetadata().get("asyncToolName"));
    }

    @Test
    public void shouldApplyApprovalOnlyToSubmitTool() {
        AsyncToolSpec spec = spec(new StubSubTools());
        AgentToolApprovalRequest approvalRequest = new AgentToolApprovalRequest();
        approvalRequest.setApprovalPrompt("确认提交任务？");
        spec.setApprovalRequired(true);
        spec.setApprovalRequest(approvalRequest);

        List<AgentToolSpec> specs = new AsyncToolSpecExpander().expandSpecs(spec);

        Assert.assertTrue(specs.get(0).isApprovalRequired());
        Assert.assertEquals("确认提交任务？", specs.get(0).getApprovalRequest().getApprovalPrompt());
        Assert.assertFalse(specs.get(1).isApprovalRequired());
        Assert.assertFalse(specs.get(2).isApprovalRequired());
        Assert.assertFalse(specs.get(3).isApprovalRequired());
        Assert.assertFalse(specs.get(4).isApprovalRequired());
    }

    @Test
    public void shouldSubmitAndEmitSubmittedEvent() {
        List<AgentRuntimeEvent> events = new ArrayList<>();
        AgentToolResult result = invoke("demo_task_submit", Map.of("input", "hello"), new StubSubTools(), events);

        Assert.assertTrue(result.isSuccess());
        Assert.assertTrue(result.getModelContent().contains("task_id: task-1"));
        Assert.assertEquals("task-1", result.getMetadata().get("taskId"));
        Assert.assertEquals(AgentRuntimeEventType.ASYNC_TOOL_SUBMITTED, events.get(0).getEventType());
        Assert.assertEquals("task-1", events.get(0).getPayload().get("taskId"));
    }

    @Test
    public void shouldFailWhenSubmitDoesNotReturnTaskId() {
        List<AgentRuntimeEvent> events = new ArrayList<>();
        AgentToolResult result = invoke("demo_task_submit", Map.of(), new StubSubTools() {
            @Override
            public AsyncToolSubmitResult submit(Map<String, Object> arguments, AgentToolContext context) {
                AsyncToolSubmitResult submitResult = super.submit(arguments, context);
                submitResult.setTaskId(null);
                return submitResult;
            }
        }, events);

        Assert.assertFalse(result.isSuccess());
        Assert.assertTrue(result.getErrorMessage().contains("taskId"));
        Assert.assertEquals(AgentRuntimeEventType.ASYNC_TOOL_FAILED, events.get(0).getEventType());
    }

    @Test
    public void shouldObserveWithCursorLimitAndContext() {
        AtomicReference<AsyncToolObserveRequest> requestRef = new AtomicReference<>();
        AtomicReference<AgentToolContext> contextRef = new AtomicReference<>();
        StubSubTools subTools = new StubSubTools() {
            @Override
            public AsyncToolTaskView observe(AsyncToolObserveRequest request, AgentToolContext context) {
                requestRef.set(request);
                contextRef.set(context);
                return super.observe(request, context);
            }
        };
        List<AgentRuntimeEvent> events = new ArrayList<>();

        AgentToolResult result = invoke("demo_task_observe", Map.of("taskId", "task-1", "cursor", 7, "limit", 999),
                subTools, events);

        AsyncToolTaskView view = (AsyncToolTaskView) result.getDisplayContent();
        Assert.assertEquals(Long.valueOf(7), requestRef.get().getCursor());
        Assert.assertEquals(Integer.valueOf(100), requestRef.get().getLimit());
        Assert.assertEquals(Long.valueOf(8), view.getNextCursor());
        Assert.assertEquals("request-1", contextRef.get().getRequestId());
        Assert.assertEquals(AgentRuntimeEventType.ASYNC_TOOL_OBSERVED, events.get(0).getEventType());
    }

    @Test
    public void shouldWrapInvalidArgumentsAndEmitFailedEvent() {
        List<AgentRuntimeEvent> events = new ArrayList<>();

        AgentToolResult result = invoke("demo_task_observe", Map.of("taskId", "task-1", "cursor", "bad"),
                new StubSubTools(), events);

        Assert.assertFalse(result.isSuccess());
        Assert.assertEquals(AgentRuntimeEventType.ASYNC_TOOL_FAILED, events.get(0).getEventType());
        Assert.assertEquals("FAILED", result.getMetadata().get("status"));
    }

    @Test
    public void shouldReturnObservationFromResultWhenRunningAndFinalResultWhenSucceeded() {
        StubSubTools running = new StubSubTools();
        AgentToolResult runningResult = invoke("demo_task_result", Map.of("taskId", "task-1"), running, new ArrayList<>());
        AsyncToolTaskView runningView = (AsyncToolTaskView) runningResult.getDisplayContent();

        StubSubTools succeeded = new StubSubTools() {
            @Override
            public AsyncToolTaskView result(AsyncToolResultRequest request, AgentToolContext context) {
                AsyncToolTaskView view = super.result(request, context);
                view.setStatus(AsyncToolTaskStatus.SUCCEEDED);
                view.setResult("done");
                return view;
            }
        };
        AgentToolResult successResult = invoke("demo_task_result", Map.of("taskId", "task-1"), succeeded, new ArrayList<>());
        AsyncToolTaskView successView = (AsyncToolTaskView) successResult.getDisplayContent();

        Assert.assertEquals(AsyncToolTaskStatus.RUNNING, runningView.getStatus());
        Assert.assertEquals(AsyncToolTaskStatus.SUCCEEDED, successView.getStatus());
        Assert.assertEquals("done", successView.getResult());
        Assert.assertTrue(successView.getTerminal());
        Assert.assertTrue(successView.getResultAvailable());
        Assert.assertFalse(runningResult.getModelContent().contains("result:"));
        Assert.assertTrue(successResult.getModelContent().contains("result_available: true"));
        Assert.assertTrue(successResult.getModelContent().contains("result: done"));
    }

    @Test
    public void shouldExposeCompletedObservationResultToModelContentAsJson() {
        List<AgentRuntimeEvent> events = new ArrayList<>();
        LinkedHashMap<String, Object> businessResult = new LinkedHashMap<>();
        businessResult.put("answer", "42");
        businessResult.put("items", List.of(1, 2));
        StubSubTools succeeded = new StubSubTools() {
            @Override
            public AsyncToolTaskView observe(AsyncToolObserveRequest request, AgentToolContext context) {
                AsyncToolTaskView view = super.observe(request, context);
                view.setStatus(AsyncToolTaskStatus.SUCCEEDED);
                view.setResult(businessResult);
                return view;
            }
        };

        AgentToolResult result = invoke("demo_task_observe", Map.of("taskId", "task-1"), succeeded, events);
        AsyncToolTaskView view = (AsyncToolTaskView) result.getDisplayContent();

        Assert.assertEquals(businessResult, view.getResult());
        Assert.assertTrue(result.getModelContent().contains("result_available: true"));
        Assert.assertTrue(result.getModelContent().contains("result: {\"answer\":\"42\",\"items\":[1,2]}"));
        Assert.assertEquals(Boolean.TRUE, events.get(0).getPayload().get("resultAvailable"));
        Assert.assertFalse(events.get(0).getPayload().containsKey("result"));
    }

    @Test
    public void shouldCancelAndRepresentUnsupportedCancelAsFailure() {
        AgentToolResult success = invoke("demo_task_cancel", Map.of("taskId", "task-1"), new StubSubTools(), new ArrayList<>());
        AgentToolResult failure = invoke("demo_task_cancel", Map.of("taskId", "task-1"), new StubSubTools() {
            @Override
            public AsyncToolCancelResult cancel(AsyncToolCancelRequest request, AgentToolContext context) {
                AsyncToolCancelResult result = super.cancel(request, context);
                result.setErrorMessage("不支持取消");
                return result;
            }
        }, new ArrayList<>());

        Assert.assertTrue(success.isSuccess());
        Assert.assertFalse(failure.isSuccess());
        Assert.assertEquals("不支持取消", failure.getErrorMessage());
    }

    @Test
    public void shouldListTasksWithoutPagination() {
        AgentToolResult result = invoke("demo_task_list", Map.of("status", "running"), new StubSubTools(), new ArrayList<>());
        AsyncToolTaskListResult list = (AsyncToolTaskListResult) result.getDisplayContent();

        Assert.assertTrue(result.isSuccess());
        Assert.assertEquals(1, list.getTasks().size());
        Assert.assertFalse(Arrays.stream(AsyncToolTaskListResult.class.getDeclaredFields())
                .anyMatch(field -> field.getName().contains("PageToken")));
    }

    @Test
    public void shouldWrapExceptionAndEmitFailedEvent() {
        List<AgentRuntimeEvent> events = new ArrayList<>();
        AgentToolResult result = invoke("demo_task_observe", Map.of("taskId", "task-1"), new StubSubTools() {
            @Override
            public AsyncToolTaskView observe(AsyncToolObserveRequest request, AgentToolContext context) {
                throw new IllegalStateException("boom");
            }
        }, events);

        Assert.assertFalse(result.isSuccess());
        Assert.assertEquals("boom", result.getErrorMessage());
        Assert.assertEquals(AgentRuntimeEventType.ASYNC_TOOL_FAILED, events.get(0).getEventType());
    }

    @Test
    public void shouldTimeoutAndEmitFailedEvent() {
        List<AgentRuntimeEvent> events = new ArrayList<>();
        AsyncToolOptions options = AsyncToolOptions.defaults();
        options.setObserveTimeout(Duration.ofMillis(20));
        StubSubTools subTools = new StubSubTools() {
            @Override
            public AsyncToolTaskView observe(AsyncToolObserveRequest request, AgentToolContext context) {
                try {
                    Thread.sleep(200);
                } catch (InterruptedException ignored) {
                    Thread.currentThread().interrupt();
                }
                return super.observe(request, context);
            }
        };
        AgentToolResult result = invoke("demo_task_observe", Map.of("taskId", "task-1"), subTools, events, options);

        Assert.assertFalse(result.isSuccess());
        Assert.assertTrue(result.getErrorMessage().contains("timed out"));
        Assert.assertEquals(AgentRuntimeEventType.ASYNC_TOOL_FAILED, events.get(0).getEventType());
    }

    @Test
    public void shouldNotPutLargePayloadIntoModelContentOrEventPayload() {
        List<AgentRuntimeEvent> events = new ArrayList<>();
        AgentToolResult result = invoke("demo_task_submit", Map.of(), new StubSubTools() {
            @Override
            public AsyncToolSubmitResult submit(Map<String, Object> arguments, AgentToolContext context) {
                AsyncToolSubmitResult result = super.submit(arguments, context);
                result.getPayload().put("large", "x".repeat(5000));
                return result;
            }
        }, events);

        Assert.assertFalse(result.getModelContent().contains("xxxxx"));
        Assert.assertFalse(events.get(0).getPayload().containsKey("large"));
        Assert.assertTrue(((AsyncToolSubmitResult) result.getDisplayContent()).getPayload().containsKey("large"));
    }

    @Test
    public void shouldTrimLargeResultForModelContentAndKeepFullDisplayContent() {
        List<AgentRuntimeEvent> events = new ArrayList<>();
        AsyncToolOptions options = AsyncToolOptions.defaults();
        options.setMaxModelContentLength(160);
        String largeResult = "x".repeat(5000);
        StubSubTools succeeded = new StubSubTools() {
            @Override
            public AsyncToolTaskView result(AsyncToolResultRequest request, AgentToolContext context) {
                AsyncToolTaskView view = super.result(request, context);
                view.setStatus(AsyncToolTaskStatus.SUCCEEDED);
                view.setResult(Map.of("large", largeResult));
                return view;
            }
        };

        AgentToolResult result = invoke("demo_task_result", Map.of("taskId", "task-1"), succeeded, events, options);
        AsyncToolTaskView view = (AsyncToolTaskView) result.getDisplayContent();

        Assert.assertTrue(result.getModelContent().length() <= options.getMaxModelContentLength() + 3);
        Assert.assertTrue(result.getModelContent().contains("result_available: true"));
        Assert.assertFalse(result.getModelContent().contains(largeResult));
        Assert.assertEquals(Map.of("large", largeResult), view.getResult());
        Assert.assertEquals(Boolean.TRUE, events.get(0).getPayload().get("resultAvailable"));
        Assert.assertFalse(events.get(0).getPayload().containsKey("result"));
    }

    private AgentToolResult invoke(String toolName,
                                   Map<String, Object> arguments,
                                   AsyncSubTools subTools,
                                   List<AgentRuntimeEvent> events) {
        return invoke(toolName, arguments, subTools, events, AsyncToolOptions.defaults());
    }

    private AgentToolResult invoke(String toolName,
                                   Map<String, Object> arguments,
                                   AsyncSubTools subTools,
                                   List<AgentRuntimeEvent> events,
                                   AsyncToolOptions options) {
        AsyncToolSpec spec = spec(subTools);
        spec.setOptions(options);
        AsyncToolSpecExpander expander = new AsyncToolSpecExpander(Executors.newCachedThreadPool());
        AgentToolContext context = context(events);
        return expander.expandInvokers(spec).get(toolName).invoke(arguments, context);
    }

    private AsyncToolSpec spec(AsyncSubTools subTools) {
        AsyncToolSpec spec = new AsyncToolSpec();
        spec.setName("demo_task");
        spec.setDescription("Demo async task.");
        spec.setSubTools(subTools);
        spec.setSubmitParametersSchema(Map.of("type", "object", "properties", Map.of("input", Map.of("type", "string"))));
        return spec;
    }

    private AgentToolContext context(List<AgentRuntimeEvent> events) {
        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");
        context.setEventEmitter(events::add);
        return context;
    }

    private static class StubSubTools implements AsyncSubTools {

        @Override
        public AsyncToolSubmitResult submit(Map<String, Object> arguments, AgentToolContext context) {
            AsyncToolSubmitResult result = new AsyncToolSubmitResult();
            result.setTaskId("task-1");
            result.setStatus(AsyncToolTaskStatus.RUNNING);
            result.setCursor(0L);
            result.setSummary("submitted");
            return result;
        }

        @Override
        public AsyncToolTaskView observe(AsyncToolObserveRequest request, AgentToolContext context) {
            AsyncToolTaskView view = new AsyncToolTaskView();
            view.setTaskId(request.getTaskId());
            view.setStatus(AsyncToolTaskStatus.RUNNING);
            view.setCursor(request.getCursor());
            view.setNextCursor(request.getCursor() == null ? 1L : request.getCursor() + 1);
            view.setSummary("running");
            AsyncToolTaskEvent event = new AsyncToolTaskEvent();
            event.setSequence(view.getNextCursor());
            event.setType("TASK_LOG");
            event.setText("progress");
            event.setCreatedAt(Instant.now());
            view.setEvents(List.of(event));
            return view;
        }

        @Override
        public AsyncToolTaskView result(AsyncToolResultRequest request, AgentToolContext context) {
            AsyncToolObserveRequest observeRequest = new AsyncToolObserveRequest();
            observeRequest.setTaskId(request.getTaskId());
            observeRequest.setCursor(request.getCursor());
            observeRequest.setLimit(request.getLimit());
            return observe(observeRequest, context);
        }

        @Override
        public AsyncToolCancelResult cancel(AsyncToolCancelRequest request, AgentToolContext context) {
            AsyncToolCancelResult result = new AsyncToolCancelResult();
            result.setTaskId(request.getTaskId());
            result.setStatus(AsyncToolTaskStatus.CANCELLING);
            result.setMessage("cancelling");
            return result;
        }

        @Override
        public AsyncToolTaskListResult list(AsyncToolListRequest request, AgentToolContext context) {
            AsyncToolTaskSummary summary = new AsyncToolTaskSummary();
            summary.setTaskId("task-1");
            summary.setStatus(AsyncToolTaskStatus.RUNNING);
            summary.setSummary("running");
            AsyncToolTaskListResult result = new AsyncToolTaskListResult();
            result.setTasks(List.of(summary));
            return result;
        }
    }
}
