// 一段文本经过 Embedding 后的样子
String text = "Spring Boot 是 Java 生态最流行的微服务框架";
float[] embedding = embeddingModel.embed(text);
// embedding = [0.0234, -0.0891, 0.1456, ..., -0.0321]  // 1536 个浮点数

// Self-Attention 的直觉理解（伪代码）
// 对于句子中的每个词，计算它与其他所有词的相关性分数
for (Token query : tokens) {
    for (Token key : tokens) {
        double attention = dot(query.vector, key.vector) / sqrt(dimension);
        // "Spring" 和 "Boot" 的注意力分数高（经常一起出现）
        // "Spring" 和 "框架" 的注意力分数中等（语义相关）
        // "Spring" 和 "微" 的注意力分数低（关系弱）
    }
}

public static double cosineSimilarity(float[] a, float[] b) {
    double dotProduct = 0, normA = 0, normB = 0;
    for (int i = 0; i < a.length; i++) {
        dotProduct += a[i] * b[i];
        normA += a[i] * a[i];
        normB += b[i] * b[i];
    }
    return dotProduct / (Math.sqrt(normA) * Math.sqrt(normB));
}
// 结果范围: [-1, 1]，1 表示完全相同，0 表示无关，-1 表示完全相反

public static double euclideanDistance(float[] a, float[] b) {
    double sum = 0;
    for (int i = 0; i < a.length; i++) {
        sum += Math.pow(a[i] - b[i], 2);
    }
    return Math.sqrt(sum);
}
// 结果范围: [0, +∞)，0 表示完全相同

public static double dotProduct(float[] a, float[] b) {
    double sum = 0;
    for (int i = 0; i < a.length; i++) {
        sum += a[i] * b[i];
    }
    return sum;
}

// 1. 固定大小分块（最简单，但最粗糙）
public List<String> fixedSizeChunk(String text, int chunkSize, int overlap) {
    List<String> chunks = new ArrayList<>();
    for (int i = 0; i < text.length(); i += chunkSize - overlap) {
        chunks.add(text.substring(i, Math.min(i + chunkSize, text.length())));
    }
    return chunks;
}

// 2. 语义分块（按语义边界切分）
public List<String> semanticChunk(String text) {
    // 先按段落切分
    String[] paragraphs = text.split("\n\n");
    List<String> chunks = new ArrayList<>();
    StringBuilder current = new StringBuilder();
    
    for (String para : paragraphs) {
        // 如果当前块加上新段落超过阈值，先保存当前块
        if (current.length() + para.length() > 512) {
            if (current.length() > 0) {
                chunks.add(current.toString().trim());
                current = new StringBuilder();
            }
        }
        current.append(para).append("\n\n");
    }
    if (current.length() > 0) {
        chunks.add(current.toString().trim());
    }
    return chunks;
}

// 3. 递归分块（LangChain 推荐方式）
public List<String> recursiveChunk(String text, int chunkSize) {
    // 优先按大分隔符切分，不够再按小分隔符
    String[] separators = {"\n\n", "\n", "。", "，", " "};
    return recursiveSplit(text, separators, 0, chunkSize);
}

@Configuration
public class EmbeddingConfig {

    @Bean
    public EmbeddingModel embeddingModel() {
        // 使用 OpenAI
        return new OpenAiEmbeddingModel(
            new OpenAiApi(System.getenv("OPENAI_API_KEY"))
        );
    }
}

@Service
public class SemanticSearchService {

    @Autowired
    private EmbeddingModel embeddingModel;

    /**
     * 计算两段文本的语义相似度
     */
    public double similarity(String text1, String text2) {
        float[] vec1 = embeddingModel.embed(text1);
        float[] vec2 = embeddingModel.embed(text2);
        return cosineSimilarity(vec1, vec2);
    }

    /**
     * 批量 Embedding（更高效，减少 API 调用次数）
     */
    public List<float[]> batchEmbed(List<String> texts) {
        EmbeddingResponse response = embeddingModel.embedForResponse(texts);
        return response.getResults().stream()
            .map(r -> r.getOutput())
            .collect(Collectors.toList());
    }

    /**
     * 构建简单的内存向量索引
     */
    public List<SearchResult> search(String query, List<Document> documents, int topK) {
        float[] queryVec = embeddingModel.embed(query);
        
        return documents.stream()
            .map(doc -> new SearchResult(
                doc,
                cosineSimilarity(queryVec, doc.getEmbedding())
            ))
            .sorted((a, b) -> Double.compare(b.score(), a.score()))
            .limit(topK)
            .collect(Collectors.toList());
    }

    private double cosineSimilarity(float[] a, float[] b) {
        double dot = 0, normA = 0, normB = 0;
        for (int i = 0; i < a.length; i++) {
            dot += a[i] * b[i];
            normA += a[i] * a[i];
            normB += b[i] * b[i];
        }
        return dot / (Math.sqrt(normA) * Math.sqrt(normB));
    }
}

record Document(String id, String content, float[] embedding) {}
record SearchResult(Document document, double score) {}

// 配置 Embedding 模型
EmbeddingModel model = OpenAiEmbeddingModel.builder()
    .apiKey(System.getenv("OPENAI_API_KEY"))
    .modelName("text-embedding-3-small")
    .build();

