package com.sonic.cow.service.impl;

import com.alibaba.fastjson.JSON;
import com.baomidou.mybatisplus.core.toolkit.CollectionUtils;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.sonic.cow.domain.bo.TextGenImagePromptParse;
import com.sonic.cow.domain.bo.TextPromptParse;
import com.sonic.cow.domain.input.AiGenImageInput;
import com.sonic.cow.enums.PromptTypeEnum;
import com.sonic.cow.service.AiGenImageService;
import com.sonic.cow.service.CommonSendMqService;
import com.sonic.cow.service.GenAiUserContentService;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.util.List;
import java.util.Map;

import static com.sonic.cow.utils.Constant.COMMON_IMAGE_PROMPT;

@Slf4j
@Service
public class AiGenImageServiceImpl implements AiGenImageService {
    @Autowired
    private CommonSendMqService commonSendMqService;
    @Autowired
    private GenAiUserContentService aiUserContentService;


    @Override
    public void aiGenImage(AiGenImageInput input, String batchNo, List<String> taskIdList) {
        log.info("AiGenImageServiceImpl input:{}", JSON.toJSONString(input));
        String imageStylePrompt = input.getImageStylePrompt();
        String content = input.getContent();
        //基图url
        String imageUrl = input.getImageUrl();
        //基图提示词
        String initPrompt = input.getInitPrompt();

        if (StringUtils.isNotBlank(imageUrl) && StringUtils.isNotBlank(initPrompt)) {
            //调用AI 合并新老形象描述
            try {
                Map<String, Object> paramMap = Maps.newHashMap();
                paramMap.put("initPrompt", initPrompt);
                paramMap.put("content", content);
                TextPromptParse textPromptParse = aiUserContentService.genContent(PromptTypeEnum.MERGE_NEW_OLD_IMAGE_DESC.name(), paramMap, TextPromptParse.class);
                if (textPromptParse != null) {
                    content = textPromptParse.getContent();
                }
            } catch (Exception e) {
                log.error("AiGenImageServiceImpl imageGenText error", e);
            }
        }
        //初始化6组prompt
        List<String> promptList = Lists.newArrayList();
        try {
            //调用AI 获取6组不同的prompt  形象风格prompt+形象描述内容+通用图像提示词 组成新的prompt
            String sex = input.getSex() == null ? "" : input.getSex();
            Integer age = input.getAge() == null ? 20 : input.getAge();
            //拼接年龄性别
            String sexAgeStr = sex + "," + age + ",";
            Map<String, Object> genPromptMap = Maps.newHashMap();
            genPromptMap.put("content", content);
            genPromptMap.put("imageStylePrompt", imageStylePrompt);
            genPromptMap.put("commonImagePrompt", COMMON_IMAGE_PROMPT);
            genPromptMap.put("sex", sex);
            genPromptMap.put("age", age);
            genPromptMap.put("introduction", input.getIntroduction() == null ? "" : input.getIntroduction());
            TextGenImagePromptParse textGenImagePromptParse = aiUserContentService.genContent(PromptTypeEnum.TEXT_TO_IMAGE_PROMPT.name(), genPromptMap, TextGenImagePromptParse.class);
            //拼接年龄性别
            promptList.add(sexAgeStr + textGenImagePromptParse.getPrompt1());
            promptList.add(sexAgeStr + textGenImagePromptParse.getPrompt2());
            promptList.add(sexAgeStr + textGenImagePromptParse.getPrompt3());
            promptList.add(sexAgeStr + textGenImagePromptParse.getPrompt4());
            promptList.add(sexAgeStr + textGenImagePromptParse.getPrompt5());
            promptList.add(sexAgeStr + textGenImagePromptParse.getPrompt6());
        } catch (Exception e) {
            log.error("AiGenImageServiceImpl textGenImagePromptParse error", e);
        }
        //创建6组任务，发送MQ处理（根据6组不同的prompt生成6张图片，上传到S3,并鉴黄）
        log.info("AiGenImageServiceImpl promptList:{}", JSON.toJSONString(promptList));
        //没有，直接返回
        if (CollectionUtils.isEmpty(promptList)) {
            return;
        }
        for (int i = 0; i < promptList.size(); i++) {
            String prompt = promptList.get(i);
            //任务id
            String taskId = taskIdList.get(i);
            //发送MQ处理
            if (StringUtils.isEmpty(imageUrl)) {
                //文生图
                commonSendMqService.aiTextGenImageSendMq(batchNo, taskId, prompt, imageUrl);
            } else {
                //图生图
                commonSendMqService.aiImageGenImageSendMq(batchNo, taskId, prompt, imageUrl);
            }
        }
    }
}
