package org.dromara.system.ai.service.impl;

import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.weaviate.WeaviateEmbeddingStore;
import io.weaviate.client.Config;
import io.weaviate.client.WeaviateClient;
import io.weaviate.client.base.Result;
import io.weaviate.client.v1.data.model.WeaviateObject;
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 org.dromara.common.core.utils.StringUtils;
import org.dromara.system.XinferenceClient;
import org.dromara.system.ai.domain.bo.AiStoreEmbeddingBo;
import org.dromara.system.ai.service.VectorStoreService;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.dromara.system.domain.AiKnowledgeAttachDo;
import org.dromara.system.domain.vo.AiKnowledgeAttachVo;
import org.dromara.system.mapper.AiKnowledgeAttachMapper;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;


import java.io.IOException;
import java.lang.reflect.Field;
import java.util.*;

/**
 * @author destiny
 * @description:
 * @date 2025/9/17 18:15
 */
@Slf4j
@Service
@RequiredArgsConstructor
public class VectorStoreServiceImpl implements VectorStoreService {

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

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

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

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

    private EmbeddingStore<TextSegment> embeddingStore;

    private final AiKnowledgeAttachMapper aiKnowledgeAttachMapper;

    @Override
    public void storeEmbeddings(AiStoreEmbeddingBo storeEmbeddingBo) {
        String kid = storeEmbeddingBo.getKid();
        String fullClassName = weaviateClassName + kid;

        // 初始化 Weaviate 客户端
        WeaviateClient client = new WeaviateClient(new Config(weaviateProtocol, String.format("%s:%s", weaviateHost, weaviatePort)));
        createSchema(kid, storeEmbeddingBo.getVectorModelName());

        List<String> chunkList = storeEmbeddingBo.getChunkList();
        String docId = storeEmbeddingBo.getDocId();
        if (StringUtils.isBlank(docId)) {
            throw new IllegalArgumentException("docId 不能为空");
        }

        XinferenceClient xinferenceClient = new XinferenceClient(storeEmbeddingBo.getBaseUrl());
        String modelName = storeEmbeddingBo.getEmbeddingModelName();
        if ("quentinz/bge-large-zh-v1.5".contains(modelName)) {
            modelName = "bge-large-zh-v1.5";
        }

        for (int i = 0; i < chunkList.size(); i++) {
            String text = chunkList.get(i);
            if (text == null || text.trim().isEmpty()) {
                System.out.println("跳过空文本: " + i);
                continue;
            }
            List<Double> vector;
            try {
                vector = xinferenceClient.getEmbedding(modelName, text);
            } catch (IOException e) {
                throw new RuntimeException("生成向量失败: " + e.getMessage(), e);
            }

            // 创建 Weaviate 对象
            Map<String, Object> properties = new HashMap<>();
            properties.put("text", text);
            properties.put("indexFilterable", true);
            properties.put("indexSearchable", true);
            properties.put("docId", docId);
            properties.put("filename", storeEmbeddingBo.getFileName());
            properties.put("chunkIndex", i);

            HashMap<String, Object> hashMap = new HashMap<>();
            hashMap.put("docId", docId);
            hashMap.put("filename", storeEmbeddingBo.getFileName());
            hashMap.put("chunkIndex", i);
            ObjectMapper mapper = new ObjectMapper();
            try {
                String jsonString = mapper.writeValueAsString(hashMap);
                properties.put("metadata", jsonString);
            } catch (JsonProcessingException e) {
                throw new RuntimeException(e);
            }


            try {
                Result<WeaviateObject> result = client.data().creator()
                    .withClassName(fullClassName)
                    .withProperties(properties)
                    .withVector(toFloatArray(vector))
                    .run();
                if (result.hasErrors()) {
                    throw new RuntimeException("存储对象失败: " + result.getError().toString());
                }
            } catch (Exception e) {
                throw new RuntimeException("存储向量失败: " + e.getMessage(), e);
            }
        }
    }

