package com.rj.qdrant;

import io.qdrant.client.QdrantClient;
import io.qdrant.client.QdrantGrpcClient;
import io.qdrant.client.grpc.Collections.Distance;
import io.qdrant.client.grpc.Collections.VectorParams;
import io.qdrant.client.grpc.Points.PointStruct;
import io.qdrant.client.grpc.Points.ScoredPoint;
import io.qdrant.client.grpc.Points.SearchPoints;
import io.qdrant.client.grpc.Points.WithPayloadSelector;
import io.qdrant.client.grpc.QdrantOuterClass.HealthCheckReply;
import org.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.Assumptions;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.parallel.Execution;
import org.junit.jupiter.api.parallel.ExecutionMode;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.time.Duration;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;

import static io.qdrant.client.PointIdFactory.id;
import static io.qdrant.client.ValueFactory.value;
import static io.qdrant.client.VectorsFactory.vectors;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertNotNull;

/**
 * 实战向检索示例（可运行的 JUnit 测试）。
 * <p>
 * 流程：准备中文样本段落 → 为每条构造「主题块」模拟向量（同主题向量在空间中彼此接近）→ 写入 Qdrant（Cosine）
 * → 用查询向量做 topK 搜索 → 日志输出排名与相似度 score。
 * </p>
 * <p>
 * 真实业务中应把 {@link SampleChunk#vector()} 替换为 Embedding 模型（如 bge、OpenAI、DashScope 等）对 {@code title + text} 的编码结果；
 * 本例仅用确定性数学向量演示检索与打分行为。
 * </p>
 * <p>
 * 连接：{@code -Dqdrant.host}、{@code -Dqdrant.grpc.port}（默认与 {@link QdrantClientVerificationTest} 一致的环境可改属性）。
 * </p>
 */
@Execution(ExecutionMode.SAME_THREAD)
class QdrantPracticalSearchScenarioTest {

    private static final Logger log = LoggerFactory.getLogger(QdrantPracticalSearchScenarioTest.class);

    private static final String HOST = System.getProperty("qdrant.host", "192.168.1.44");
    private static final int GRPC_PORT = Integer.getInteger("qdrant.grpc.port", 6335);
    private static final Duration RPC_TIMEOUT = Duration.ofSeconds(35);

    /** 向量维度：按块划分主题，便于无模型情况下的可解释模拟 */
    private static final int DIM = 48;

    private static final int BLOCK = 16;
    private static final int THEME_SERVICE = 0;
    private static final int THEME_PRODUCT = 1;
    private static final int THEME_OTHER = 2;

    private static final String COLLECTION = "java_practical_" + UUID.randomUUID().toString().replace("-", "");

    private static QdrantClient client;

    /**
     * 一条可检索样本：业务上对应「一段可被向量化的文本」及其向量。
     */
    private record SampleChunk(long pointId, String title, String text, String category, int theme, int salt) {
        float[] vector() {
            return docEmbedding(theme, salt);
        }
    }

    @BeforeAll
    static void connectAndCreateCollection() throws Exception {
        log.info("======== Qdrant 实战示例：连接并建库 ========");
        log.info("gRPC {}:{} ，集合: {} ，维度: {}", HOST, GRPC_PORT, COLLECTION, DIM);
        try {
            client = new QdrantClient(QdrantGrpcClient.newBuilder(HOST, GRPC_PORT, false).build());
            HealthCheckReply health =
                    client.healthCheckAsync(RPC_TIMEOUT).get(RPC_TIMEOUT.toMillis(), TimeUnit.MILLISECONDS);
            assertNotNull(health);
            assertFalse(health.getTitle().isBlank());

            VectorParams params =
                    VectorParams.newBuilder().setSize(DIM).setDistance(Distance.Cosine).build();
            long rpcMs = RPC_TIMEOUT.toMillis();
            boolean exists =
                    client.collectionExistsAsync(COLLECTION, RPC_TIMEOUT).get(rpcMs, TimeUnit.MILLISECONDS);
            if (exists) {
                client.deleteCollectionAsync(COLLECTION, RPC_TIMEOUT).get(rpcMs, TimeUnit.MILLISECONDS);
            }
            client.createCollectionAsync(COLLECTION, params, RPC_TIMEOUT).get(rpcMs, TimeUnit.MILLISECONDS);
            log.info("[成功] 集合已就绪: {}", COLLECTION);
        } catch (ExecutionException | TimeoutException | InterruptedException e) {
            log.warn("无法连接 Qdrant: {}", e.getMessage());
            Assumptions.abort("Qdrant 不可达 (%s:%d): %s".formatted(HOST, GRPC_PORT, e.getMessage()));
        }
    }

