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.common.utils.RedisLock;
import com.sonic.cow.client.ImageGenImageClient;
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.cow.utils.RedisKeyUtils;
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 AiImageGenImageHandler implements EventHandler, InitializingBean {

    @Autowired
    private EventConsumer eventConsumer;
    @Autowired
    private TextGenImageClient textGenImageClient;
    @Autowired
    private ImageGenImageClient imageGenImageClient;
    @Autowired
    private S3Client s3Client;
    @Autowired
    private GenImageTaskService genImageTaskService;
    @Autowired
    private RedisKeyUtils redisKeyUtils;
    @Autowired
    private RedisLock.RedisWrapper redisWrapper;
    @Autowired
    private Seedream4GenImageClient seedream4GenImageClient;

    @Override
    public void onEvent(Event event) {
        AiGenImagePayload payload = event.normalizedData(AiGenImagePayload.class);
        log.info("AiImageGenImageHandler payload:{}", payload);
//        String imageUrl = payload.getImageUrl();
//        if (StringUtils.isNotEmpty(imageUrl)) {
//            //图生图时才加锁,针对批次号加锁
//            String key = redisKeyUtils.aiGenImageLockKey(payload.getBatchNo());
//            RedisLock redisLock = new RedisLock(key, redisWrapper);
//            redisLock.tryAcquireRun(25 * 1000, 25 * 1000, () -> {
//                aiGenImage(payload);
//                return true;
//            });
//        }
        aiGenImage(payload);
    }

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

    private void aiGenImage(AiGenImagePayload payload) {
        String prompt = payload.getPrompt();
        String taskId = payload.getTaskId();
        //基图
        String imageUrl = payload.getImageUrl();
        if (StringUtils.isEmpty(taskId) || StringUtils.isEmpty(prompt)) {
            return;
        }
        //生成的图片
        String base64Json = null;
        String asyncTaskId = "0";
        //完成状态
        String status = "COMPLETED";
        String failReason = null;
        try {
            //调用模型-同步图生图
//            base64Json = imageGenImageClient.imageGenImage(imageUrl, prompt);
            base64Json = seedream4GenImageClient.genImage(imageUrl, prompt);
        } catch (Exception e) {
            log.error("AiImageGenImageHandler 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("AiImageGenImageHandler 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(2);
        genImageTaskService.genImageCallback(genImageCallbackInput);
    }
}
