package com.easyagents.agent.runtime.agentscope;

import com.easyagents.agent.runtime.AgentRuntimeException;
import com.easyagents.agent.runtime.AgentRuntimeExecutionContext;
import com.easyagents.agent.runtime.event.*;
import com.easyagents.agent.runtime.knowledge.*;
import io.agentscope.core.message.TextBlock;
import io.agentscope.core.rag.Knowledge;
import io.agentscope.core.rag.model.Document;
import io.agentscope.core.rag.model.DocumentMetadata;
import io.agentscope.core.rag.model.RetrieveConfig;
import reactor.core.publisher.Mono;
import reactor.core.publisher.Sinks;

import java.util.*;

/**
 * 将运行时知识库检索器适配为一个聚合 AgentScope Knowledge。
 */
public class AgentScopeKnowledgeAdapter {

    /**
     * 创建聚合 Knowledge。
     *
     * @param request 运行请求
     * @return 聚合 Knowledge；未配置知识库时返回 null
     */
    public Knowledge createAggregateKnowledge(AgentRuntimeExecutionContext request) {
        return createAggregateKnowledge(request, (Sinks.Many<AgentRuntimeEvent>) null);
    }

    /**
     * 创建带事件 sink 的聚合 Knowledge。
     *
     * @param request 运行请求
     * @param eventSink 事件 sink
     * @return 聚合 Knowledge；未配置知识库时返回 null
     */
    public Knowledge createAggregateKnowledge(AgentRuntimeExecutionContext request, Sinks.Many<AgentRuntimeEvent> eventSink) {
        return createAggregateKnowledge(request, fixedHolder(request, eventSink));
    }

    /**
     * 创建可读取当前运行轮次事件出口的聚合 Knowledge。
     *
     * @param request 运行时级上下文
     * @param turnContextHolder 当前运行轮次上下文持有器
     * @return 聚合 Knowledge；未配置知识库时返回 null
     */
    public Knowledge createAggregateKnowledge(AgentRuntimeExecutionContext request,
                                              AgentRuntimeTurnContextHolder turnContextHolder) {
        if (request.getAgentDefinition().getKnowledgeSpecs().isEmpty()) {
            return null;
        }
        return new AggregateKnowledge(request, turnContextHolder);
    }

    private AgentRuntimeTurnContextHolder fixedHolder(AgentRuntimeExecutionContext request,
                                                      Sinks.Many<AgentRuntimeEvent> eventSink) {
        AgentRuntimeTurnContextHolder holder = new AgentRuntimeTurnContextHolder();
        AgentRuntimeEventBridge bridge = new AgentRuntimeEventBridge(request, holder);
        holder.set(new AgentRuntimeTurnContext(null, eventSink, bridge));
        return holder;
    }

    /**
     * 将运行时文档转换为 AgentScope 文档。
     *
     * @param documents 运行时文档
     * @return AgentScope 文档
     */
    public List<Document> toDocuments(List<AgentKnowledgeDocument> documents) {
        List<Document> converted = new ArrayList<>();
        if (documents == null) {
            return converted;
        }
        for (AgentKnowledgeDocument document : documents) {
            converted.add(toDocument(document));
        }
        return converted;
    }

    private Document toDocument(AgentKnowledgeDocument document) {
        Map<String, Object> payload = new LinkedHashMap<>();
        payload.put("documentId", document.getDocumentId());
        payload.put("documentName", document.getDocumentName());
        payload.put("chunkId", document.getChunkId());
        payload.put("sourceUri", document.getSourceUri());
        payload.put("knowledgeMetadata", document.getKnowledgeMetadata());
        payload.put("documentMetadata", document.getMetadata());
        payload.putAll(document.getMetadata());
        DocumentMetadata metadata = DocumentMetadata.builder()
                .content(TextBlock.builder().text(safeContent(document)).build())
                .docId(safeDocumentId(document))
                .chunkId(safeChunkId(document))
                .payload(payload)
                .build();
        Document converted = new Document(metadata);
        converted.setScore(document.getScore());
        return converted;
    }