    @AfterAll
    static void tearDown() throws Exception {
        if (client == null) {
            return;
        }
        try {
            client.deleteCollectionAsync(COLLECTION, RPC_TIMEOUT).get(RPC_TIMEOUT.toMillis(), TimeUnit.MILLISECONDS);
            log.info("已删除集合 {}", COLLECTION);
        } catch (Exception e) {
            log.info("清理集合时忽略: {}", e.getMessage());
        }
        client.close();
        log.info("======== Qdrant 实战示例结束 ========");
    }

    @Test
    void ingestSampleCorpus_andRunSimilaritySearch() throws Exception {
        List<SampleChunk> corpus = buildSampleCorpus();
        log.info("--- 1) 样本数据（{} 条）---", corpus.size());
        for (SampleChunk c : corpus) {
            log.info("  id={} | {} | {}", c.pointId, c.category, c.title);
        }

        List<PointStruct> points = new ArrayList<>();
        for (SampleChunk c : corpus) {
            float[] v = c.vector();
            points.add(
                    PointStruct.newBuilder()
                            .setId(id(c.pointId))
                            .setVectors(vectors(v))
                            .putAllPayload(
                                    Map.of(
                                            "title", value(c.title),
                                            "category", value(c.category),
                                            "text", value(c.text)))
                            .build());
        }

        long rpcMs = RPC_TIMEOUT.toMillis();
        log.info("--- 2) 写入 Qdrant（upsert）---");
        client.upsertAsync(COLLECTION, points, RPC_TIMEOUT).get(rpcMs, TimeUnit.MILLISECONDS);
        long n =
                client.countAsync(COLLECTION, RPC_TIMEOUT).get(rpcMs, TimeUnit.MILLISECONDS);
        assertEquals(corpus.size(), n);
        log.info("[成功] 已写入 {} 条向量点", n);

        // 查询1：向量贴近「客服 / 退换货」主题块
        float[] queryService = queryVectorForTheme(THEME_SERVICE, 901);
        log.info("--- 3) 模拟实战查询 A：主题≈智能客服（退换货、售后）---");
        runSearchAndLog("查询A-客服场景", queryService, 8);

        List<ScoredPoint> hitsA =
                client
                        .searchAsync(
                                SearchPoints.newBuilder()
                                        .setCollectionName(COLLECTION)
                                        .addAllVector(floatList(queryService))
                                        .setLimit(5)
                                        .setWithPayload(WithPayloadSelector.newBuilder().setEnable(true).build())
                                        .build(),
                                RPC_TIMEOUT)
                        .get(rpcMs, TimeUnit.MILLISECONDS);
        assertFalse(hitsA.isEmpty());
        for (int i = 0; i < Math.min(3, hitsA.size()); i++) {
            String cat = hitsA.get(i).getPayloadMap().get("category").getStringValue();
            assertEquals(
                    "智能客服",
                    cat,
                    "前 3 条命中应均为「智能客服」类样本（依赖模拟向量主题块）");
        }

        // 查询2：向量贴近「产品知识」主题块
        float[] queryProduct = queryVectorForTheme(THEME_PRODUCT, 902);
        log.info("--- 4) 模拟实战查询 B：主题≈产品规格 / 参数 ---");
        runSearchAndLog("查询B-产品知识", queryProduct, 8);

        List<ScoredPoint> hitsB =
                client
                        .searchAsync(
                                SearchPoints.newBuilder()
                                        .setCollectionName(COLLECTION)
                                        .addAllVector(floatList(queryProduct))
                                        .setLimit(3)
                                        .setWithPayload(WithPayloadSelector.newBuilder().setEnable(true).build())
                                        .build(),
                                RPC_TIMEOUT)
                        .get(rpcMs, TimeUnit.MILLISECONDS);
        assertFalse(hitsB.isEmpty());
        assertEquals(
                "产品知识",
                hitsB.get(0).getPayloadMap().get("category").getStringValue(),
                "Top1 应为产品类文档");
        log.info("--- 全部断言通过 ---");
    }

