package edu.whut.config;

import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.ai.ollama.OllamaEmbeddingModel;
import org.springframework.ai.ollama.api.OllamaApi;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.ai.vectorstore.pgvector.PgVectorStore;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Primary;
import org.springframework.jdbc.core.JdbcTemplate;

/**
 * 配置 Ollama 客户端、Embedding 模型、向量存储和文本拆分器等 Bean
 * 注意（1.0.0-M6）：
 * - 不再使用 OllamaChatClient，改为 OllamaChatModel
 * - Embedding 使用 EmbeddingModel 接口 + OllamaEmbeddingModel
 * - PgVectorStore 与 SimpleVectorStore 使用 builder 方式创建
 */
@Configuration
public class OllamaConfig {

    /**
     * 创建 Ollama API 客户端，负责与 Ollama 服务的基础通信
     */
    @Bean
    public OllamaApi ollamaApi(@Value("${spring.ai.ollama.base-url}") String baseUrl) {
        return new OllamaApi(baseUrl);
    }

    /**
     * 对话模型；控制器里通过 OllamaOptions 指定具体 chat 模型（如 deepseek-r1:1.5b）
     */
    @Bean
    public OllamaChatModel ollamaChatModel(OllamaApi ollamaApi) {
        // 如需设置默认对话模型，可在这里 .defaultOptions(OllamaOptions.builder().model("xxx").build())
        return OllamaChatModel.builder()
                .ollamaApi(ollamaApi)
                .build();
    }

    /**
     * 文本拆分器：基于 Token 切分
     */
    @Bean
    public TokenTextSplitter tokenTextSplitter() {
        return new TokenTextSplitter();
    }

    /**
     * EmbeddingModel（Ollama）：使用配置项 spring.ai.rag.embed 指定 embedding 模型，如 nomic-embed-text
     */
    @Bean
    @Primary
    public EmbeddingModel embeddingModel(
            @Value("${spring.ai.rag.embed}") String embedModel,
            OllamaApi ollamaApi) {
        return OllamaEmbeddingModel.builder()
                .ollamaApi(ollamaApi)
                .defaultOptions(OllamaOptions.builder().model(embedModel).build())
                .build();
    }

    /**
     * 简单内存向量存储（in-memory）
     */
    @Bean
    public SimpleVectorStore vectorStore(EmbeddingModel embeddingModel) {
        return SimpleVectorStore.builder(embeddingModel).build();
    }

    /**
     * 基于 PostgreSQL pgvector 的持久化向量存储
     * 如需自定义表名可追加 .vectorTableName("your_table")
     */
    @Bean
    public PgVectorStore pgVectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel) {
        return PgVectorStore.builder(jdbcTemplate, embeddingModel).build();
    }
}