    /**
     * 获取 AgentScope 要求的非空文档 ID。
     *
     * @param document 知识文档
     * @return 非空文档 ID
     */
    private String safeDocumentId(AgentKnowledgeDocument document) {
        if (document.getDocumentId() != null && !document.getDocumentId().isBlank()) {
            return document.getDocumentId();
        }
        if (document.getChunkId() != null && !document.getChunkId().isBlank()) {
            return document.getChunkId();
        }
        return "knowledge-document";
    }

    /**
     * 获取 AgentScope 要求的非空分片 ID。
     *
     * @param document 知识文档
     * @return 非空分片 ID
     */
    private String safeChunkId(AgentKnowledgeDocument document) {
        if (document.getChunkId() != null && !document.getChunkId().isBlank()) {
            return document.getChunkId();
        }
        if (document.getDocumentId() != null && !document.getDocumentId().isBlank()) {
            return document.getDocumentId();
        }
        return "0";
    }

    /**
     * 获取 AgentScope 要求的非空文档内容。
     *
     * @param document 知识文档
     * @return 文档内容
     */
    private String safeContent(AgentKnowledgeDocument document) {
        return document.getContent() == null ? "" : document.getContent();
    }

    /**
     * 将检索调用分发到多个知识源的聚合 Knowledge 实现。
     */
    private class AggregateKnowledge implements Knowledge {

        private final AgentRuntimeExecutionContext request;
        private final AgentRuntimeTurnContextHolder turnContextHolder;

        private AggregateKnowledge(AgentRuntimeExecutionContext request, AgentRuntimeTurnContextHolder turnContextHolder) {
            this.request = request;
            this.turnContextHolder = turnContextHolder;
        }

        /**
         * 忽略文档新增，因为知识库索引由 EasyFlow 负责。
         *
         * @param documents 文档列表
         * @return 完成信号
         */
        @Override
        public Mono<Void> addDocuments(List<Document> documents) {
            return Mono.error(new UnsupportedOperationException(
                "Easy-Agents agent runtime knowledge does not support addDocuments. Use external knowledge service instead."));
        }

        /**
         * 从已配置的知识源检索文档。
         *
         * @param query 查询
         * @param config 检索配置
         * @return 文档列表
         */
        @Override
        public Mono<List<Document>> retrieve(String query, RetrieveConfig config) {
            return Mono.fromCallable(() -> retrieveAll(query, config));
        }

        /**
         * 检索并合并所有已配置的知识源。
         *
         * @param query 查询
         * @param config 检索配置
         * @return 合并后的文档
         */
        private List<Document> retrieveAll(String query, RetrieveConfig config) {
            List<AgentKnowledgeDocument> allDocuments = new ArrayList<>();
            int globalLimit = config == null || config.getLimit() <= 0 ? 5 : config.getLimit();
            double globalThreshold = config == null ? 0D : config.getScoreThreshold();
            for (AgentKnowledgeSpec spec : request.getAgentDefinition().getKnowledgeSpecs()) {
                AgentKnowledgeRetriever retriever = request.getKnowledgeRetrievers().get(spec.getKnowledgeId());
                if (retriever == null) {
                    throw new AgentRuntimeException("Knowledge retriever is required: " + spec.getKnowledgeId());
                }
                AgentKnowledgeRetrievalRequest retrievalRequest = new AgentKnowledgeRetrievalRequest();
                retrievalRequest.setQuery(query);
                retrievalRequest.setLimit(spec.getLimit());
                retrievalRequest.setScoreThreshold(Math.max(spec.getScoreThreshold(), globalThreshold));
                retrievalRequest.setKnowledgeSpec(spec);
                AgentRuntimeExecutionContext currentRequest = currentRequest();
                retrievalRequest.setRuntimeContext(currentRequest.getRuntimeContext());
                retrievalRequest.getMetadata().put("traceId", currentRequest.getTraceId());
                retrievalRequest.getMetadata().put("sessionId", currentRequest.getSessionId());
                AgentKnowledgeRetrievalResult result = retriever.retrieve(retrievalRequest);
                if (result == null || result.getDocuments() == null) {
                    emitKnowledgeRetrievalEvent(query, spec, retrievalRequest, new ArrayList<>());
                    continue;
                }
                emitKnowledgeRetrievalEvent(query, spec, retrievalRequest, result.getDocuments());
                for (AgentKnowledgeDocument document : result.getDocuments()) {
                    preserveKnowledgeMetadata(spec, document);
                    allDocuments.add(document);
                }
            }
            allDocuments.sort(Comparator.comparing(
                    AgentKnowledgeDocument::getScore,
                    Comparator.nullsLast(Comparator.reverseOrder())
            ));
            if (allDocuments.size() > globalLimit) {
                allDocuments = new ArrayList<>(allDocuments.subList(0, globalLimit));
            }
            return toDocuments(allDocuments);
        }

