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.AiModelCapabilityEnum;
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;

import static cn.iocoder.yudao.framework.common.exception.util.ServiceExceptionUtil.exception;
import static cn.iocoder.yudao.module.ai.enums.ErrorCodeConstants.MODEL_VISION_NOT_ENABLED;

@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) {
        // 图片识别必须走视觉能力的模型：文本模型不接受图片输入，回退过去只会报错
        AiModelDO visionModel = modelService.getModelListByStatusAndType(
                        CommonStatusEnum.ENABLE.getStatus(), AiModelTypeEnum.CHAT.getType(), null)
                .stream()
                .filter(m -> AiModelCapabilityEnum.VISION.getCapability().equals(m.getCapability()))
                .findFirst()
                .orElseThrow(() -> exception(MODEL_VISION_NOT_ENABLED));

        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;
    }

}