package com.xly.config;

import com.fasterxml.jackson.databind.DeserializationFeature;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.SerializationFeature;
import com.fasterxml.jackson.datatype.jsr310.JavaTimeModule;
import com.xly.agent.*;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.ollama.OllamaChatModel;
import dev.langchain4j.model.ollama.OllamaStreamingChatModel;
import dev.langchain4j.service.AiServices;
import org.springframework.beans.factory.annotation.Qualifier;
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 java.time.Duration;

/**
 * 模型与 Agent 统一装配（langchain4j 1.14.0 新 API：ChatModel / StreamingChatModel）
 * 约定：
 *  - 路由/SQL/参数提取类任务 → 低温（确定性）
 *  - 闲聊类任务             → 高温（多样性）
 */
@Configuration
public class ModelConfig {

    @Value("${langchain4j.ollama.base-url}")
    private String baseUrl;

    @Value("${langchain4j.ollama.chat-model-name}")
    private String chatModelName;

    @Value("${langchain4j.ollama.sql-model-name}")
    private String sqlModelName;

    // ========================================================
    //  一、私有 builder 工厂：消除重复配置
    // ========================================================

    /** 非流式 Ollama 模型的公共构造 */
    private OllamaChatModel.OllamaChatModelBuilder chatBuilder(String modelName,
                                                               double temperature,
                                                               double topP,
                                                               long timeoutSeconds,
                                                               int maxRetries) {
        return OllamaChatModel.builder()
                .baseUrl(baseUrl)
                .modelName(modelName)
                .temperature(temperature)
                .topP(topP)
                .timeout(Duration.ofSeconds(timeoutSeconds))
                .maxRetries(maxRetries);
    }

    /** 流式 Ollama 模型的公共构造 */
    private OllamaStreamingChatModel.OllamaStreamingChatModelBuilder streamingBuilder(String modelName,
                                                                                      double temperature,
                                                                                      double topP,
                                                                                      int numPredict,
                                                                                      long timeoutSeconds) {
        return OllamaStreamingChatModel.builder()
                .baseUrl(baseUrl)
                .modelName(modelName)
                .temperature(temperature)
                .topP(topP)
                .numPredict(numPredict)
                .timeout(Duration.ofSeconds(timeoutSeconds));
    }

    // ========================================================
    //  二、模型 Bean
    // ========================================================

    /** 主对话模型：路由/场景/方法选择等确定性任务共用 */
    @Bean
    @Primary
    public OllamaChatModel chatLanguageModel() {
        return chatBuilder(chatModelName, 0.1, 0.95, 120, 2)
                .logRequests(true)
                .logResponses(true)
                .build();
    }

    /** 自由闲聊（非流式） */
    @Bean("chatiModel")
    public OllamaChatModel chatiModel() {
        return chatBuilder(chatModelName, 0.7, 0.9, 60, 2).build();
    }

    /** 自由闲聊（流式） */
    @Bean("chatiStreamingModel")
    public OllamaStreamingChatModel chatiStreamingModel() {
        return streamingBuilder(chatModelName, 0.7, 0.9, 512, 60).build();
    }

    /** SQL 专用（非流式，零温 + 大输出） */
    @Bean("sqlChatModel")
    public OllamaChatModel sqlChatModel() {
        return chatBuilder(sqlModelName, 0.0, 0.95, 120, 3)
                .numPredict(4096)
                .build();
    }

    /** 主流式对话 */
    @Bean("streamingChatModel")
    @Primary
    public OllamaStreamingChatModel streamingChatModel() {
        return streamingBuilder(chatModelName, 0.3, 0.9, 1024, 60).build();
    }

    /** SQL 流式 */
    @Bean("streamingSqlModel")
    public OllamaStreamingChatModel streamingSqlModel() {
        return streamingBuilder(sqlModelName, 0.2, 0.95, 2048, 120).build();
    }

    // ========================================================
    //  三、JSON
    // ========================================================

    @Bean
    @Primary
    public ObjectMapper objectMapper() {
        ObjectMapper mapper = new ObjectMapper();
        mapper.registerModule(new JavaTimeModule());
        mapper.disable(SerializationFeature.WRITE_DATES_AS_TIMESTAMPS);
        mapper.configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false);
        return mapper;
    }

    // ========================================================
    //  四、Agent Bean
    // ========================================================

    /** 动态 SQL Agent */
    @Bean
    public DynamicTableNl2SqlAiAgent dynamicTableNl2SqlAiAgent(
            @Qualifier("sqlChatModel") OllamaChatModel sqlModel) {
        return AiServices.builder(DynamicTableNl2SqlAiAgent.class)
                .chatModel(sqlModel)
                .chatMemoryProvider(memoryId -> MessageWindowChatMemory.withMaxMessages(10))
                .build();
    }

    /**
     * 闲聊 Agent：会话隔离由 @MemoryId + OperableChatMemoryProvider 完成，
     * 全局共用一个实例即可
     */
    @Bean
    public ChatiAgent chatiAgent(
            @Qualifier("chatiModel") OllamaChatModel chatiModel,
            @Qualifier("chatiStreamingModel") OllamaStreamingChatModel chatiStreamingModel,
            OperableChatMemoryProvider operableChatMemoryProvider) {
        return AiServices.builder(ChatiAgent.class)
                .chatModel(chatiModel)
                .streamingChatModel(chatiStreamingModel)
                .chatMemoryProvider(operableChatMemoryProvider)
                .maxSequentialToolsInvocations(1)
                .build();
    }

    /** 场景选择 Agent（一级路由） */
    @Bean
    public SceneSelectorAiAgent sceneSelectorAiAgent(
            @Qualifier("chatLanguageModel") OllamaChatModel chatLanguageModel) {
        return AiServices.builder(SceneSelectorAiAgent.class)
                .chatModel(chatLanguageModel)
                .chatMemoryProvider(memoryId -> MessageWindowChatMemory.withMaxMessages(10))
                .maxSequentialToolsInvocations(1)
                .build();
    }

    /**
     * 方法选择 Agent（二级路由）：从场景内方法树中路由出 sMethodNo。
     * 分类任务，用极短记忆窗口避免被上文惯性带偏。
     */
    @Bean
    public SecMethodAiAgent secMethodAiAgent(
            @Qualifier("chatLanguageModel") OllamaChatModel chatLanguageModel) {
        return AiServices.builder(SecMethodAiAgent.class)
                .chatModel(chatLanguageModel)
                .chatMemoryProvider(memoryId -> MessageWindowChatMemory.withMaxMessages(2))
                .maxSequentialToolsInvocations(1)
                .build();
    }

    /**
     * 方法选择 Agent（二级路由）：从数据中获取到对应的 添加的部件名称，删除的部件名称
     * 分类任务，用极短记忆窗口避免被上文惯性带偏。
     */
    @Bean
    public DynamicPartAiAgent dynamicPartAiAgent(
            @Qualifier("chatLanguageModel") OllamaChatModel chatLanguageModel) {
        return AiServices.builder(DynamicPartAiAgent.class)
                .chatModel(chatLanguageModel)
                .chatMemoryProvider(memoryId -> MessageWindowChatMemory.withMaxMessages(2))
                .maxSequentialToolsInvocations(1)
                .build();
    }

}