ai-paas-agent/src/main/java/edu/nju/software/aipaasagent/service/ChatCompletionService.java

package edu.nju.software.aipaasagent.service;

import edu.nju.software.aipaasagent.agent.manager.AgentFactory;
import edu.nju.software.aipaasagent.agent.core.base.BaseAgent;
import edu.nju.software.aipaasagent.agent.config.AgentConfigRegistry;
import edu.nju.software.aipaasagent.agent.config.AgentConfiguration;
import edu.nju.software.aipaasagent.dto.ChatCompletionRequest;
import edu.nju.software.aipaasagent.dto.ChatCompletionResponse;
import edu.nju.software.aipaasagent.dto.StreamEvent;
import lombok.extern.slf4j.Slf4j;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

import java.time.Instant;
import java.util.List;
import java.util.UUID;

/**
 * Chat Completion 服务
 * 处理 OpenAI 兼容的聊天完成请求
 * 支持流式和非流式输出
 */
@Slf4j
@Service
public class ChatCompletionService {

    @Autowired
    private AgentFactory agentFactory;

    @Autowired
    private AgentConfigRegistry agentConfigRegistry;

    @Autowired
    private MultimodalFileService multimodalFileService;

    /**
     * 非流式聊天完成
     *
     * @param request 聊天完成请求
     * @return 聊天完成响应
     */
    public ChatCompletionResponse chatCompletion(ChatCompletionRequest request) {
        String threadId = getOrCreateThreadId(request.getThreadId());
        String agentId = getOrCreateAgentId(request.getAgentId());
        
        log.info("非流式聊天完成 - threadId: {}, agentId: {}", threadId, agentId);

        // 获取或创建 Agent
        BaseAgent agent = getOrCreateAgent();
        
        // 转换消息格式并处理多模态文件
        String userMessage = extractLastUserMessage(request.getMessages());
        
        // 处理多模态文件（如果提供了 file_ids）
        if (request.getFileIds() != null && !request.getFileIds().isEmpty()) {
            userMessage = processMultimodalContent(userMessage, request.getFileIds());
        }
        
        String content = agent.doChat(userMessage, threadId, Boolean.TRUE.equals(request.getIsThink()));

        // 获取模型名称：接口参数 > ConfigMap 配置
        String model = resolveModelName(request.getModel(), agentId);

        // 构建响应
        return buildChatCompletionResponse(content, threadId, agentId, model);
    }

    /**
     * 流式聊天完成
     *
     * @param request 聊天完成请求
     * @return SSE 流
     */
    public Flux<ChatCompletionResponse> streamChatCompletion(ChatCompletionRequest request) {
        String threadId = getOrCreateThreadId(request.getThreadId());
        String agentId = getOrCreateAgentId(request.getAgentId());
        
        log.info("流式聊天完成 - threadId: {}, agentId: {}", threadId, agentId);

        // 获取或创建 Agent
        BaseAgent agent = getOrCreateAgent();
        
        // 转换消息格式并处理多模态文件
        String userMessage = extractLastUserMessage(request.getMessages());
        
        // 处理多模态文件（如果提供了 file_ids）
        if (request.getFileIds() != null && !request.getFileIds().isEmpty()) {
            userMessage = processMultimodalContent(userMessage, request.getFileIds());
        }

        // 生成响应 ID
        String responseId = "chatcmpl-" + UUID.randomUUID().toString().replace("-", "");
        
        // 获取模型名称：接口参数 > ConfigMap 配置
        String model = resolveModelName(request.getModel(), agentId);

        Flux<String> streamContent = agent.doStreamChat(userMessage, threadId, Boolean.TRUE.equals(request.getIsThink()));
        
        return streamContent
            .index()
            .map(tuple -> {
                long index = tuple.getT1();
                String content = tuple.getT2();
                
                return ChatCompletionResponse.builder()
                    .id(responseId)
                    .object("chat.completion.chunk")
                    .created(Instant.now().getEpochSecond())
                    .model(model)
                    .choices(List.of(
                        ChatCompletionResponse.Choice.builder()
                            .index((int) index)
                            .delta(ChatCompletionResponse.Delta.builder()
                                .content(content)
                                .build())
                            .finishReason(null)
                            .build()
                    ))
                    .threadId(threadId)
                    .agentId(agentId)
                    .build();
            })
            .concatWith(Flux.just(
                // 发送结束标记
                ChatCompletionResponse.builder()
                    .id(responseId)
                    .object("chat.completion.chunk")
                    .created(Instant.now().getEpochSecond())
                    .model(model)
                    .choices(List.of(
                        ChatCompletionResponse.Choice.builder()
                            .index(0)
                            .delta(new ChatCompletionResponse.Delta())
                            .finishReason("stop")
                            .build()
                    ))
                    .threadId(threadId)
                    .agentId(agentId)
                    .build()
            ));
    }
    
