package com.emoon.openplatform.service.impl;

import cn.hutool.core.map.MapUtil;
import com.emoon.common.xinference.XinferenceClient;
import com.emoon.openplatform.domain.bo.QueryVectorBo;
import com.emoon.openplatform.service.IVectorStoreService;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.weaviate.WeaviateEmbeddingStore;
import io.weaviate.client.WeaviateClient;
import io.weaviate.client.base.Result;
import io.weaviate.client.v1.graphql.model.GraphQLResponse;
import io.weaviate.client.v1.graphql.query.Raw;
import io.weaviate.client.v1.schema.model.Property;
import io.weaviate.client.v1.schema.model.Schema;
import io.weaviate.client.v1.schema.model.WeaviateClass;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;

import java.io.IOException;
import java.util.*;
import java.util.stream.Collectors;

/**
 * @author destiny
 * @description:
 * @date 2025/10/13 16:52
 */
@Slf4j
@Service
@RequiredArgsConstructor
public class VectorStoreServiceImpl implements IVectorStoreService {

    private final WeaviateClient weaviateClient;

    private EmbeddingStore<TextSegment> embeddingStore;

    @Value("${weaviate.protocol}")
    private String weaviateProtocol;

    @Value("${weaviate.host}")
    private String weaviateHost;

    @Value("${weaviate.port}")
    private String weaviatePort;

    @Value("${weaviate.className:LocalKnowledge}")
    private String weaviateClassName;

    @Value("${file.base-url}")
    private String fileBaseUrl;

    /**
     * 执行向量查询操作，根据输入的查询向量在Weaviate中进行语义搜索
     *
     * @param queryVectorBo 查询向量业务对象，包含查询所需的各种参数
     * @return 包含查询结果的Map列表，每个Map代表一个匹配的文档片段
     */
    @Override
    public List<Map<String, Object>> getQueryVector(QueryVectorBo queryVectorBo) {
        // 构造完整的类名并创建对应的schema
        String fullClassName = weaviateClassName + queryVectorBo.getKid();
        createSchema(queryVectorBo.getKid(), queryVectorBo.getVectorModelName());

        // 将查询向量转换为字符串格式，用于GraphQL查询
        String vectorString = getDoubles(queryVectorBo).stream()
            .map(Double::floatValue)
            .toList()
            .stream()
            .map(String::valueOf)
            .collect(Collectors.joining(", "));

        // 执行基于向量的语义搜索，使用nearVector进行相似度匹配
        Result<GraphQLResponse> result = weaviateClient.graphQL().raw().withQuery(String.format(
            """
                {
                    Get {%n
                        %s(nearVector: {vector: [%s], certainty: 0.7}, limit: %d) {%n
                            text%n
                            docId%n
                            filename%n
                            chunkIndex%n
                            _additional { id distance }%n
                        }%n
                    }%n
                }
                """, fullClassName, vectorString, queryVectorBo.getMaxResults()
        )).run();

        // 提取查询结果数据
        Map<String, Object> dataMap = extractDataMap(result);
        if (MapUtil.isEmpty(dataMap)) {
            return Collections.emptyList();
        }

        // 获取指定类别的查询结果
        List<Map<String, Object>> classData = (List<Map<String, Object>>) dataMap.get(fullClassName);
        if (classData == null || classData.isEmpty()) {
            return Collections.emptyList();
        }
        return classData;
    }