    @Override
    public void createSchema(String kid, String modelName) {
        String fullClassName = weaviateClassName + kid;

        // Initialize Weaviate client
        WeaviateClient client = new WeaviateClient(new Config(weaviateProtocol, String.format("%s:%s", weaviateHost, weaviatePort)));

        // Check if schema exists
        Result<Schema> existingSchema = client.schema().getter().run();
        if (existingSchema.hasErrors()) {
            throw new RuntimeException("Failed to get schema: " + existingSchema.getError().toString());
        }

        // Check if class exists
        boolean classExists = existingSchema.getResult().getClasses().stream()
            .anyMatch(c -> c.getClassName().equals(fullClassName));
        if (!classExists) {
            WeaviateClass weaviateClass = WeaviateClass.builder()
                .className(fullClassName)
                .properties(Arrays.asList(
                    Property.builder().name("text").dataType(Arrays.asList("text")).build(),
                    Property.builder().name("indexFilterable").dataType(Arrays.asList("boolean")).build(),
                    Property.builder().name("indexSearchable").dataType(Arrays.asList("boolean")).build(),
                    Property.builder().name("docId").dataType(Arrays.asList("string")).build(),
                    Property.builder().name("filename").dataType(Arrays.asList("string")).build(),
                    Property.builder().name("metadata").dataType(Arrays.asList("string")).build(),
                    Property.builder().name("chunkIndex").dataType(Arrays.asList("int")).build()
                ))
                .vectorizer("none")
                .build();
            Schema.builder()
                .classes(Collections.singletonList(weaviateClass))
                .build();
            client.schema().classCreator().withClass(weaviateClass).run();
        } else {
            WeaviateClass existingClass = existingSchema.getResult().getClasses().stream()
                .filter(c -> c.getClassName().equals(fullClassName))
                .findFirst()
                .orElse(null);
            boolean hasDocId = existingClass.getProperties().stream()
                .anyMatch(p -> p.getName().equals("docId"));
            if (!hasDocId) {
                Property docIdProperty = Property.builder()
                    .name("docId")
                    .dataType(Arrays.asList("string"))
                    .build();
                client.schema().propertyCreator()
                    .withClassName(fullClassName)
                    .withProperty(docIdProperty)
                    .run();
            }
            boolean filename = existingClass.getProperties().stream()
                .anyMatch(p -> p.getName().equals("filename"));
            if (!filename) {
                Property docIdProperty = Property.builder()
                    .name("filename")
                    .dataType(Arrays.asList("string"))
                    .build();
                client.schema().propertyCreator()
                    .withClassName(fullClassName)
                    .withProperty(docIdProperty)
                    .run();
            }
            boolean chunkIndex = existingClass.getProperties().stream()
                .anyMatch(p -> p.getName().equals("chunkIndex"));
            if (!chunkIndex) {
                Property docIdProperty = Property.builder()
                    .name("chunkIndex")
                    .dataType(Arrays.asList("int"))
                    .build();
                client.schema().propertyCreator()
                    .withClassName(fullClassName)
                    .withProperty(docIdProperty)
                    .run();
            }

        }

        embeddingStore = WeaviateEmbeddingStore.builder()
            .scheme(weaviateProtocol)
            .host(weaviateHost)
            .objectClass(fullClassName)
            .avoidDups(true)
            .consistencyLevel("ALL")
            .build();
    }

    public static Float[] toFloatArray(List<Double> doubleList) {
        Float[] floatArray = new Float[doubleList.size()];
        for (int i = 0; i < doubleList.size(); i++) {
            floatArray[i] = doubleList.get(i).floatValue();
        }
        return floatArray;
    }

    public void removeByDocId(String docId) {
        AiKnowledgeAttachVo aiKnowledgeAttachVo = aiKnowledgeAttachMapper.selectVoOne(
            new LambdaQueryWrapper<AiKnowledgeAttachDo>().eq(AiKnowledgeAttachDo::getDocCode, docId));
        String kid = aiKnowledgeAttachVo.getKnowledgeCode();
        createSchema(kid, "");

        try {
            // 反射拿到私有字段 client
            Field clientField = embeddingStore.getClass().getDeclaredField("client");
            clientField.setAccessible(true);
            Object client = clientField.get(embeddingStore);

            // 反射拿到私有字段 objectClass
            Field objectClassField = embeddingStore.getClass().getDeclaredField("objectClass");
            objectClassField.setAccessible(true);
            String objectClass = (String) objectClassField.get(embeddingStore);

            // 下面需要强转 client 为 io.weaviate.client.WeaviateClient
            WeaviateClient weaviateClient = (WeaviateClient) client;

            // 构造删除请求
            weaviateClient.batch()
                .objectsBatchDeleter()
                .withClassName(objectClass)
                .withWhere(io.weaviate.client.v1.filters.WhereFilter.builder()
                    .path("docId")
                    .operator("Equal")
                    .valueText(docId)
                    .build())
                .run();

        } catch (NoSuchFieldException | IllegalAccessException e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    public void removeById(String id, String modelName) {
        createSchema(id, modelName);
        // 根据条件删除向量数据
        embeddingStore.remove(id);
    }
}