    /**
     * SSE 流式聊天完成（返回 StreamEvent 格式）
     * 用于 stream=true + isThink=true 的情况
     * 
     * @param request 聊天完成请求
     * @return StreamEvent 流
     */
    public Flux<StreamEvent> streamChatCompletionSSE(ChatCompletionRequest request) {
        String threadId = getOrCreateThreadId(request.getThreadId());
        String agentId = getOrCreateAgentId(request.getAgentId());
        
        log.info("SSE流式聊天完成 - threadId: {}, agentId: {}", threadId, agentId);

        // 获取或创建 Agent
        BaseAgent agent = getOrCreateAgent();
        
        // 转换消息格式并处理多模态文件
        String userMessage = extractLastUserMessage(request.getMessages());
        
        // 处理多模态文件（如果提供了 file_ids）
        if (request.getFileIds() != null && !request.getFileIds().isEmpty()) {
            userMessage = processMultimodalContent(userMessage, request.getFileIds());
        }

        // 调用 SSE 流式方法
        return agent.doStreamChatSSE(userMessage, threadId, Boolean.TRUE.equals(request.getIsThink()));
    }

    /**
     * 获取或创建 Agent
     */
    private BaseAgent getOrCreateAgent() {
        // 先尝试获取已有实例
        BaseAgent agent = agentFactory.getCurrentAgent();
        if (agent == null) {
            // 未创建则创建新实例
            agent = agentFactory.createAgent();
        }
        return agent;
    }

    /**
     * 获取或创建 Thread ID
     */
    private String getOrCreateThreadId(String threadId) {
        if (threadId == null || threadId.isEmpty()) {
            return "thread-" + UUID.randomUUID().toString().replace("-", "");
        }
        return threadId;
    }

    /**
     * 获取或创建 Agent ID
     */
    private String getOrCreateAgentId(String agentId) {
        if (agentId == null || agentId.isEmpty()) {
            return "seeCoderManus";
        }
        return agentId;
    }

    /**
     * 解析模型名称
     * 优先级：接口参数 > ConfigMap 配置
     *
     * @param requestModel 请求中的模型名称（可选）
     * @param agentId      Agent ID
     * @return 最终使用的模型名称
     */
    private String resolveModelName(String requestModel, String agentId) {
        // 如果接口传了 model，优先使用
        if (requestModel != null && !requestModel.isEmpty()) {
            log.debug("使用接口传入的模型名称: {}", requestModel);
            return requestModel;
        }
        
        // 否则从 ConfigMap 获取
        AgentConfiguration config = agentConfigRegistry.getConfig(agentId);
        if (config != null && config.getModel() != null && config.getModel().getName() != null) {
            String configModel = config.getModel().getName();
            log.debug("使用 ConfigMap 配置的模型名称: {}", configModel);
            return configModel;
        }
        
        // 默认模型
        log.warn("未找到模型配置，使用默认模型: gemma3:4b");
        return "gemma3:4b";
    }