    @Override
    public List<Map<String, Object>> getQueryVector(QueryVectorBo queryVectorBo, double score) {
        // 构造完整的类名并创建对应的schema
        String fullClassName = weaviateClassName + queryVectorBo.getKid();
        createSchema(queryVectorBo.getKid(), queryVectorBo.getVectorModelName());

        // 将查询向量转换为字符串格式，用于GraphQL查询
        String vectorString = getDoubles(queryVectorBo).stream()
            .map(Double::floatValue)
            .toList()
            .stream()
            .map(String::valueOf)
            .collect(Collectors.joining(", "));

        // 主查询：语义搜索 Top-K

        // 执行主查询
        Result<GraphQLResponse> result = weaviateClient.graphQL().raw().withQuery(String.format("""
            {
                Get {%n
                    %s(nearVector: {vector: [%s], certainty: %f}, limit: %d) {%n
                        text%n
                        docId%n
                        filename%n
                        chunkIndex%n
                        _additional { id distance }%n
                    }%n
               }%n
            }
            """, fullClassName, vectorString, score, queryVectorBo.getMaxResults())).run();

        // 提取并处理查询结果
        Map<String, Object> dataMap = extractDataMap(result);
        if (dataMap == null) {
            return Collections.emptyList();
        }

        List<Map<String, Object>> classData = (List<Map<String, Object>>) dataMap.get(fullClassName);
        if (classData == null || classData.isEmpty()) {
            return Collections.emptyList();
        }

        // 用于按 docId 收集所有相关的 chunkIndex -> text 映射
        Map<String, Map<Integer, String>> docIdToChunkMap = new HashMap<>();

        // 第一步：遍历主查询结果，获取每个 chunk 及其上下文
        for (Map<String, Object> targetChunk : classData) {
            String docId = (String) targetChunk.get("docId");
            if (docId == null) continue;

            Number chunkIndexNum = (Number) targetChunk.get("chunkIndex");
            int targetIndex = chunkIndexNum != null ? chunkIndexNum.intValue() : -1;
            if (targetIndex < 0) continue;

            // 确保该 docId 的 map 存在
            Map<Integer, String> chunkMap = docIdToChunkMap.computeIfAbsent(docId, k -> new HashMap<>());

            // 添加当前 chunk
            String text = (String) targetChunk.get("text");
            if (text != null) {
                chunkMap.put(targetIndex, text);
            }

            // 查询前一个 (targetIndex - 1)
            Map<String, Object> prev = queryAdjacentChunk(fullClassName, docId, targetIndex - 1);
            if (prev != null && prev.containsKey("text")) {
                Integer idx = ((Number) prev.get("chunkIndex")).intValue();
                chunkMap.put(idx, (String) prev.get("text"));
            }

            // 查询后一个 (targetIndex + 1)
            Map<String, Object> next = queryAdjacentChunk(fullClassName, docId, targetIndex + 1);
            if (next != null && next.containsKey("text")) {
                Integer idx = ((Number) next.get("chunkIndex")).intValue();
                chunkMap.put(idx, (String) next.get("text"));
            }
        }

        // 第二步：对每个 docId，合并连续 chunkIndex，生成结果
        List<Map<String, Object>> results = new ArrayList<>();

        for (Map.Entry<String, Map<Integer, String>> entry : docIdToChunkMap.entrySet()) {
            String docId = entry.getKey();
            Map<Integer, String> chunkMap = entry.getValue();

            // 提取所有 index 并排序
            List<Integer> indices = new ArrayList<>(chunkMap.keySet());
            indices.sort(Integer::compareTo);

            if (indices.isEmpty()) continue;

            // 合并连续区间
            List<int[]> ranges = new ArrayList<>();
            int start = indices.get(0), end = indices.get(0);

            for (int i = 1; i < indices.size(); i++) {
                if (indices.get(i) == end + 1) {
                    end = indices.get(i); // 连续
                } else {
                    ranges.add(new int[]{start, end}); // 不连续，保存
                    start = end = indices.get(i);
                }
            }
            ranges.add(new int[]{start, end}); // 最后一段

            // 为每个连续区间生成结果
            for (int[] range : ranges) {
                int startIdx = range[0];
                int endIdx = range[1];

                // 拼接文本
                StringBuilder fullText = new StringBuilder();
                for (int idx = startIdx; idx <= endIdx; idx++) {
                    String text = chunkMap.get(idx);
                    if (text != null) {
                        if (fullText.length() > 0) fullText.append(" ");
                        fullText.append(text);
                    }
                }

                // 获取 filename（从原始数据中找任意一个匹配的）
                String filename = null;
                for (int idx = startIdx; idx <= endIdx; idx++) {
                    Map<String, Object> found = findChunkInClassData(classData, docId, idx);
                    if (found != null && found.get("filename") != null) {
                        filename = (String) found.get("filename");
                        break;
                    }
                }

                // 构造返回结果
                Map<String, Object> resultMap = new HashMap<>();
                resultMap.put("filename", filename != null ? filename : "unknown");
                resultMap.put("text", fullText.toString());
                resultMap.put("docId", docId);
                resultMap.put("downloadUrl", fileBaseUrl + "/api/file/view?docId=" + docId);
                resultMap.put("chunkIndexRange", startIdx + " - " + endIdx);

                results.add(resultMap);
            }
        }

        return results;
    }


