package com.sonic.cow.service.impl;

import com.alibaba.fastjson.JSONObject;
import com.baomidou.mybatisplus.core.toolkit.Wrappers;
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import com.google.common.collect.Lists;
import com.sonic.common.auth.domains.Session;
import com.sonic.common.utils.LogUtils;
import com.sonic.cow.client.ContextChatClient;
import com.sonic.cow.dao.GenImageTaskDao;
import com.sonic.cow.domain.bo.AiUserCacheInfo;
import com.sonic.cow.domain.bo.LimitBo;
import com.sonic.cow.domain.entity.GenImageTask;
import com.sonic.cow.domain.entity.GenImageTaskRecord;
import com.sonic.cow.domain.input.AiGenImageInput;
import com.sonic.cow.domain.input.GenImageCallbackInput;
import com.sonic.cow.domain.input.GenImageInput;
import com.sonic.cow.domain.output.AiGenImageOutput;
import com.sonic.cow.domain.output.GenImageListOutput;
import com.sonic.cow.domain.output.GenImageOutput;
import com.sonic.cow.enums.BizResultCode;
import com.sonic.cow.enums.CreateImageType;
import com.sonic.cow.enums.SexEnums;
import com.sonic.cow.enums.ToastResultCode;
import com.sonic.cow.service.AiGenImageService;
import com.sonic.cow.service.GenImageTaskRecordService;
import com.sonic.cow.service.GenImageTaskService;
import com.sonic.cow.utils.*;
import com.sonic.frog.lib.client.AiChatInfoClient;
import com.sonic.frog.lib.client.AiClient;
import com.sonic.frog.lib.client.AiUserAlbumClient;
import com.sonic.frog.lib.enums.HeartbeatLevelEnum;
import com.sonic.frog.lib.output.AiChatInfoOutput;
import com.sonic.frog.lib.output.AiInfoApiOutput;
import com.sonic.lion.lib.client.PayClient;
import com.sonic.lion.lib.enums.BizType;
import com.sonic.lion.lib.input.BalanceCheckoutInput;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;

import java.time.LocalDateTime;
import java.util.List;
import java.util.UUID;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

/**
 * 生成图片任务
 *
 * @author code
 */
@Slf4j
@Service
public class GenImageTaskServiceImpl extends ServiceImpl<GenImageTaskDao, GenImageTask> implements GenImageTaskService {

    @Autowired
    private GenImageTaskDao genImageTaskDao;
    @Autowired
    private GenImageTaskRecordService genImageTaskRecordService;
    @Autowired
    private RedisKeyUtils redisKeyUtils;
    @Autowired
    private LimitUtils limitUtils;
    @Autowired
    private CacheUtils cacheUtils;
    @Autowired
    private StringRedisTemplate stringRedisTemplate;
    @Autowired
    private AiGenImageService aiGenImageService;
    @Autowired
    private AiClient aiClient;
    @Autowired
    private AiUserAlbumClient aiUserAlbumClient;
    @Autowired
    private PayClient payClient;
    @Autowired
    private AiChatInfoClient aiChatInfoClient;
    @Autowired
    private ContextChatClient contextChatClient;

    private static final Integer GEN_IMAGE_MAX_SIZE = 10;

    private static final Integer EXP_TIME = 24 * 60 * 60;

    /**
     * 异步任务处理的线程池
     */
    private ExecutorService executor = new ThreadPoolExecutor(10,  // 核心线程数（假设CPU核心数为8，IO密集型任务）
            30,  // 最大线程数
            60,
            TimeUnit.SECONDS,
            new LinkedBlockingQueue<>(500),  // 限制队列大小为500
            new ThreadPoolExecutor.CallerRunsPolicy() // 拒绝策略：让提交者线程执行
    );