    /**
     * 提取最后一条用户消息
     */
    private String extractLastUserMessage(List<ChatCompletionRequest.ChatMessage> messages) {
        if (messages == null || messages.isEmpty()) {
            return "";
        }
        
        // 找到最后一条用户消息
        for (int i = messages.size() - 1; i >= 0; i--) {
            ChatCompletionRequest.ChatMessage message = messages.get(i);
            if ("user".equals(message.getRole())) {
                return message.getContent();
            }
        }
        
        // 如果没有用户消息，返回最后一条消息的内容
        return messages.get(messages.size() - 1).getContent();
    }

    /**
     * 处理多模态内容
     * 将文件内容整合到用户消息中
     *
     * @param userMessage 用户原始消息
     * @param fileIds     文件 ID 列表
     * @return 整合后的消息内容
     */
    private String processMultimodalContent(String userMessage, List<String> fileIds) {
        log.info("处理多模态内容 - fileIds: {}", fileIds);
        
        // 读取文件内容
        List<MultimodalFileService.FileContent> fileContents = multimodalFileService.readFiles(fileIds);
        
        if (fileContents.isEmpty()) {
            log.warn("未能读取任何文件内容 - fileIds: {}", fileIds);
            return userMessage;
        }
        
        StringBuilder multimodalContent = new StringBuilder();
        
        // 添加用户原始消息
        if (userMessage != null && !userMessage.isEmpty()) {
            multimodalContent.append("用户消息：\n").append(userMessage).append("\n\n");
        }
        
        // 添加文件内容
        multimodalContent.append("文件内容：\n");
        
        for (int i = 0; i < fileContents.size(); i++) {
            MultimodalFileService.FileContent fileContent = fileContents.get(i);
            multimodalContent.append("--- 文件 ").append(i + 1).append(" ---\n");
            multimodalContent.append("文件名：").append(fileContent.getFileName()).append("\n");
            multimodalContent.append("文件类型：").append(fileContent.getContentType()).append("\n");
            
            if (fileContent.isText()) {
                // 文本文件，直接添加内容
                multimodalContent.append("文件内容：\n").append(fileContent.getTextContent()).append("\n");
            } else if (fileContent.isImage()) {
                // 图片文件，添加 Base64 编码的 Data URI
                // 注意：某些模型支持直接处理图片，这里将图片作为 Base64 提供
                multimodalContent.append("图片内容（Base64）：\n");
                multimodalContent.append(fileContent.getDataUri()).append("\n");
            } else {
                // 其他类型文件，添加 Base64 内容
                multimodalContent.append("文件内容（Base64）：\n");
                multimodalContent.append(fileContent.getBase64Content()).append("\n");
            }
            
            multimodalContent.append("\n");
        }
        
        String result = multimodalContent.toString();
        log.info("多模态内容处理完成 - 原始消息长度: {}, 处理后长度: {}", 
                userMessage != null ? userMessage.length() : 0, result.length());
        
        return result;
    }

    /**
     * 构建聊天完成响应
     */
    private ChatCompletionResponse buildChatCompletionResponse(
            String content, 
            String threadId, 
            String agentId,
            String model) {
        
        String responseId = "chatcmpl-" + UUID.randomUUID().toString().replace("-", "");
        
        return ChatCompletionResponse.builder()
            .id(responseId)
            .object("chat.completion")
            .created(Instant.now().getEpochSecond())
            .model(model != null ? model : "default-model")
            .choices(List.of(
                ChatCompletionResponse.Choice.builder()
                    .index(0)
                    .message(ChatCompletionResponse.Message.builder()
                        .role("assistant")
                        .content(content)
                        .build())
                    .finishReason("stop")
                    .build()
            ))
            .usage(ChatCompletionResponse.Usage.builder()
                .promptTokens(0)  // TODO: 实现 token 计数
                .completionTokens(0)
                .totalTokens(0)
                .build())
            .threadId(threadId)
            .agentId(agentId)
            .build();
    }
}