package com.sonic.cow.service.impl;

import com.baomidou.mybatisplus.core.toolkit.CollectionUtils;
import com.baomidou.mybatisplus.core.toolkit.StringUtils;
import com.baomidou.mybatisplus.core.toolkit.Wrappers;
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import com.google.common.collect.Lists;
import com.sonic.cow.dao.GenImageTaskRecordDao;
import com.sonic.cow.domain.entity.GenImageTask;
import com.sonic.cow.domain.entity.GenImageTaskRecord;
import com.sonic.cow.domain.output.GenImageListOutput;
import com.sonic.cow.enums.ToastResultCode;
import com.sonic.cow.service.GenImageTaskRecordService;
import com.sonic.cow.utils.BeanConvert;
import com.sonic.cow.utils.MD5Util;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;

import java.time.LocalDateTime;
import java.util.List;
import java.util.stream.Collectors;

/**
 * 生成图片任务子项表
 *
 * @author code
 */
@Slf4j
@Service
public class GenImageTaskRecordServiceImpl extends ServiceImpl<GenImageTaskRecordDao, GenImageTaskRecord> implements GenImageTaskRecordService {

    @Override
    public void saveRecord(GenImageTask genImageTask, List<String> taskIdList) {
        List<GenImageTaskRecord> list = Lists.newArrayList();
        for (String taskId : taskIdList) {
            GenImageTaskRecord genImageTaskRecord = new GenImageTaskRecord();
            genImageTaskRecord.setTaskId(genImageTask.getId());
            genImageTaskRecord.setBatchNo(genImageTask.getBatchNo());
            genImageTaskRecord.setUserId(genImageTask.getUserId());
            genImageTaskRecord.setThirdTaskId(taskId);
            genImageTaskRecord.setPrompt(null);
            //先给一个默认值1，好走索引
            genImageTaskRecord.setImageUrlMd5("1");
            genImageTaskRecord.setStatus(GenImageTaskRecord.Status.PENDING);
            genImageTaskRecord.setCompleted(false);
            genImageTaskRecord.setCreateTime(LocalDateTime.now());
            genImageTaskRecord.setEditTime(LocalDateTime.now());
            list.add(genImageTaskRecord);
        }
        if (CollectionUtils.isNotEmpty(list)) {
            saveBatch(list);
        }
    }

    @Override
    public GenImageTaskRecord getByThirdTaskId(String thirdTaskId) {
        GenImageTaskRecord genImageTaskRecord = getOne(Wrappers.<GenImageTaskRecord>lambdaQuery()
                .eq(GenImageTaskRecord::getThirdTaskId, thirdTaskId));
        return genImageTaskRecord;
    }

    @Override
    public void genImageCompleted(Long id, String imageUrl, String status, String prompt, String asyncTaskId, String failReason, Integer type) {
        String imageUrlMd5 = null;
        if (StringUtils.isNotEmpty(imageUrl)) {
            imageUrlMd5 = MD5Util.digest(imageUrl);
        }
        update(Wrappers.<GenImageTaskRecord>lambdaUpdate()
                .set(GenImageTaskRecord::getImageUrl, imageUrl)
                .set(GenImageTaskRecord::getImageUrlMd5, imageUrlMd5)
                .set(GenImageTaskRecord::getCompleted, "PENDING".equals(status) ? false : true)
                .set(GenImageTaskRecord::getStatus, status)
                .set(GenImageTaskRecord::getFailReason, failReason)
                .set(GenImageTaskRecord::getPrompt, prompt)
                .set(GenImageTaskRecord::getAsyncTaskId, asyncTaskId)
                .set(GenImageTaskRecord::getType, type)
                .eq(GenImageTaskRecord::getId, id));
    }

    @Override
    public Integer queryCompletedTaskCountByTaskId(Long taskId) {
        int count = count(Wrappers.<GenImageTaskRecord>lambdaQuery()
                .eq(GenImageTaskRecord::getTaskId, taskId)
                .eq(GenImageTaskRecord::getCompleted, true));
        return count;
    }

    @Override
    public Integer queryCompletedTaskCountByBatchNo(String batchNo) {
        int count = count(Wrappers.<GenImageTaskRecord>lambdaQuery()
                .eq(GenImageTaskRecord::getBatchNo, batchNo)
                .eq(GenImageTaskRecord::getCompleted, true));
        return count;
    }


    @Override
    public List<GenImageListOutput> queryByBatchNo(String batchNo) {
        List<GenImageTaskRecord> list = list(Wrappers.<GenImageTaskRecord>lambdaQuery()
                .select(GenImageTaskRecord::getImageUrl, GenImageTaskRecord::getStatus, GenImageTaskRecord::getPrompt)
                .eq(GenImageTaskRecord::getBatchNo, batchNo));
        List<GenImageListOutput> outputList = BeanConvert.copeList(list, GenImageListOutput.class);
        return outputList;
    }

    @Override
    public String getInitPromptByImageUrl(Long userId, String imageUrl) {
        if (StringUtils.isEmpty(imageUrl)) {
            return null;
        }
        //将url转换成MD5的字段进行查询
        String imageUrlMd5 = MD5Util.digest(imageUrl);
        GenImageTaskRecord genImageTaskRecord = getOne(Wrappers.<GenImageTaskRecord>lambdaQuery()
                .eq(GenImageTaskRecord::getUserId, userId)
                .eq(GenImageTaskRecord::getImageUrlMd5, imageUrlMd5)
                .last(" limit 1")
        );
        return genImageTaskRecord != null ? genImageTaskRecord.getPrompt() : null;
    }

    @Override
    public void checkImageIsAIGenerated(Long userId, List<String> imgUrlList) {
        ToastResultCode.IMAGE_NOT_ILLEGAL.check(CollectionUtils.isEmpty(imgUrlList));
        List<String> imgUrlMd5List = imgUrlList.stream().map(imgUrl -> MD5Util.digest(imgUrl)).collect(Collectors.toList());
        int count = count(Wrappers.<GenImageTaskRecord>lambdaQuery().eq(GenImageTaskRecord::getUserId, userId).in(GenImageTaskRecord::getImageUrlMd5, imgUrlMd5List));
        ToastResultCode.IMAGE_NOT_ILLEGAL.check(imgUrlList.size() != count);
    }
}
