package cn.iocoder.yudao.module.ai.api.chat;

import cn.iocoder.yudao.framework.common.enums.CommonStatusEnum;
import cn.iocoder.yudao.module.ai.dal.dataobject.model.AiModelDO;
import cn.iocoder.yudao.module.ai.enums.model.AiModelTypeEnum;
import cn.iocoder.yudao.module.ai.service.model.AiModelService;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.content.Media;
import org.springframework.core.io.ByteArrayResource;
import org.springframework.stereotype.Service;
import org.springframework.util.MimeTypeUtils;

import java.util.Base64;
import java.util.List;

@Service
@Slf4j
public class AiChatApiImpl implements AiChatApi {

    @Resource
    private AiModelService modelService;

    @Override
    public String chat(String systemPrompt, String userMessage) {
        AiModelDO defaultModel = modelService.getRequiredDefaultModel(AiModelTypeEnum.CHAT.getType());
        ChatModel chatModel = modelService.getChatModel(defaultModel.getId());
        Prompt prompt = new Prompt(List.of(
                new SystemMessage(systemPrompt),
                new UserMessage(userMessage)));
        String result = chatModel.call(prompt).getResult().getOutput().getText();
        log.info("AI Chat 调用完成，模型: {}, 回复长度: {}", defaultModel.getModel(),
                result != null ? result.length() : 0);
        return result;
    }

    @Override
    public String chatWithImage(String systemPrompt, String userMessage, String imageBase64, String mimeType) {
        // 优先选择视觉模型（模型名含 vl），否则回退到默认对话模型
        List<AiModelDO> chatModels = modelService.getModelListByStatusAndType(
                CommonStatusEnum.ENABLE.getStatus(), AiModelTypeEnum.CHAT.getType(), null);
        AiModelDO visionModel = chatModels.stream()
                .filter(m -> m.getModel() != null && m.getModel().contains("vl"))
                .findFirst()
                .orElseGet(() -> modelService.getRequiredDefaultModel(AiModelTypeEnum.CHAT.getType()));

        ChatModel chatModel = modelService.getChatModel(visionModel.getId());
        byte[] imageBytes = Base64.getDecoder().decode(imageBase64);
        Media media = new Media(MimeTypeUtils.parseMimeType(mimeType),
                new ByteArrayResource(imageBytes));
        UserMessage imageMessage = UserMessage.builder()
                .text(userMessage)
                .media(List.of(media))
                .build();
        Prompt prompt = new Prompt(List.of(
                new SystemMessage(systemPrompt),
                imageMessage));
        String result = chatModel.call(prompt).getResult().getOutput().getText();
        log.info("AI Chat(多模态) 调用完成，模型: {}, 回复长度: {}", visionModel.getModel(),
                result != null ? result.length() : 0);
        return result;
    }

}