package edu.whut.config;

import org.springframework.ai.ollama.OllamaChatClient;
import org.springframework.ai.ollama.OllamaEmbeddingClient;
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.PgVectorStore;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.jdbc.core.JdbcTemplate;

/**
 * 配置 Ollama 和 OpenAI 客户端，以及向量存储和文本拆分器等 Bean
 */
@Configuration
public class OllamaConfig {

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


    /**
     * 基于 OllamaApi 实例创建对话客户端，用于与 Ollama 模型进行对话交互
     */
    @Bean
    public OllamaChatClient ollamaChatClient(OllamaApi ollamaApi) {
        return new OllamaChatClient(ollamaApi);
    }

    /**
     * 文本拆分器，根据 Token 划分长文本，便于分段处理或嵌入计算
     */
    @Bean
    public TokenTextSplitter tokenTextSplitter() {
        return new TokenTextSplitter();
    }

    /**
     * 创建一个简单的向量存储 (in-memory)，可根据配置选择 Ollama 或 OpenAI 的嵌入模型
     */
    @Bean
    public SimpleVectorStore vectorStore(
            @Value("${spring.ai.rag.embed}") String model,
            OllamaApi ollamaApi) {

        // 固定使用 Ollama 的嵌入客户端（不再走 OpenAI 分支）
        OllamaEmbeddingClient embeddingClient = new OllamaEmbeddingClient(ollamaApi);
        embeddingClient.withDefaultOptions(
                OllamaOptions.create().withModel(model)
        );
        return new SimpleVectorStore(embeddingClient);
    }

    /**
     * 创建基于 PostgreSQL pgvector 扩展的持久化向量存储，可根据配置选择 Ollama 或 OpenAI 模型
     */
    @Bean
    public PgVectorStore pgVectorStore(
            @Value("${spring.ai.rag.embed}") String model,
            OllamaApi ollamaApi,
            JdbcTemplate jdbcTemplate) {

        // 固定使用 Ollama 的嵌入客户端（不再走 OpenAI 分支）
        OllamaEmbeddingClient embeddingClient = new OllamaEmbeddingClient(ollamaApi);
        embeddingClient.withDefaultOptions(
                OllamaOptions.create().withModel(model)
        );
        return new PgVectorStore(jdbcTemplate, embeddingClient);
    }

}