package org.dromara.knowledge.vectorstore;

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 WeaviateVectorStoreService implements VectorStoreService {
    @Value("${ruoyi.ai.knowledge.vector-store.weaviate.base-url:http://localhost:8080}")
    private String baseUrl;
    @Value("${ruoyi.ai.knowledge.vector-store.weaviate.api-key:}")
    private String apiKey;
    @Value("${ruoyi.ai.knowledge.vector-store.weaviate.class-name:KnowledgeVector}")
    private String className;

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

    public WeaviateVectorStoreService(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> obj = new HashMap<>();
        obj.put("class", className);
        obj.put("id", vectorId);
        obj.put("properties", Map.of(
                "knowledgeId", knowledgeId,
                "docId", docId,
                "fragmentId", fragmentId,
                "content", content,
                "model", modelName
        ));
        List<Double> vec = new ArrayList<>();
        for (double v : vector) vec.add(v);
        obj.put("vector", vec);
        String body = JsonUtils.toJsonString(obj);
        post("/v1/objects", 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);
        String gql = "{" +
                "  Get {" +
                "    " + className + "(nearVector: {vector: " + JsonUtils.toJsonString(vec) + "}, limit: " + (topK != null ? topK : 10) + ", where: {path: [\"knowledgeId\"], operator: Equal, valueString: \"" + knowledgeId + "\"}) {" +
                "      fragmentId docId content _additional { distance id }" +
                "    }" +
                "  }" +
                "}";
        String body = JsonUtils.toJsonString(Map.of("query", gql));
        String resp = post("/v1/graphql", body);
        Map<?, ?> parsed = JsonUtils.parseObject(resp, new com.fasterxml.jackson.core.type.TypeReference<Map<String, Object>>() {
        });
        Map<?, ?> data = parsed != null ? (Map<?, ?>) parsed.get("data") : null;
        Map<?, ?> get = data != null ? (Map<?, ?>) data.get("Get") : null;
        List<?> list = get != null ? (List<?>) get.get(className) : null;
        List<KnowledgeSearchResultVo> results = new ArrayList<>();
        if (list != null) {
            int rank = 1;
            for (Object o : list) {
                if (!(o instanceof Map<?, ?> m)) continue;
                KnowledgeSearchResultVo vo = new KnowledgeSearchResultVo();
                vo.setFragmentId(String.valueOf(m.get("fragmentId")));
                vo.setDocId(String.valueOf(m.get("docId")));
                String content = String.valueOf(m.get("content"));
                if (content == null || content.equals("null")) {
                    KnowledgeFragment frag = fragmentService.getById(vo.getFragmentId());
                    if (frag != null) {
                        content = frag.getContent();
                        vo.setFragmentIndex(frag.getIdx());
                        vo.setCreateTime(frag.getCreateTime());
                    }
                } else {
                    KnowledgeFragment frag = fragmentService.getById(vo.getFragmentId());
                    if (frag != null) {
                        vo.setFragmentIndex(frag.getIdx());
                        vo.setCreateTime(frag.getCreateTime());
                    }
                }
                vo.setContent(content);
                Map<?, ?> add = (Map<?, ?>) m.get("_additional");
                Double dist = add != null ? toDouble(add.get("distance")) : null;
                double score = dist != null ? (1.0 - Math.min(1.0, dist)) : 0.0;
                if (threshold != null && score < threshold) continue;
                vo.setSimilarityScore(score);
                vo.setRanking(rank++);
                results.add(vo);
            }
        }
        return results;
    }

    @Override
    public void deleteByKnowledgeId(String knowledgeId) {
        String gql = "{ " +
                "  Aggregate { " +
                "    " + className + "(where: {path: [\"knowledgeId\"], operator: Equal, valueString: \"" + knowledgeId + "\"}) { meta { count } }" +
                "  }" +
                "}";
        post("/v1/graphql", JsonUtils.toJsonString(Map.of("query", gql)));
    }

    @Override
    public void rebuildIndex(String knowledgeId) {
    }

    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("Authorization", "Bearer " + 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;
        }
    }
}
