package com.easyagents.agent.runtime.agentscope;

import com.easyagents.agent.runtime.AgentDefinition;
import com.easyagents.agent.runtime.AgentRuntimeExecutionContext;
import com.easyagents.agent.runtime.event.AgentRuntimeEvent;
import com.easyagents.agent.runtime.event.AgentRuntimeEventBridge;
import com.easyagents.agent.runtime.event.AgentRuntimeEventType;
import com.easyagents.agent.runtime.event.AgentRuntimeTurnContext;
import com.easyagents.agent.runtime.event.AgentRuntimeTurnContextHolder;
import com.easyagents.agent.runtime.hitl.AgentToolApprovalCoordinator;
import com.easyagents.agent.runtime.knowledge.AgentKnowledgeDocument;
import com.easyagents.agent.runtime.knowledge.AgentKnowledgeRegistration;
import com.easyagents.agent.runtime.knowledge.AgentKnowledgeRetrievalResult;
import com.easyagents.agent.runtime.knowledge.AgentKnowledgeSpec;
import com.easyagents.agent.runtime.knowledge.AgentKnowledgeToolNames;
import com.easyagents.agent.runtime.tool.AgentToolSpec;
import io.agentscope.core.message.ToolResultBlock;
import io.agentscope.core.message.ToolUseBlock;
import io.agentscope.core.tool.ToolCallParam;
import io.agentscope.core.tool.Toolkit;
import org.junit.Assert;
import org.junit.Test;
import reactor.core.publisher.Sinks;

import java.time.Duration;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * {@link AgentScopeKnowledgeAdapter} 回归测试。
 */
public class AgentScopeKnowledgeAdapterTest {

    /**
     * 验证每个知识库生成独立工具，工具名、描述和 Schema 向模型暴露完整信息。
     */
    @Test
    public void shouldCreateOneToolPerKnowledge() {
        AgentKnowledgeSpec first = knowledgeSpec("knowledge-1", "homeinn_faq", 5);
        first.setName("如家 FAQ");
        first.setDescription("如家酒店入住、退房和会员服务常见问题");
        AgentKnowledgeSpec second = knowledgeSpec("knowledge-2", "hotel_policy", 3);
        AgentRuntimeExecutionContext context = executionContext(List.of(first, second));
        context.setKnowledgeRegistrations(List.of(
                registration(first, 1, 0.9D),
                registration(second, 1, 0.8D)));

        List<AgentToolSpec> toolSpecs = new AgentScopeKnowledgeAdapter().createToolSpecs(context);

        Assert.assertEquals(2, toolSpecs.size());
        Assert.assertEquals("retrieve_knowledge_homeinn_faq", toolSpecs.get(0).getName());
        Assert.assertTrue(toolSpecs.get(0).getDescription().contains("如家 FAQ"));
        Assert.assertTrue(toolSpecs.get(0).getDescription().contains("入住、退房"));
        Assert.assertEquals(List.of("query"), toolSpecs.get(0).getParametersSchema().get("required"));
        Assert.assertFalse(String.valueOf(toolSpecs.get(0).getParametersSchema()).contains("limit"));
    }

    /**
     * 验证选中一个知识库工具时只调用对应 Retriever，并使用绑定 limit 和阈值。
     */
    @Test
    public void shouldDispatchOnlyToSelectedKnowledgeRetriever() {
        AgentKnowledgeSpec first = knowledgeSpec("knowledge-1", "homeinn_faq", 2);
        first.setScoreThreshold(0.5D);
        AgentKnowledgeSpec second = knowledgeSpec("knowledge-2", "hotel_policy", 2);
        AgentRuntimeExecutionContext context = executionContext(List.of(first, second));
        AtomicInteger firstCalls = new AtomicInteger();
        AtomicInteger secondCalls = new AtomicInteger();
        context.setKnowledgeRegistrations(List.of(
                new AgentKnowledgeRegistration(first, request -> {
                    firstCalls.incrementAndGet();
                    Assert.assertEquals("如家几点退房", request.getQuery());
                    Assert.assertEquals(2, request.getLimit());
                    Assert.assertEquals(0.5D, request.getScoreThreshold(), 0.0001D);
                    return AgentKnowledgeRetrievalResult.of(List.of(
                            document("knowledge-1", 0.9D),
                            document("knowledge-1-low", 0.2D)));
                }),
                new AgentKnowledgeRegistration(second, request -> {
                    secondCalls.incrementAndGet();
                    return AgentKnowledgeRetrievalResult.of(List.of(document("knowledge-2", 0.8D)));
                })));
        RegisteredKnowledgeTools registered = register(context);

        ToolResultBlock result = registered.toolkit().getTool("retrieve_knowledge_homeinn_faq")
                .callAsync(toolCall("retrieve_knowledge_homeinn_faq", "如家几点退房"))
                .block();

        Assert.assertNotNull(result);
        Assert.assertEquals(1, firstCalls.get());
        Assert.assertEquals(0, secondCalls.get());
        Assert.assertEquals(1, result.getMetadata().get("documentCount"));
    }

