package com.ruoyi.ai.config;

import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.pinecone.PineconeEmbeddingStore;
import dev.langchain4j.store.embedding.pinecone.PineconeServerlessIndexConfig;
import io.pinecone.clients.Index;
import io.pinecone.clients.Pinecone;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.io.IOException;
import java.util.function.Supplier;

/**
 * 向量存储配置
 */
@Slf4j
@Configuration
public class EmbeddingStoreConfig {

    /** Pinecone 连接重试次数：构造期会同步发起 HTTP/2 握手，偶发 "Required SETTINGS preface not received" */
    private static final int CONNECT_ATTEMPTS = 3;

    /** 重试间隔基数（毫秒），按次数线性退避 */
    private static final long RETRY_BACKOFF_MILLIS = 1500L;

    @Value("${pinecone.api-key:pcsk_4SJLnh_tNB3wSLJU8tc4E5P28PcXX8eCLdURqZpVhg1FMV8CRYxjneWdzqRdB5Ftqooi9}")
    private String pineconeApiKey;

    @Value("${pinecone.index:xiaozhi-index}")
    private String indexName;

    @Value("${pinecone.namespace:knowledge-base}")
    private String namespace;

    @Bean
    public Pinecone pinecone() {
        return new Pinecone.Builder(pineconeApiKey).build();
    }

    @Bean
    public Index pineconeIndex(Pinecone pinecone) {
        return retry("获取 Pinecone 索引连接", () -> pinecone.getIndexConnection(indexName));
    }

    @Bean
    public EmbeddingStore<TextSegment> embeddingStore(EmbeddingModel embeddingModel) {
        return retry("初始化 Pinecone 向量存储", () -> PineconeEmbeddingStore.builder()
                .apiKey(pineconeApiKey)
                .index(indexName)
                .nameSpace(namespace)
                .createIndex(PineconeServerlessIndexConfig.builder()
                        .cloud("AWS")
                        .region("us-east-1")
                        .dimension(embeddingModel.dimension())
                        .build())
                .build());
    }

    /**
     * 对 Pinecone 启动期连接做有限次重试。
     * <p>
     * Pinecone 客户端固定走 HTTP/2，构造 EmbeddingStore / 获取索引连接时会同步调用控制面接口，
     * 偶发握手失败（java.io.IOException: Required SETTINGS preface not received）会以运行时异常抛出，
     * 直接导致整个 ApplicationContext 加载失败。这里仅对这类瞬时网络异常重试，其它异常立即抛出。
     */
    private <T> T retry(String action, Supplier<T> supplier) {
        RuntimeException last = null;
        for (int attempt = 1; attempt <= CONNECT_ATTEMPTS; attempt++) {
            try {
                return supplier.get();
            } catch (RuntimeException e) {
                if (!isTransientConnectFailure(e) || attempt == CONNECT_ATTEMPTS) {
                    throw e;
                }
                last = e;
                log.warn("{}失败（第 {}/{} 次），{}ms 后重试：{}",
                        action, attempt, CONNECT_ATTEMPTS, RETRY_BACKOFF_MILLIS * attempt, rootMessage(e));
                sleep(RETRY_BACKOFF_MILLIS * attempt);
            }
        }
        throw last;
    }

    /** 沿 cause 链判断是否为 HTTP/2 握手类的瞬时连接失败 */
    private boolean isTransientConnectFailure(Throwable e) {
        for (Throwable t = e; t != null; t = t.getCause()) {
            if (t instanceof IOException) {
                return true;
            }
            String message = t.getMessage();
            if (message != null && (message.contains("SETTINGS preface")
                    || message.contains("Connection reset")
                    || message.contains("timed out"))) {
                return true;
            }
        }
        return false;
    }

    private String rootMessage(Throwable e) {
        Throwable root = e;
        while (root.getCause() != null && root.getCause() != root) {
            root = root.getCause();
        }
        return root.getMessage();
    }

    private void sleep(long millis) {
        try {
            Thread.sleep(millis);
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
        }
    }
}