package com.xly.config;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.SerializationFeature;
import com.fasterxml.jackson.datatype.jsr310.JavaTimeModule;
import com.xly.agent.DynamicTableNl2SqlAiAgent;
import com.xly.agent.SceneSelectorAiAgent;
import dev.langchain4j.model.ollama.OllamaChatModel;
import dev.langchain4j.model.ollama.OllamaStreamingChatModel;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
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;

@Configuration
public class ModelConfig {

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

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

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

    // ====================== 主对话模型 ======================
    @Bean
    @Primary
    public OllamaChatModel chatLanguageModel() {
        return OllamaChatModel.builder()
                .baseUrl(chatModelUrl)
                .modelName(chatModelName)
                .temperature(0.1)
                .topP(0.95)
                .timeout(Duration.ofSeconds(120))
                .maxRetries(2)
                .logRequests(true)
                .logResponses(true)
                .build();
    }

    // ====================== 自由闲聊模型 ======================
    @Bean("chatiModel")
    public OllamaChatModel chatiModel() {
        return OllamaChatModel.builder()
                .baseUrl(chatModelUrl)
                .modelName(chatModelName)
                .temperature(0.7)
                .topP(0.9)
                .timeout(Duration.ofSeconds(60))
                .maxRetries(2)
                .build();
    }

    // ====================== SQL 专用模型 ======================
    @Bean("sqlChatModel")
    public OllamaChatModel sqlChatModel() {
        return OllamaChatModel.builder()
                .baseUrl(chatModelUrl)
                .modelName(sqlModelName)
                .temperature(0.0)
                .topP(0.95)
                .numPredict(4096)
                .timeout(Duration.ofSeconds(120))
                .maxRetries(3)
                .build();
    }

    // ====================== 流式对话模型 ======================
    @Bean("streamingChatModel")
    @Primary
    public OllamaStreamingChatModel streamingChatModel() {
        return OllamaStreamingChatModel.builder()
                .baseUrl(chatModelUrl)
                .modelName(chatModelName)
                .temperature(0.3)
                .topP(0.9)
                .numPredict(1024)
                .timeout(Duration.ofSeconds(60))
                .build();
    }

    // ====================== 流式 SQL 模型 ======================
    @Bean("streamingSqlModel")
    public OllamaStreamingChatModel streamingSqlModel() {
        return OllamaStreamingChatModel.builder()
                .baseUrl(chatModelUrl)
                .modelName(sqlModelName)
                .temperature(0.2)
                .topP(0.95)
                .numPredict(2048)
                .timeout(Duration.ofSeconds(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(com.fasterxml.jackson.databind.DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false);
        return mapper;
    }

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

    // ====================== 场景选择 Agent ======================
    @Bean
    public SceneSelectorAiAgent sceneSelectorAiAgent(
            @Qualifier("chatLanguageModel") OllamaChatModel chatLanguageModel) {
        return AiServices.builder(SceneSelectorAiAgent.class)
                .chatModel(chatLanguageModel)
                .chatMemoryProvider(memoryId -> MessageWindowChatMemory.withMaxMessages(10))
                .build();
    }
}