    private static void runSearchAndLog(String queryLabel, float[] queryVector, int topK) throws Exception {
        long rpcMs = RPC_TIMEOUT.toMillis();
        List<ScoredPoint> hits =
                client
                        .searchAsync(
                                SearchPoints.newBuilder()
                                        .setCollectionName(COLLECTION)
                                        .addAllVector(floatList(queryVector))
                                        .setLimit(topK)
                                        .setWithPayload(WithPayloadSelector.newBuilder().setEnable(true).build())
                                        .build(),
                                RPC_TIMEOUT)
                        .get(rpcMs, TimeUnit.MILLISECONDS);

        log.info("[{}] top{} 检索结果（Cosine 距离下 score 越大通常越相似，以服务端返回为准）:", queryLabel, topK);
        int rank = 1;
        for (ScoredPoint p : hits) {
            var payload = p.getPayloadMap();
            String title = payload.getOrDefault("title", value("")).getStringValue();
            String category = payload.getOrDefault("category", value("")).getStringValue();
            String text = payload.getOrDefault("text", value("")).getStringValue();
            String preview = text.length() > 60 ? text.substring(0, 60) + "…" : text;
            log.info(
                    "  #{} | pointId={} | score={} | category={} | title={}",
                    rank++,
                    p.getId().getNum(),
                    p.getScore(),
                    category,
                    title);
            log.info("      snippet: {}", preview);
        }
    }

    private static List<Float> floatList(float[] v) {
        List<Float> list = new ArrayList<>(v.length);
        for (float x : v) {
            list.add(x);
        }
        return list;
    }

    private static List<SampleChunk> buildSampleCorpus() {
        List<SampleChunk> list = new ArrayList<>();
        long id = 1001;
        list.add(
                new SampleChunk(
                        id++,
                        "退换货政策说明",
                        "签收七日内商品未使用且包装完好，可申请无理由退货；质量问题由我方承担运费。",
                        "智能客服",
                        THEME_SERVICE,
                        1));
        list.add(
                new SampleChunk(
                        id++,
                        "物流进度查询",
                        "请提供订单号，可在「我的订单」中查看实时物流节点与预计送达时间。",
                        "智能客服",
                        THEME_SERVICE,
                        2));
        list.add(
                new SampleChunk(
                        id++,
                        "电子发票开具",
                        "订单完成后 24 小时内可申请增值税普通发票，发送至预留邮箱。",
                        "智能客服",
                        THEME_SERVICE,
                        3));
        list.add(
                new SampleChunk(
                        id++,
                        "旗舰手机显示屏参数",
                        "采用 6.7 英寸 AMOLED，分辨率 3200×1440，支持 120Hz 高刷与 HDR10+。",
                        "产品知识",
                        THEME_PRODUCT,
                        4));
        list.add(
                new SampleChunk(
                        id++,
                        "整机保修期限",
                        "主机保修 24 个月，电池与充电器保修 12 个月，人为损坏不在保修范围。",
                        "产品知识",
                        THEME_PRODUCT,
                        5));
        list.add(
                new SampleChunk(
                        id++,
                        "周末天气闲聊",
                        "周六多云转晴，气温 18～26℃，适合户外活动。",
                        "闲聊",
                        THEME_OTHER,
                        6));
        return list;
    }

    /** 主题块基底：不同主题占据不同维度区间，便于无 Embedding 时的聚类演示 */
    private static float[] topicTemplate(int theme) {
        float[] v = new float[DIM];
        int start = theme * BLOCK;
        for (int i = 0; i < BLOCK; i++) {
            v[start + i] = 1f + 0.03f * (float) Math.sin(i + theme * 17);
        }
        l2Normalize(v);
        return v;
    }

    private static float[] docEmbedding(int theme, int salt) {
        float[] v = topicTemplate(theme);
        for (int i = 0; i < DIM; i++) {
            v[i] += 0.02f * (float) Math.sin(salt * 13L + i * 5);
        }
        l2Normalize(v);
        return v;
    }

    /** 查询向量：与某主题对齐并加独立扰动，模拟「用户问题」经 Embedding 后落在主题附近 */
    private static float[] queryVectorForTheme(int theme, int querySalt) {
        float[] v = topicTemplate(theme);
        for (int i = 0; i < DIM; i++) {
            v[i] += 0.04f * (float) Math.sin(querySalt * 7L + i * 3);
        }
        l2Normalize(v);
        return v;
    }

    private static void l2Normalize(float[] v) {
        double sum = 0;
        for (float x : v) {
            sum += (double) x * x;
        }
        if (sum < 1e-12) {
            return;
        }
        float inv = (float) (1.0 / Math.sqrt(sum));
        for (int i = 0; i < v.length; i++) {
            v[i] *= inv;
        }
    }
}