    /**
     * 验证最终知识库事件与阈值过滤、排序和截断后的模型可见文档一致。
     */
    @Test
    public void retrievalEventShouldMatchFinalToolDocuments() {
        AgentKnowledgeSpec spec = knowledgeSpec("knowledge-1", "homeinn_faq", 2);
        spec.setScoreThreshold(0.5D);
        AgentRuntimeExecutionContext context = executionContext(List.of(spec));
        context.setKnowledgeRegistrations(List.of(new AgentKnowledgeRegistration(spec, request ->
                AgentKnowledgeRetrievalResult.of(List.of(
                        document("third", 0.7D),
                        document("first", 0.9D),
                        document("filtered", 0.2D),
                        document("second", 0.8D))))));
        RegisteredKnowledgeTools registered = register(context);

        ToolResultBlock result = registered.toolkit().getTool("retrieve_knowledge_homeinn_faq")
                .callAsync(toolCall("retrieve_knowledge_homeinn_faq", "query"))
                .block();
        List<AgentRuntimeEvent> events = registered.eventSink().asFlux()
                .filter(event -> event.getEventType() == AgentRuntimeEventType.KNOWLEDGE_RETRIEVAL)
                .take(1)
                .collectList()
                .block(Duration.ofSeconds(1));

        Assert.assertNotNull(result);
        Assert.assertNotNull(events);
        Assert.assertEquals(1, events.size());
        AgentRuntimeEvent event = events.get(0);
        Assert.assertEquals(2, event.getPayload().get("documentCount"));
        @SuppressWarnings("unchecked")
        List<Map<String, Object>> eventDocuments =
                (List<Map<String, Object>>) event.getPayload().get("documents");
        @SuppressWarnings("unchecked")
        List<Map<String, Object>> resultDocuments =
                (List<Map<String, Object>>) result.getMetadata().get("documents");
        Assert.assertEquals(resultDocuments, eventDocuments);
        Assert.assertEquals("first", eventDocuments.get(0).get("documentId"));
        Assert.assertEquals("second", eventDocuments.get(1).get("documentId"));
    }

    /**
     * 验证超长运行名会稳定压缩到 Function Call 长度上限。
     */
    @Test
    public void longRuntimeNameShouldProduceStableBoundedToolName() {
        String runtimeName = "knowledge_" + "a".repeat(80);

        String first = AgentKnowledgeToolNames.build(runtimeName);
        String second = AgentKnowledgeToolNames.build(runtimeName);

        Assert.assertEquals(first, second);
        Assert.assertEquals(AgentKnowledgeToolNames.MAX_TOOL_NAME_LENGTH, first.length());
        Assert.assertTrue(first.startsWith(AgentKnowledgeToolNames.PREFIX));
    }

    private RegisteredKnowledgeTools register(AgentRuntimeExecutionContext context) {
        AgentScopeKnowledgeAdapter adapter = new AgentScopeKnowledgeAdapter();
        List<AgentToolSpec> toolSpecs = adapter.createToolSpecs(context);
        Toolkit toolkit = new Toolkit();
        Sinks.Many<AgentRuntimeEvent> eventSink = Sinks.many().replay().all();
        AgentRuntimeTurnContextHolder holder = new AgentRuntimeTurnContextHolder();
        AgentRuntimeEventBridge eventBridge = new AgentRuntimeEventBridge(context, holder);
        holder.set(new AgentRuntimeTurnContext(context, eventSink, eventBridge));
        adapter.registerTools(context, toolSpecs, toolkit, new AgentScopeToolAdapter(),
                AgentToolApprovalCoordinator.disabled(), holder);
        return new RegisteredKnowledgeTools(toolkit, eventSink);
    }

    private ToolCallParam toolCall(String toolName, String query) {
        ToolUseBlock toolUseBlock = ToolUseBlock.builder()
                .id("call-1")
                .name(toolName)
                .input(Map.of("query", query))
                .build();
        return ToolCallParam.builder()
                .toolUseBlock(toolUseBlock)
                .input(toolUseBlock.getInput())
                .build();
    }

    private AgentRuntimeExecutionContext executionContext(List<AgentKnowledgeSpec> specs) {
        AgentDefinition definition = new AgentDefinition();
        definition.setAgentId("agent-1");
        definition.setKnowledgeSpecs(specs);
        AgentRuntimeExecutionContext context = new AgentRuntimeExecutionContext();
        context.setAgentDefinition(definition);
        context.setTraceId("trace-1");
        context.setSessionId("session-1");
        return context;
    }

    private AgentKnowledgeSpec knowledgeSpec(String knowledgeId, String runtimeName, int limit) {
        AgentKnowledgeSpec spec = new AgentKnowledgeSpec();
        spec.setKnowledgeId(knowledgeId);
        spec.setRuntimeName(runtimeName);
        spec.setName(knowledgeId);
        spec.setLimit(limit);
        return spec;
    }

    private AgentKnowledgeRegistration registration(AgentKnowledgeSpec spec,
                                                    int documentCount,
                                                    double firstScore) {
        return new AgentKnowledgeRegistration(spec, request ->
                AgentKnowledgeRetrievalResult.of(documents(spec.getKnowledgeId(), documentCount, firstScore)));
    }

    private List<AgentKnowledgeDocument> documents(String knowledgeId, int count, double firstScore) {
        List<AgentKnowledgeDocument> documents = new ArrayList<>();
        for (int index = 0; index < count; index++) {
            documents.add(document(knowledgeId + "-doc-" + index, firstScore - (index * 0.01D)));
        }
        return documents;
    }

    private AgentKnowledgeDocument document(String documentId, double score) {
        AgentKnowledgeDocument document = new AgentKnowledgeDocument();
        document.setDocumentId(documentId);
        document.setDocumentName("FAQ");
        document.setChunkId(documentId + "-chunk");
        document.setContent("content-" + documentId);
        document.setScore(score);
        return document;
    }

    private record RegisteredKnowledgeTools(Toolkit toolkit,
                                            Sinks.Many<AgentRuntimeEvent> eventSink) {
    }
}