// 单条 Embedding
Response<Embedding> response = model.embed("Spring Boot 微服务框架");
float[] vector = response.content().vector();

// 批量 Embedding
List<TextSegment> segments = List.of(
    TextSegment.from("Spring Boot 是 Java 微服务框架"),
    TextSegment.from("React 是前端 JavaScript 框架"),
    TextSegment.from("Kubernetes 是容器编排平台")
);
Response<List<Embedding>> batchResponse = model.embedAll(segments);

// 使用 EmbeddingStore 存储和检索
EmbeddingStore<TextSegment> store = new InMemoryEmbeddingStore<>();

// 添加向量
for (int i = 0; i < segments.size(); i++) {
    store.add(batchResponse.content().get(i), segments.get(i));
}

// 语义搜索
float[] queryEmbedding = model.embed("Java 后端开发").content().vector();
List<EmbeddingMatch<TextSegment>> results = store.findRelevant(
    Embedding.from(queryEmbedding), 3, 0.7  // top 3, 最低相似度 0.7
);

for (EmbeddingMatch<TextSegment> match : results) {
    System.out.printf("相似度: %.4f | 内容: %s%n", 
        match.score(), match.embedded().text());
}

// Spring AI + pgvector（PostgreSQL 扩展）
@Bean
public VectorStore vectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel) {
    return new PgVectorStore(jdbcTemplate, embeddingModel,
        PgVectorStore.PgVectorStoreOptions.builder()
            .dimensions(1536)
            .indexType(HNSW)  // 近似最近邻索引
            .distanceType(COSINE_DISTANCE)
            .build()
    );
}

// 存储文档
vectorStore.add(List.of(
    new Document("Spring Boot 是 Java 微服务框架", Map.of("category", "backend")),
    new Document("React 是前端 JavaScript 框架", Map.of("category", "frontend"))
));

// 语义检索
List<Document> results = vectorStore.similaritySearch(
    SearchRequest.query("Java 后端开发")
        .withTopK(5)
        .withSimilarityThreshold(0.7)
);

// ❌ 错误：逐条调用（100 次 HTTP 请求）
for (String text : texts) {
    float[] vec = model.embed(text);  // 每次都是一次 API 调用
}

// ✅ 正确：批量调用（1 次 HTTP 请求）
List<List<float>> vectors = model.embedAll(texts);  // 一次搞定

@Service
public class CachedEmbeddingService {

    private final EmbeddingModel delegate;
    private final Cache<String, float[]> cache;

    public CachedEmbeddingService(EmbeddingModel delegate) {
        this.delegate = delegate;
        this.cache = Caffeine.newBuilder()
            .maximumSize(10_000)
            .expireAfterAccess(Duration.ofHours(24))
            .build();
    }

    public float[] embed(String text) {
        return cache.get(text, key -> delegate.embed(key));
    }
}

// 使用 256 维（而非默认的 3072 维）
float[] fullVector = model.embed("text");  // 3072 维
float[] reduced = Arrays.copyOf(fullVector, 256);  // 256 维

// OpenAI API 原生支持（更推荐，模型内部做了优化）
EmbeddingResponse response = embeddingModel.call(
    new EmbeddingRequest(List.of("text"), 
        Map.of("dimensions", "256"))
);

// DJL 加载本地 Embedding 模型
Criteria<String, float[]> criteria = Criteria.builder()
    .setTypes(String.class, float[].class)
    .optModelUrls("djl://ai.djl.huggingface.onnxruntime/BAAI/bge-small-zh-v1.5")
    .optEngine("OnnxRuntime")
    .build();

ZooModel<String, float[]> model = criteria.loadModel();
Predictor<String, float[]> predictor = model.newPredictor();
float[] embedding = predictor.predict("Spring Boot 微服务框架");

// ❌ 入库用 OpenAI，查询用 BGE——维度不同，相似度无意义
float[] stored = openaiModel.embed("文档内容");    // 1536 维
float[] query = bgeModel.embed("用户查询");        // 1024 维
// cosineSimilarity(stored, query) → 胡乱计算，结果无意义

// ✅ 始终使用同一个模型
float[] stored = embeddingModel.embed("文档内容");  // 1536 维
float[] query = embeddingModel.embed("用户查询");   // 1536 维

public static float[] normalize(float[] vector) {
    double norm = 0;
    for (float v : vector) norm += v * v;
    norm = Math.sqrt(norm);
    float[] result = new float[vector.length];
    for (int i = 0; i < vector.length; i++) {
        result[i] = (float) (vector[i] / norm);
    }
    return result;
}

@Service
public class ResilientEmbeddingService {

    private final EmbeddingModel delegate;
    private final RetryTemplate retryTemplate;

    public ResilientEmbeddingService(EmbeddingModel delegate) {
        this.delegate = delegate;
        this.retryTemplate = RetryTemplate.builder()
            .maxAttempts(3)
            .exponentialBackoff(1000, 2.0, 10000)
            .retryOn(EmbeddingException.class)
            .build();
    }

    public float[] embed(String text) {
        return retryTemplate.execute(ctx -> delegate.embed(text));
    }
}