    @Transactional(rollbackFor = Exception.class)
    @Override
    public GenImageOutput genImage(Session session, GenImageInput input) {
        //处理生日为 age
        if (StringUtils.isNotEmpty(input.getBirthday())) {
            //将生日转换成年龄
            input.setAge(DateUtils.calculateAge(input.getBirthday()));
        }
        //类型
        CreateImageType createImageType = input.getType();

        Long currentUserId = session.getUserId();
        String message = ToastResultCode.GEN_IMAGE_LIMIT_ERROR.getErrorMsg();
        //TODO 限流的逻辑：用户限流、IP限流、设备限流
//        LimitBo userLimitBo = limitUtils.defaultLimitCheckByKeyV2(redisKeyUtils.genImageUserLimit(currentUserId, input.getHl()), GEN_IMAGE_MAX_SIZE, EXP_TIME);
//        ToastResultCode.GEN_IMAGE_LIMIT_ERROR.check(userLimitBo.getLimitBl(), CheckUtils.getErrorMessage(userLimitBo, message));
//        LimitBo deviceLimitBo = limitUtils.defaultLimitCheckByKeyV2(redisKeyUtils.genImageDeviceLimit(input.getDeviceId(), input.getHl()), GEN_IMAGE_MAX_SIZE, EXP_TIME);
//        ToastResultCode.GEN_IMAGE_LIMIT_ERROR.check(deviceLimitBo.getLimitBl(), CheckUtils.getErrorMessage(deviceLimitBo, message));
//        LimitBo ipLimitBo = limitUtils.defaultLimitCheckByKeyV2(redisKeyUtils.genImageIpLimit(input.getIp(), input.getHl()), GEN_IMAGE_MAX_SIZE, EXP_TIME);
//        ToastResultCode.GEN_IMAGE_LIMIT_ERROR.check(ipLimitBo.getLimitBl(), CheckUtils.getErrorMessage(ipLimitBo, message));
        //个人介绍 敏感词检测
//        contextChatClient.nsfwCheck(input.getIntroduction());
        //生成批次号
        String batchNo = CheckUtils.genUUIDAnEncode();
        //生成6个任务id
        List<String> taskIdList = Lists.newArrayList();
        int taskCount = 6;
        for (int i = 0; i < taskCount; i++) {
            taskIdList.add(UUID.randomUUID().toString());
        }
        //解析提示词创建生成图片的异步任务，得到任务id列表
        AiGenImageInput aiGenImageInput = new AiGenImageInput();
        aiGenImageInput.setImageStylePrompt(input.getImageStylePrompt());
        aiGenImageInput.setContent(input.getContent());
        if (input.getHl() != null && input.getHl()) {
            //图片主人id
            Long imageUserId = currentUserId;
            //权限判断，查询该aiId是否是当前用户的
            AiInfoApiOutput aiInfo = aiClient.getAiInfo(input.getAiId());
            if (CreateImageType.CREATE_AI_IMAGE.equals(createImageType)) {
                //创建AI形象，1分钟内20次的限制
                boolean b = limitUtils.defaultLimitCheckByKey(redisKeyUtils.aiImageLimitKey(currentUserId), 20, 60);
                ToastResultCode.REQUEST_LIMIT_ERROR.check(b);
            }
            //编辑形象，相像图片时才做权限判断
            if (CreateImageType.EDIT_AI_IMAGE.equals(createImageType) || CreateImageType.ALBUM.equals(createImageType)) {
                //是否本人操作
                ToastResultCode.SYS_PERMISSION_DENIED.check(aiInfo != null && !aiInfo.getUserId().equals(currentUserId));
                //创建次数检验，创作次数扣减
                aiUserAlbumClient.useCreateCount(currentUserId);
            }
            //创作自定义背景，需要支付20Coin
            if (CreateImageType.BACKGROUND.equals(createImageType)) {
                //检测心动等级有没达到10级，10级才能创建背景
                AiChatInfoOutput aiChatInfo = aiChatInfoClient.getAiChatInfo(currentUserId, input.getAiId());
                List<HeartbeatLevelEnum> unlockHearbeatLevelList = aiChatInfo.getUnlockHearbeatLevelList();
                ToastResultCode.SYS_PERMISSION_DENIED.check(!unlockHearbeatLevelList.contains(HeartbeatLevelEnum.LEVEL_10));
                //发起支付
                createBackgroundPay(currentUserId);
                //创建背景时需要使用AI用户的形象图片主人id
                imageUserId = aiInfo.getUserId();
            }
            //查询出AI用户的形象图片
            aiGenImageInput.setImageUrl(aiInfo != null ? aiInfo.getBaseImageUrl() : null);
            //基图的initPrompt
            aiGenImageInput.setInitPrompt(genImageTaskRecordService.getInitPromptByImageUrl(imageUserId, aiGenImageInput.getImageUrl()));
        } else {
            aiGenImageInput.setImageReferenceUrl(input.getImageReferenceUrl());
        }
        //处理入参中的性别的问题
        if (StringUtils.isNotEmpty(input.getSex())) {
            aiGenImageInput.setSex(SexEnums.getSex(input.getSex()));
        }
        aiGenImageInput.setAge(input.getAge());
        aiGenImageInput.setIntroduction(input.getIntroduction());
        aiGenImageInput.setBirthday(input.getBirthday());

        String traceId = LogUtils.getTraceId();
        //线程池异步执行
        executor.execute(() -> {
            try {
                LogUtils.setTraceId(traceId);
                //调用Ai生成图片
                aiGenImageService.aiGenImage(aiGenImageInput, batchNo, taskIdList);
            } catch (Exception e) {
                log.error("aiGenImage error:", e);
            }
        });

        //保存数据库的逻辑
        GenImageTask genImageTask = new GenImageTask();
        genImageTask.setUserId(currentUserId);
        genImageTask.setBatchNo(batchNo);
        genImageTask.setTaskCount(taskCount);
        genImageTask.setCompletedCount(0);
        genImageTask.setPollingCount(0);
        genImageTask.setStatus(GenImageTask.Status.PENDING);
        genImageTask.setEndpoint(input.getEndpoint());
        genImageTask.setDeviceId(input.getDeviceId());
        genImageTask.setIpAddress(input.getIp());
        genImageTask.setCreateTime(LocalDateTime.now());
        genImageTask.setEditTime(LocalDateTime.now());
        save(genImageTask);

        //保存明细数据
        genImageTaskRecordService.saveRecord(genImageTask, taskIdList);

        //组装出参数据
        GenImageOutput output = new GenImageOutput();
        output.setBatchNo(batchNo);
        return output;
    }