    /**
     * 根据查询向量执行语义搜索，返回最相似的结果列表
     *
     * @param queryVectorBo 查询向量相关的业务对象，包含kid、向量模型名称等信息
     * @param score         相似度阈值，用于过滤搜索结果
     * @param isRetrieval   是否为检索模式（该参数在当前实现中未使用）
     * @return 包含搜索结果的Map列表，每个Map包含text、docId、filename、chunkIndex等字段
     */
    @Override
    public List<Map<String, Object>> getQueryVector(QueryVectorBo queryVectorBo, double score, boolean isRetrieval) {
        // 构造完整的类名并创建对应的schema
        String fullClassName = weaviateClassName + queryVectorBo.getKid();
        createSchema(queryVectorBo.getKid(), queryVectorBo.getVectorModelName());

        // 将查询向量转换为字符串格式，用于GraphQL查询
        String vectorString = getDoubles(queryVectorBo).stream()
            .map(Double::floatValue)
            .toList()
            .stream()
            .map(String::valueOf)
            .collect(Collectors.joining(", "));

        // 主查询：语义搜索 Top-K

        // 执行主查询
        Result<GraphQLResponse> result = weaviateClient.graphQL().raw().withQuery(String.format("""
            {
                Get {%n
                    %s(nearVector: {vector: [%s], certainty: %f}, limit: %d) {%n
                        text%n
                        docId%n
                        filename%n
                        chunkIndex%n
                        _additional { id distance }%n
                    }%n
               }%n
            }
            """, fullClassName, vectorString, score, queryVectorBo.getMaxResults())).run();

        // 提取并处理查询结果
        Map<String, Object> dataMap = extractDataMap(result);
        if (dataMap == null) return Collections.emptyList();

        List<Map<String, Object>> classData = (List<Map<String, Object>>) dataMap.get(fullClassName);
        if (classData == null || classData.isEmpty()) return Collections.emptyList();
        return classData;
    }


    /**
     * 创建 Weaviate schema，如果指定的类不存在则创建新类，否则检查并补充必要的属性。
     *
     * @param knowledgeCode 知识库编码，用于构建完整的类名（weaviateClassName + knowledgeCode）
     * @param modelName     模型名称
     */
    @Override
    public void createSchema(String knowledgeCode, String modelName) {
        // 构造完整的 Weaviate 类名
        String fullClassName = weaviateClassName + knowledgeCode;

        // 获取现有的 schema 信息
        Result<Schema> existingSchema = weaviateClient.schema().getter().run();

        // 检查获取 schema 是否出错
        if (existingSchema.hasErrors()) {
            throw new RuntimeException("Failed to get schema: " + existingSchema.getError().toString());
        }

        // 判断目标 class 是否已存在
        if (existingSchema.getResult().getClasses().stream()
            .noneMatch(clazz -> clazz.getClassName().equals(fullClassName))) {

            // 如果类不存在，则创建新的 Weaviate 类，并定义其属性
            WeaviateClass weaviateClass = WeaviateClass.builder()
                .className(fullClassName)
                .properties(Arrays.asList(
                    Property.builder().name("text").dataType(List.of("text")).build(),
                    Property.builder().name("indexFilterable").dataType(List.of("boolean")).build(),
                    Property.builder().name("indexSearchable").dataType(List.of("boolean")).build(),
                    Property.builder().name("docId").dataType(List.of("string")).build(),
                    Property.builder().name("filename").dataType(List.of("string")).build(),
                    Property.builder().name("metadata").dataType(List.of("string")).build(),
                    Property.builder().name("chunkIndex").dataType(List.of("int")).build()
                ))
                .vectorizer("none")
                .build();

            // 执行类创建操作
            weaviateClient.schema().classCreator().withClass(weaviateClass).run();
        } else {
            // 类已存在，检查并添加缺失的属性

            // 获取当前已存在的类定义
            WeaviateClass existingClass = existingSchema.getResult().getClasses().stream()
                .filter(c -> c.getClassName().equals(fullClassName))
                .findFirst()
                .orElse(null);

            // 检查是否存在 docId 属性，如不存在则添加
            boolean hasDocId = existingClass.getProperties().stream()
                .anyMatch(p -> p.getName().equals("docId"));
            if (!hasDocId) {
                Property docIdProperty = Property.builder()
                    .name("docId")
                    .dataType(Arrays.asList("string"))
                    .build();
                weaviateClient.schema().propertyCreator()
                    .withClassName(fullClassName)
                    .withProperty(docIdProperty)
                    .run();
            }

            // 检查是否存在 filename 属性，如不存在则添加
            boolean filename = existingClass.getProperties().stream()
                .anyMatch(p -> p.getName().equals("filename"));
            if (!filename) {
                Property docIdProperty = Property.builder()
                    .name("filename")
                    .dataType(Arrays.asList("string"))
                    .build();
                weaviateClient.schema().propertyCreator()
                    .withClassName(fullClassName)
                    .withProperty(docIdProperty)
                    .run();
            }

            // 检查是否存在 chunkIndex 属性，如不存在则添加
            boolean chunkIndex = existingClass.getProperties().stream()
                .anyMatch(p -> p.getName().equals("chunkIndex"));
            if (!chunkIndex) {
                Property docIdProperty = Property.builder()
                    .name("chunkIndex")
                    .dataType(Arrays.asList("int"))
                    .build();
                weaviateClient.schema().propertyCreator()
                    .withClassName(fullClassName)
                    .withProperty(docIdProperty)
                    .run();
            }
        }

        // 初始化 embeddingStore，用于后续向量存储操作
        embeddingStore = WeaviateEmbeddingStore.builder()
            .scheme(weaviateProtocol)
            .host(String.format("%s:%s", weaviateHost, weaviatePort))
            .objectClass(fullClassName)
            .avoidDups(true)
            .consistencyLevel("ALL")
            .build();
    }