        /**
         * 在单条文档上保留知识库级元数据。
         *
         * @param spec 知识库声明
         * @param document 文档
         */
        private void preserveKnowledgeMetadata(AgentKnowledgeSpec spec, AgentKnowledgeDocument document) {
            Map<String, Object> knowledgeMetadata = new LinkedHashMap<>(spec.getMetadata());
            knowledgeMetadata.put("knowledgeId", spec.getKnowledgeId());
            knowledgeMetadata.put("knowledgeName", spec.getName());
            knowledgeMetadata.put("retrievalMode", spec.getRetrievalMode().name());
            knowledgeMetadata.putAll(document.getKnowledgeMetadata());
            document.setKnowledgeMetadata(knowledgeMetadata);
            document.getMetadata().putIfAbsent("knowledgeId", spec.getKnowledgeId());
            document.getMetadata().putIfAbsent("knowledgeName", spec.getName());
        }

        /**
         * 发射知识库检索旁路事件，供聊天界面展示检索过程。
         *
         * <p>知识库检索本身属于 AgentScope RAG 主线路，返回的 Document 会继续进入
         * AgentScope 的上下文注入流程；这里发出的 {@code KNOWLEDGE_RETRIEVAL}
         * 只是旁路告知调用方，不会回写 memory，也不会参与模型消息序列。</p>
         *
         * @param query 查询
         * @param spec 知识库声明
         * @param retrievalRequest 检索请求
         * @param documents 检索文档
         */
        private void emitKnowledgeRetrievalEvent(String query,
                                                 AgentKnowledgeSpec spec,
                                                 AgentKnowledgeRetrievalRequest retrievalRequest,
                                                 List<AgentKnowledgeDocument> documents) {
            AgentRuntimeEvent event = currentEventBridge().event(AgentRuntimeEventType.KNOWLEDGE_RETRIEVAL);
            event.getPayload().put("query", query);
            event.getPayload().put("knowledgeId", spec.getKnowledgeId());
            event.getPayload().put("knowledgeName", spec.getName());
            event.getPayload().put("knowledgeType", spec.getMetadata().get("knowledgeType"));
            event.getPayload().put("faqCollection", spec.getMetadata().get("faqCollection"));
            event.getPayload().put("limit", retrievalRequest.getLimit());
            event.getPayload().put("scoreThreshold", retrievalRequest.getScoreThreshold());
            event.getPayload().put("documentCount", documents == null ? 0 : documents.size());
            event.getPayload().put("documents", documentSummaries(documents));
            currentEventBridge().emit(event);
        }

        private AgentRuntimeExecutionContext currentRequest() {
            return turnContextHolder == null ? request : turnContextHolder.executionContext(request);
        }

        private AgentRuntimeEventBridge currentEventBridge() {
            if (turnContextHolder != null && turnContextHolder.eventBridge().isPresent()) {
                return turnContextHolder.eventBridge().get();
            }
            return new AgentRuntimeEventBridge(request, turnContextHolder);
        }

        /**
         * 构建用于事件展示的命中片段，保留前端引注需要的原始 chunk 内容。
         *
         * @param documents 检索文档
         * @return 命中片段列表
         */
        private List<Map<String, Object>> documentSummaries(List<AgentKnowledgeDocument> documents) {
            List<Map<String, Object>> summaries = new ArrayList<>();
            if (documents == null) {
                return summaries;
            }
            for (AgentKnowledgeDocument document : documents) {
                Map<String, Object> summary = new LinkedHashMap<>();
                summary.put("documentId", document.getDocumentId());
                summary.put("documentName", document.getDocumentName());
                summary.put("chunkId", document.getChunkId());
                summary.put("chunkContent", document.getContent());
                summary.put("score", document.getScore());
                summary.put("sourceUri", document.getSourceUri());
                summary.put("metadata", document.getMetadata());
                summaries.add(summary);
            }
            return summaries;
        }
    }
}