    /**
     * 自定义背景付费
     *
     * @param currentUserId
     */
    private void createBackgroundPay(Long currentUserId) {
        //需要调用支付服务扣钱20Coin
        String orderNo = KeyGenerator.instance().generatorUniqueKey("");
        BalanceCheckoutInput balanceCheckoutInput = BalanceCheckoutInput.builder()
                .platform("Balance")
                .outTradeNo(orderNo)
                .bizType(BizType.CREATE_AI_IMAGE)
                .name(BizType.CREATE_AI_IMAGE.getDesc())
                //付款人
                .srcAccountId(currentUserId)
                //收款人
                .desAccountId(-1L)
                //总金额
                .productAmount(7000L)
                //折扣金额
                .promoAmount(0L)
                .build();
        payClient.checkoutToB(balanceCheckoutInput);
    }

    @Override
    public List<GenImageListOutput> getImage(Long userId, String batchNo) {
        //图片已完成数
        int completedCount = genImageTaskRecordService.queryCompletedTaskCountByBatchNo(batchNo);
        //图片未处理完成时来获取结果，做时间判断和限流控制
        if (completedCount == 0) {
            //解析出批次号中的时间字段，超过5分钟的批次号直接快速返回掉不做任何处理
            Long createTime = CheckUtils.uuidDecodeToTime(batchNo);
            if (createTime + 300 < System.currentTimeMillis() / 1000) {
                BizResultCode.GEN_IMAGE_TIMEOUT.check(true);
            }
            //redis接口请求限流 每5s内只能请求一次
            boolean limitBl = limitUtils.defaultLimitCheckByKey(redisKeyUtils.genImagePollingLimit(batchNo), 1, 5);
            if (limitBl) {
                return Lists.newArrayList();
            }
        }
        //从缓存中获取数据，回调时清理缓存
        //查询数据库获取图片生成结果
        List<GenImageListOutput> outputList = cacheUtils.getCacheListAndSet(redisKeyUtils.getGenImageCacheKey(batchNo), GenImageListOutput.class, () -> {
            List<GenImageListOutput> queryOutputList = genImageTaskRecordService.queryByBatchNo(batchNo);
            //如果鉴黄了，不返回地址
            for (GenImageListOutput genImageListOutput : queryOutputList) {
                if (GenImageTaskRecord.Status.NSFW.equals(genImageListOutput.getStatus())) {
                    genImageListOutput.setImageUrl(null);
                }
            }
            return JSONObject.toJSONString(queryOutputList);
        }, 60);
        //更新任务表的心跳时间字段，需要根据心跳时间来关闭掉任务
        genImageTaskDao.updateHeartBeat(batchNo);
        return outputList;
    }

    @Override
    public void genImageCallback(GenImageCallbackInput input) {
        GenImageTaskRecord genImageTaskRecord = genImageTaskRecordService.getByThirdTaskId(input.getThirdTaskId());
        if (genImageTaskRecord == null) {
            return;
        }
        //更新数据库状态
        genImageTaskRecordService.genImageCompleted(genImageTaskRecord.getId(), input.getImageUrl(), input.getStatus(), input.getPrompt(), input.getAsyncTaskId(), input.getFailReason(), input.getType());
        //查询已生成任务完成的数据总数
        int completedCount = genImageTaskRecordService.queryCompletedTaskCountByTaskId(genImageTaskRecord.getTaskId());
        //统计生成数量，判断数据库的已完成条数已经和当前已完成条数计算出状态并进行更新
        genImageTaskDao.updateCompletedCount(genImageTaskRecord.getTaskId(), completedCount);
        //清理缓存数据
        stringRedisTemplate.delete(redisKeyUtils.getGenImageCacheKey(genImageTaskRecord.getBatchNo()));
    }

    @Override
    public void delGenImageTask(Long userId, String batchNo) {
        //更新数据库状态为已停止，防止被定时任务扫描到进行处理
        update(Wrappers.<GenImageTask>lambdaUpdate()
                .set(GenImageTask::getStatus, GenImageTask.Status.RELEASED)
                .eq(GenImageTask::getBatchNo, batchNo)
                .eq(GenImageTask::getStatus, GenImageTask.Status.PENDING));
    }

    @Override
    public Integer scanHeartBeatExpTime() {
        //考虑网络延迟的问题，这里先设置成18秒内的心跳次数3次
        LocalDateTime startTime = LocalDateTime.now().minusSeconds(18);
        List<GenImageTask> taskList = genImageTaskDao.scanHeartBeatExpTime(startTime);
        for (GenImageTask genImageTask : taskList) {
            delGenImageTask(null, genImageTask.getBatchNo());
        }
        return taskList.size();
    }

}
