package com.agent.management.controller;

import com.agent.management.common.Result;
import com.agent.management.config.KbProperties;
import com.agent.management.service.VectorStoreService;
import lombok.Data;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.document.Document;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;

/**
 * 知识库检索 API（向量召回）
 *
 * <p>对查询文本向量化后从 Milvus 召回 TopK 条相似分块，附带元数据用于
 * 在 UI 展示来源文档、分类、分块序号，或在工作流节点中后处理拼装上下文。</p>
 */
@Slf4j
@RestController
@RequestMapping("/api/kb/search")
@RequiredArgsConstructor
public class KbSearchController {

    private final VectorStoreService vectorStoreService;
    private final KbProperties kbProps;

    @PostMapping
    public Result<List<SearchHit>> search(@RequestBody SearchReq req) {
        if (!kbProps.isEnabled()) {
            return Result.error(503, "知识库子系统已禁用");
        }
        if (req.getQuery() == null || req.getQuery().trim().isEmpty()) {
            return Result.error(400, "query 不能为空");
        }

        int topK = req.getTopK() != null && req.getTopK() > 0 ? req.getTopK() : 5;
        Double threshold = req.getSimilarityThreshold();  // null 表示使用默认值
        String filter = buildFilter(req.getCategoryId());

        List<Document> docs = vectorStoreService.similaritySearch(
                req.getQuery().trim(), topK, threshold, filter);

        List<SearchHit> hits = new ArrayList<>(docs.size());
        for (Document d : docs) {
            hits.add(SearchHit.from(d));
        }
        return Result.success(hits);
    }

    /**
     * 构造元数据过滤表达式：categoryId 精确匹配（Spring AI 字符串语法）
     */
    private String buildFilter(Long categoryId) {
        if (categoryId == null) {
            return null;
        }
        return "categoryId == " + categoryId;
    }

    @Data
    public static class SearchReq {
        /** 查询文本（必填） */
        private String query;
        /** 返回条数上限，默认 5 */
        private Integer topK;
        /** 相似度阈值（0~1），null 表示使用 VectorStore 默认值 */
        private Double similarityThreshold;
        /** 分类过滤（精确匹配元数据中的 categoryId），null 表示不过滤 */
        private Long categoryId;
    }

    @Data
    public static class SearchHit {
        /** 命中分块原文 */
        private String content;
        /** 相似度分数（来源于元数据 distance 字段，越大越相似） */
        private Double score;
        /** 文档 ID（元数据） */
        private Long documentId;
        /** 分块 ID（元数据，可用于回查 H2 中的 chunk 内容） */
        private Long chunkId;
        /** 分块序号（元数据） */
        private Integer chunkIndex;
        /** 分类 ID（元数据） */
        private Long categoryId;
        /** 来源文档名称（元数据） */
        private String docName;

        static SearchHit from(Document d) {
            SearchHit h = new SearchHit();
            h.content = d.getText();
            Map<String, Object> meta = d.getMetadata();
            if (meta != null) {
                h.score = readDouble(meta.get("distance"));
                h.documentId = readLong(meta.get("documentId"));
                h.chunkId = readLong(meta.get("chunkId"));
                h.chunkIndex = readInt(meta.get("chunkIndex"));
                h.categoryId = readLong(meta.get("categoryId"));
                h.docName = readString(meta.get("docName"));
            }
            return h;
        }

        private static Double readDouble(Object v) {
            return v instanceof Number ? ((Number) v).doubleValue() : null;
        }
        private static Long readLong(Object v) {
            if (v instanceof Number) return ((Number) v).longValue();
            if (v instanceof String s && !s.isBlank()) {
                try { return Long.parseLong(s); } catch (NumberFormatException ignored) {}
            }
            return null;
        }
        private static Integer readInt(Object v) {
            if (v instanceof Number) return ((Number) v).intValue();
            if (v instanceof String s && !s.isBlank()) {
                try { return Integer.parseInt(s); } catch (NumberFormatException ignored) {}
            }
            return null;
        }
        private static String readString(Object v) {
            return v == null ? null : v.toString();
        }
    }
}
