package com.sonic.cow.event.inner.handler;

import com.baomidou.mybatisplus.core.toolkit.StringUtils;
import com.sonic.common.event.Event;
import com.sonic.common.event.EventConsumer;
import com.sonic.common.event.EventHandler;
import com.sonic.cow.client.Seedream4GenImageClient;
import com.sonic.cow.client.TextGenImageClient;
import com.sonic.cow.domain.input.GenImageCallbackInput;
import com.sonic.cow.event.inner.EventType;
import com.sonic.cow.event.inner.payload.AiGenImagePayload;
import com.sonic.cow.service.GenImageTaskService;
import com.sonic.shark.lib.client.S3Client;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.util.Base64;

/**
 * @author chenjun
 */
@Slf4j
@Component
public class AiTextGenImageHandler implements EventHandler, InitializingBean {

    @Autowired
    private EventConsumer eventConsumer;
    @Autowired
    private TextGenImageClient textGenImageClient;
    @Autowired
    private S3Client s3Client;
    @Autowired
    private GenImageTaskService genImageTaskService;
    @Autowired
    private Seedream4GenImageClient seedream4GenImageClient;

    @Override
    public void onEvent(Event event) {
        AiGenImagePayload payload = event.normalizedData(AiGenImagePayload.class);
        log.info("AiTextGenImageHandler payload:{}", payload);
        //文生图
        aiGenImage(payload);
    }

    @Override
    public void afterPropertiesSet() {
        eventConsumer.registerHandler(EventType.AI_TEXT_GEN_IMAGE.getEventCode(), this);
    }

    private void aiGenImage(AiGenImagePayload payload) {
        String prompt = payload.getPrompt();
        String taskId = payload.getTaskId();
        if (StringUtils.isEmpty(taskId) || StringUtils.isEmpty(prompt)) {
            return;
        }
        //生成的图片
        String base64Json = null;
        String asyncTaskId = "0";
        //完成状态
        String status = "COMPLETED";
        String failReason = null;
        try {
            //调用模型-文生成图片
//            base64Json = textGenImageClient.textGenImage(prompt);
            base64Json = seedream4GenImageClient.genImage("", prompt);
        } catch (Exception e) {
            log.error("AiTextGenImageHandler error:", e);
            status = "FAILED";
            failReason = e.getMessage();
        }
        //图片上传，更新taskId状态发送MQ到业务系统处理
        String genImageUrl = null;
        if (StringUtils.isNotEmpty(base64Json)) {
            //上传到S3
            byte[] imageBytes = Base64.getDecoder().decode(base64Json);
            genImageUrl = s3Client.uploadAwsS3(imageBytes, "ROLE","jpg");
            log.info("AiTextGenImageHandler genImageUrl:{}", genImageUrl);
        }
        //更新生成图片任务状态结果
        GenImageCallbackInput genImageCallbackInput = new GenImageCallbackInput();
        genImageCallbackInput.setThirdTaskId(taskId);
        genImageCallbackInput.setPrompt(prompt);
        genImageCallbackInput.setImageUrl(genImageUrl);
        genImageCallbackInput.setStatus(status);
        genImageCallbackInput.setFailReason(failReason);
        genImageCallbackInput.setAsyncTaskId(asyncTaskId);
        genImageCallbackInput.setType(1);
        genImageTaskService.genImageCallback(genImageCallbackInput);
    }
}
