package org.dromara.knowledge.vectorstore;

import lombok.RequiredArgsConstructor;
import org.dromara.knowledge.domain.KnowledgeFragment;
import org.dromara.knowledge.domain.KnowledgeVector;
import org.dromara.knowledge.domain.vo.KnowledgeSearchResultVo;
import org.dromara.knowledge.mapper.KnowledgeVectorMapper;
import org.dromara.knowledge.service.KnowledgeFragmentService;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.stream.Collectors;

@Component
@RequiredArgsConstructor
public class DbVectorStoreService implements VectorStoreService {

    private final KnowledgeVectorMapper mapper;
    private final KnowledgeFragmentService fragmentService;

    @Override
    public void upsert(String knowledgeId, String docId, String fragmentId, String vectorId, double[] vector, String content, String modelName) {
        KnowledgeVector kv = new KnowledgeVector();
        kv.setKnowledgeId(knowledgeId);
        kv.setDocId(docId);
        kv.setFragmentId(fragmentId);
        kv.setVectorId(vectorId);
        kv.setVectorData(Arrays.toString(vector));
        kv.setDimension(vector.length);
        kv.setModelName(modelName);
        kv.setStatus(20);
        mapper.insert(kv);
    }

    @Override
    public List<KnowledgeSearchResultVo> search(String knowledgeId, double[] queryVector, String modelName, Integer topK, Double threshold) {
        List<KnowledgeVector> vectors = mapper.selectList(new com.baomidou.mybatisplus.core.conditions.query.QueryWrapper<KnowledgeVector>().lambda()
                .eq(KnowledgeVector::getKnowledgeId, knowledgeId)
                .eq(KnowledgeVector::getModelName, modelName));
        List<KnowledgeSearchResultVo> results = new ArrayList<>();
        for (KnowledgeVector v : vectors) {
            double[] vec = parse(v.getVectorData());
            if (vec == null || vec.length != queryVector.length) continue;
            double sim = cosine(queryVector, vec);
            if (threshold != null && sim < threshold) continue;
            KnowledgeFragment frag = fragmentService.getById(v.getFragmentId());
            if (frag == null) continue;
            KnowledgeSearchResultVo r = new KnowledgeSearchResultVo();
            r.setFragmentId(v.getFragmentId());
            r.setDocId(v.getDocId());
            r.setContent(frag.getContent());
            r.setSimilarityScore(sim);
            r.setFragmentIndex(frag.getIdx());
            r.setCreateTime(frag.getCreateTime());
            results.add(r);
        }
        results.sort((a, b) -> Double.compare(b.getSimilarityScore(), a.getSimilarityScore()));
        for (int i = 0; i < results.size(); i++) results.get(i).setRanking(i + 1);
        int limit = topK != null ? topK : 10;
        return results.stream().limit(limit).collect(Collectors.toList());
    }

    @Override
    public void deleteByKnowledgeId(String knowledgeId) {
        mapper.delete(new com.baomidou.mybatisplus.core.conditions.query.QueryWrapper<KnowledgeVector>().lambda()
                .eq(KnowledgeVector::getKnowledgeId, knowledgeId));
    }

    @Override
    public void rebuildIndex(String knowledgeId) {
    }

    private double[] parse(String s) {
        if (s == null || s.isBlank()) return null;
        String clean = s.replace("[", "").replace("]", "");
        String[] parts = clean.split(",");
        double[] v = new double[parts.length];
        for (int i = 0; i < parts.length; i++) v[i] = Double.parseDouble(parts[i].trim());
        return v;
    }

    private double cosine(double[] a, double[] b) {
        double dot = 0, n1 = 0, n2 = 0;
        for (int i = 0; i < a.length; i++) {
            dot += a[i] * b[i];
            n1 += a[i] * a[i];
            n2 += b[i] * b[i];
        }
        if (n1 == 0 || n2 == 0) return 0;
        return dot / (Math.sqrt(n1) * Math.sqrt(n2));
    }
}
