package org.dromara.knowledge.vectorstore;

import cn.hutool.core.util.StrUtil;
import okhttp3.*;
import org.dromara.common.json.utils.JsonUtils;
import org.dromara.knowledge.domain.KnowledgeFragment;
import org.dromara.knowledge.domain.vo.KnowledgeSearchResultVo;
import org.dromara.knowledge.service.KnowledgeFragmentService;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

import java.nio.charset.StandardCharsets;
import java.util.*;

@Component
public class QdrantVectorStoreService implements VectorStoreService {
    @Value("${ruoyi.ai.knowledge.vector-store.qdrant.base-url:http://localhost:6333}")
    private String baseUrl;
    @Value("${ruoyi.ai.knowledge.vector-store.qdrant.api-key:}")
    private String apiKey;
    @Value("${ruoyi.ai.knowledge.vector-store.qdrant.collection:knowledge_vectors}")
    private String collection;
    @Value("${ruoyi.ai.knowledge.vector-store.qdrant.distance:Cosine}")
    private String distance;
    @Value("${ruoyi.ai.knowledge.embedding.default-model:bge-large-zh}")
    private String defaultModel;

    private final OkHttpClient client = new OkHttpClient();
    private final KnowledgeFragmentService fragmentService;

    public QdrantVectorStoreService(KnowledgeFragmentService fragmentService) {
        this.fragmentService = fragmentService;
    }

    @Override
    public void upsert(String knowledgeId, String docId, String fragmentId, String vectorId, double[] vector, String content, String modelName) {
        Map<String, Object> payload = new HashMap<>();
        payload.put("knowledgeId", knowledgeId);
        payload.put("docId", docId);
        payload.put("fragmentId", fragmentId);
        payload.put("content", content);
        payload.put("model", modelName);
        List<Double> vec = new ArrayList<>();
        for (double v : vector) vec.add(v);
        Map<String, Object> point = new HashMap<>();
        point.put("id", vectorId);
        point.put("vector", vec);
        point.put("payload", payload);
        Map<String, Object> body = new HashMap<>();
        body.put("points", List.of(point));
        post("/collections/" + collection + "/points?wait=true", JsonUtils.toJsonString(body));
    }

    @Override
    public List<KnowledgeSearchResultVo> search(String knowledgeId, double[] queryVector, String modelName, Integer topK, Double threshold) {
        List<Double> vec = new ArrayList<>();
        for (double v : queryVector) vec.add(v);
        Map<String, Object> filter = Map.of("must", List.of(Map.of("key", "knowledgeId", "match", Map.of("value", knowledgeId))));
        Map<String, Object> body = new HashMap<>();
        body.put("vector", vec);
        body.put("filter", filter);
        body.put("limit", topK != null ? topK : 10);
        String resp = post("/collections/" + collection + "/points/search", JsonUtils.toJsonString(body));
        List<KnowledgeSearchResultVo> results = new ArrayList<>();
        Map<?, ?> parsed = JsonUtils.parseObject(resp, new com.fasterxml.jackson.core.type.TypeReference<Map<String, Object>>() {
        });
        Object r = parsed != null ? parsed.get("result") : null;
        if (r instanceof List<?> list) {
            int rank = 1;
            for (Object item : list) {
                if (!(item instanceof Map<?, ?> m)) continue;
                Double score = toDouble(m.get("score"));
                Map<?, ?> payload = (Map<?, ?>) m.get("payload");
                if (threshold != null && score != null && score < threshold) continue;
                KnowledgeSearchResultVo vo = new KnowledgeSearchResultVo();
                if (payload != null) {
                    String fragId = String.valueOf(payload.get("fragmentId"));
                    vo.setFragmentId(fragId);
                    vo.setDocId(String.valueOf(payload.get("docId")));
                    Object c = payload.get("content");
                    String content = c != null ? String.valueOf(c) : null;
                    if (StrUtil.isBlank(content)) {
                        var frag = fragmentService.getById(fragId);
                        if (frag != null) {
                            content = frag.getContent();
                            vo.setFragmentIndex(frag.getIdx());
                            vo.setCreateTime(frag.getCreateTime());
                        }
                    } else {
                        KnowledgeFragment frag = fragmentService.getById(fragId);
                        if (frag != null) {
                            vo.setFragmentIndex(frag.getIdx());
                            vo.setCreateTime(frag.getCreateTime());
                        }
                    }
                    vo.setContent(content);
                }
                vo.setSimilarityScore(score != null ? score : 0);
                vo.setRanking(rank++);
                results.add(vo);
            }
        }
        return results;
    }

    @Override
    public void deleteByKnowledgeId(String knowledgeId) {
        Map<String, Object> filter = Map.of("must", List.of(Map.of("key", "knowledgeId", "match", Map.of("value", knowledgeId))));
        Map<String, Object> body = new HashMap<>();
        body.put("filter", filter);
        post("/collections/" + collection + "/points/delete", JsonUtils.toJsonString(body));
    }

    @Override
    public void rebuildIndex(String knowledgeId) {
        ensureCollection(1024);
    }

    private void ensureCollection(int size) {
        String get = get("/collections/" + collection);
        if (get == null || get.contains("\"status\":\"error\"")) {
            Map<String, Object> vectors = Map.of("size", size, "distance", distance);
            Map<String, Object> body = Map.of("vectors", vectors);
            put("/collections/" + collection, JsonUtils.toJsonString(body));
        }
    }

    private String get(String path) {
        Request.Builder b = new Request.Builder().url(url(path)).get();
        if (apiKey != null && !apiKey.isBlank()) b.header("api-key", apiKey);
        try (Response resp = client.newCall(b.build()).execute()) {
            if (resp.body() == null) return null;
            return resp.body().string();
        } catch (Exception e) {
            return null;
        }
    }

    private String put(String path, String json) {
        Request.Builder b = new Request.Builder().url(url(path))
                .put(RequestBody.create(MediaType.parse("application/json"), json.getBytes(StandardCharsets.UTF_8)));
        if (apiKey != null && !apiKey.isBlank()) b.header("api-key", apiKey);
        try (Response resp = client.newCall(b.build()).execute()) {
            return Objects.requireNonNull(resp.body()).string();
        } catch (Exception e) {
            return null;
        }
    }

    private String post(String path, String json) {
        Request.Builder b = new Request.Builder().url(url(path))
                .post(RequestBody.create(MediaType.parse("application/json"), json.getBytes(StandardCharsets.UTF_8)));
        if (apiKey != null && !apiKey.isBlank()) b.header("api-key", apiKey);
        try (Response resp = client.newCall(b.build()).execute()) {
            return Objects.requireNonNull(resp.body()).string();
        } catch (Exception e) {
            return null;
        }
    }

    private String url(String path) {
        if (baseUrl.endsWith("/") && path.startsWith("/")) return baseUrl + path.substring(1);
        if (!baseUrl.endsWith("/") && !path.startsWith("/")) return baseUrl + "/" + path;
        return baseUrl + path;
    }

    private Double toDouble(Object o) {
        try {
            if (o == null) return null;
            if (o instanceof Number n) return n.doubleValue();
            return Double.parseDouble(String.valueOf(o));
        } catch (Exception e) {
            return null;
        }
    }
}