    private static List<Double> getDoubles(QueryVectorBo queryVectorBo) {
        XinferenceClient xinferenceClient = new XinferenceClient(queryVectorBo.getBaseUrl());
        String embeddingModelName = queryVectorBo.getEmbeddingModelName();
        String[] embeddingModelNameArr = embeddingModelName.split("/");
        String modelUid = embeddingModelNameArr[embeddingModelNameArr.length - 1];
        List<Double> queryVector;
        try {
            queryVector = xinferenceClient.getEmbedding(modelUid, queryVectorBo.getQuery());
        } catch (IOException e) {
            throw new RuntimeException("生成向量失败: " + e.getMessage(), e);
        }
        return queryVector;
    }

    // 提取 GraphQL 返回的 data -> Get -> className 数据
    private Map<String, Object> extractDataMap(Result<GraphQLResponse> result) {
        Object data = result.getResult().getData();
        if (!(data instanceof Map)) return null;
        Map<String, Object> dataMap = (Map<String, Object>) data;
        return (Map<String, Object>) dataMap.get("Get");
    }

    /**
     * 查询相邻 chunk（指定 docId 和 chunkIndex）
     *
     * @param className  Weaviate 中的类名，指定要查询的数据类
     * @param docId      文档ID，用于筛选特定文档的chunk
     * @param chunkIndex chunk索引，用于定位文档中的特定chunk
     * @return 返回包含chunk信息的Map对象，如果未找到则返回null
     */
    private Map<String, Object> queryAdjacentChunk(String className, String docId, int chunkIndex) {
        // 构造GraphQL查询语句，根据docId和chunkIndex精确查找对应的chunk数据
        Raw contextRaw = weaviateClient.graphQL().raw().withQuery(String.format("""
            {
              Get {
                %s(where: {
                  operator: And, operands: [
                    {path: ["docId"], operator: Equal, valueText: "%s"},
                    {path: ["chunkIndex"], operator: Equal, valueInt: %d}
                  ]
                }, limit: 1) {
                  text docId filename chunkIndex _additional { id distance }
                }
              }
            }
            """, className, docId, chunkIndex));
        // 执行GraphQL查询并获取结果
        Result<GraphQLResponse> contextResult = contextRaw.run();
        Object contextData = contextResult.getResult().getData();

        // 检查查询结果是否为空或类型不匹配
        if (!(contextData instanceof Map)) return null;

        // 解析查询结果，提取Get节点下的数据
        Map<String, Object> contextDataMap = (Map<String, Object>) contextData;
        Map<String, Object> contextGetMap = (Map<String, Object>) contextDataMap.get("Get");
        if (contextGetMap == null) return null;

        // 获取指定类名下的chunk列表
        List<Map<String, Object>> contextList = (List<Map<String, Object>>) contextGetMap.get(className);
        if (contextList == null || contextList.isEmpty()) return null;

        // 返回第一个匹配的chunk数据
        return contextList.get(0);
    }

    /**
     * 从主查询结果中查找某个 docId 和 chunkIndex 的原始 chunk（用于取 filename）
     *
     * @param classData  主查询结果数据列表，每个元素是一个包含文档信息的Map
     * @param docId      要查找的文档ID
     * @param chunkIndex 要查找的块索引
     * @return 找到的匹配chunk信息Map，如果未找到则返回null
     */
    private Map<String, Object> findChunkInClassData(List<Map<String, Object>> classData, String docId, int chunkIndex) {
        // 遍历所有chunk数据，查找匹配docId和chunkIndex的记录
        for (Map<String, Object> chunk : classData) {
            if (docId.equals(chunk.get("docId"))) {
                Number idxNum = (Number) chunk.get("chunkIndex");
                if (idxNum != null && idxNum.intValue() == chunkIndex) {
                    return chunk;
                }
            }
        }
        return null;
    }

}
