package com.cyksj.web.controller.mirror;

import com.baomidou.mybatisplus.core.conditions.update.LambdaUpdateWrapper;
import com.baomidou.mybatisplus.core.toolkit.Wrappers;
import com.cyksj.common.exception.BusinessRuntimeException;
import com.cyksj.dto.Result;
import com.cyksj.model.dto.*;
import com.cyksj.enums.GatewayResponse;
import com.cyksj.mapper.MidjourneyUserMapper;
import com.cyksj.model.entity.MidjourneyUser;
import com.cyksj.model.entity.MidjourneyUserConversation;
import com.cyksj.redis.RedisService;
import com.cyksj.service.midjourney.MidjourneyService;
import com.ejlchina.searcher.BeanSearcher;
import com.ejlchina.searcher.SearchResult;
import com.ejlchina.searcher.util.MapBuilder;
import com.ejlchina.searcher.util.MapUtils;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.bind.annotation.*;

import javax.servlet.http.HttpServletRequest;
import java.util.Date;
import java.util.List;

/**
 * @author zwhui
 * @date 2024/4/23 10:31
 */
@Slf4j
@RequiredArgsConstructor
@RestController
@RequestMapping("/applets/midjourney")
public class MidjourneyController {

    private final MidjourneyUserMapper midjourneyUserMapper;

    private final HttpServletRequest request;

    private final MidjourneyService midjourneyService;

    private final BeanSearcher beanSearcher;
    
    private final RedisService redisService;

    public MidjourneyUser getUser(){
        String userToken = request.getHeader("user-token");
        MidjourneyUser midjourneyUser = (MidjourneyUser) redisService.get(RedisService.key.MIDJOURNEY_USER.getName() + userToken);
        if (midjourneyUser == null){
            midjourneyUser = midjourneyUserMapper.selectOne(Wrappers.lambdaQuery(MidjourneyUser.class)
                    .eq(MidjourneyUser::getUserToken, userToken).last("limit 1"));
            if (midjourneyUser == null){
                throw new BusinessRuntimeException("账号不存在,请重新登录");
            }
            redisService.set(RedisService.key.MIDJOURNEY_USER.getName() + userToken, midjourneyUser, RedisService.key.MIDJOURNEY_USER.getTimeout());
            Object fastNum = redisService.get(RedisService.key.MIDJOURNEY_FAST_LIMIT.getName() + midjourneyUser.getId());
            if (fastNum == null){
                redisService.set(RedisService.key.MIDJOURNEY_FAST_LIMIT.getName() + midjourneyUser.getId(), midjourneyUser.getMjFastNum(), RedisService.key.MIDJOURNEY_FAST_LIMIT.getTimeout());
            }
            Object relaxNum = redisService.get(RedisService.key.MIDJOURNEY_RELAX_LIMIT.getName() + midjourneyUser.getId());
            if (relaxNum == null){
                if (midjourneyUser.getMjRelaxNum() != null){
                    redisService.set(RedisService.key.MIDJOURNEY_RELAX_LIMIT.getName() + midjourneyUser.getId(), midjourneyUser.getMjRelaxNum(), RedisService.key.MIDJOURNEY_RELAX_LIMIT.getTimeout());
                }
            }
        }else {
            Object fastNum = redisService.get(RedisService.key.MIDJOURNEY_FAST_LIMIT.getName() + midjourneyUser.getId());
            if (fastNum == null){
                redisService.set(RedisService.key.MIDJOURNEY_FAST_LIMIT.getName() + midjourneyUser.getId(), midjourneyUser.getMjFastNum(), RedisService.key.MIDJOURNEY_FAST_LIMIT.getTimeout());
            }else {
                midjourneyUser.setMjFastNum(Integer.parseInt(fastNum.toString()));
            }
            Object relaxNum = redisService.get(RedisService.key.MIDJOURNEY_RELAX_LIMIT.getName() + midjourneyUser.getId());
            if (relaxNum == null){
                if (midjourneyUser.getMjRelaxNum() != null){
                    redisService.set(RedisService.key.MIDJOURNEY_RELAX_LIMIT.getName() + midjourneyUser.getId(), midjourneyUser.getMjRelaxNum(), RedisService.key.MIDJOURNEY_RELAX_LIMIT.getTimeout());
                }
            }else {
                midjourneyUser.setMjRelaxNum(Integer.parseInt(relaxNum.toString()));
            }
        }
        return midjourneyUser;
    }


    /**
     * 检查用户次数
     */
    public Long checkUserLimit(MidjourneyUser user){
        Long num = 0L;
        if (user.getExpireTime().before(new Date())){
            throw BusinessRuntimeException.getInstance("账号已过期");
        }
       if (user.getMode() == 1){
           num = redisService.decr(RedisService.key.MIDJOURNEY_FAST_LIMIT.getName() + user.getId(), 1L);
           if (num < 0) {
               user.setMode(2);
               redisService.incr(RedisService.key.MIDJOURNEY_FAST_LIMIT.getName() + user.getId(), 1L);
           }
       }
       if (user.getMode() == 2){
           Object relax = redisService.get(RedisService.key.MIDJOURNEY_RELAX_LIMIT.getName() + user.getId());
           if (relax != null){
               num = redisService.decr(RedisService.key.MIDJOURNEY_RELAX_LIMIT.getName() + user.getId(), 1L);
               if (num < 0){
                   redisService.incr(RedisService.key.MIDJOURNEY_RELAX_LIMIT.getName() + user.getId(), 1L);
               }
           }else {
               num = null;
           }
       }
       return num;
    }

    /**
     * 恢复次数
     */
    public void recoverUserLimit(Long id,Integer mode,Long num){
        if (mode == 1){
            redisService.incr(RedisService.key.MIDJOURNEY_FAST_LIMIT.getName() + id, 1L);
        }
        if (mode == 2){
            if (num != null){
                redisService.incr(RedisService.key.MIDJOURNEY_RELAX_LIMIT.getName() + id, 1L);
            }
        }
    }

    /**
     * 同步数据库
     */
    public void syncUser(Long id,Integer mode,Long num){
        log.info("同步次数 id:{},mode:{},num:{}",id,mode,num);
        if (num == null){
            return;
        }
        LambdaUpdateWrapper<MidjourneyUser> wrapper = Wrappers.lambdaUpdate(MidjourneyUser.class)
                .eq(MidjourneyUser::getId, id);
        if (mode == 1){
            wrapper.set(MidjourneyUser::getMjFastNum, num);
        }
        if (mode == 2){
            wrapper.set(MidjourneyUser::getMjRelaxNum, num);
        }
        midjourneyUserMapper.update(null, wrapper);
    }


    /**
     * 查询用户信息
     */
    @GetMapping("/whoami")
    public Result<MidjourneyUser> queryUser(){
        MidjourneyUser user = getUser();
        return GatewayResponse.SUCCESS.newBuilder().toResult(user);
    }


    /**
     * 提交Imagine任务
     */
    @PostMapping("/submit/imagine")
    public Result<MidjourneyUserConversation> submitImagine(@RequestBody SubmitImagineDTO submitImagineDTO) {
        log.info("提交Imagine任务，提示：{}，base64数组长度：{}",submitImagineDTO.getPrompt(),submitImagineDTO.getBase64Array());
        MidjourneyUser user = getUser();
        Long num = checkUserLimit(user);
        if (num != null && num < 0){
            throw BusinessRuntimeException.getInstance("次数已用完");
        }
        MidjourneyUserConversation conversation;
        try {
            conversation = midjourneyService.submitImagine(user, submitImagineDTO.getPrompt(), submitImagineDTO.getBase64Array());
        } catch (Exception e) {
            recoverUserLimit(user.getId(), user.getMode(),num);
            throw BusinessRuntimeException.getInstance(e.getMessage());
        }
        syncUser(user.getId(), user.getMode(), num);
        return GatewayResponse.SUCCESS.newBuilder().toResult(conversation);
    }

    /**
     * 提交Describe任务
     */
    @PostMapping("/submit/describe")
    public Result<MidjourneyUserConversation> submitDescribe(@RequestBody SubmitDescribeDTO submitDescribeDTO) {
        log.info("提交Describe任务");
        MidjourneyUser user = getUser();
        Long num = checkUserLimit(user);
        if (num != null && num < 0){
            throw BusinessRuntimeException.getInstance("次数已用完");
        }
        MidjourneyUserConversation conversation;
        try {
            conversation = midjourneyService.submitDescribe(user, submitDescribeDTO.getBase64());
        } catch (Exception e) {
            recoverUserLimit(user.getId(), user.getMode(),num);
            throw BusinessRuntimeException.getInstance(e.getMessage());
        }
        syncUser(user.getId(), user.getMode(), num);
        return GatewayResponse.SUCCESS.newBuilder().toResult(conversation);
    }

    /**
     * 提交Blend任务
     */
    @PostMapping("/submit/blend")
    public Result<MidjourneyUserConversation> submitBlend(@RequestBody SubmitBlendDTO submitBlendDTO) {
        log.info("提交Blend任务 dimensions:{}，base64数组长度：{}", submitBlendDTO.getDimensions(), submitBlendDTO.getBase64Array().size());
        MidjourneyUser user = getUser();
        Long num = checkUserLimit(user);
        if (num != null && num < 0){
            throw BusinessRuntimeException.getInstance("次数已用完");
        }
        MidjourneyUserConversation conversation;
        try {
            conversation = midjourneyService.submitBlend(user, submitBlendDTO.getDimensions(), submitBlendDTO.getBase64Array());
        } catch (Exception e) {
            recoverUserLimit(user.getId(), user.getMode(),num);
            throw BusinessRuntimeException.getInstance(e.getMessage());
        }
        syncUser(user.getId(), user.getMode(), num);
        return GatewayResponse.SUCCESS.newBuilder().toResult(conversation);
    }

    /**
     * 提交Modal任务
     */
    @PostMapping("/submit/modal")
    public Result<MidjourneyUserConversation> submitModal(@RequestBody SubmitModalDTO submitModalDTO) {
        log.info("提交Modal任务，taskId:{},提示：{}，base64数组长度：{}",submitModalDTO.getTaskId(),submitModalDTO.getPrompt(),submitModalDTO.getMaskBase64());
        MidjourneyUser user = getUser();
        Long num = checkUserLimit(user);
        if (num != null && num < 0){
            throw BusinessRuntimeException.getInstance("次数已用完");
        }
        MidjourneyUserConversation conversation;
        try {
            conversation = midjourneyService.submitModal(user, submitModalDTO.getTaskId(), submitModalDTO.getPrompt(), submitModalDTO.getMaskBase64());
        } catch (Exception e) {
            recoverUserLimit(user.getId(), user.getMode(),num);
            throw BusinessRuntimeException.getInstance(e.getMessage());
        }
        syncUser(user.getId(), user.getMode(), num);
        return GatewayResponse.SUCCESS.newBuilder().toResult(conversation);
    }

    /**
     * 提交Shorten任务
     */
    @PostMapping("/submit/shorten")
    public Result<MidjourneyUserConversation> submitShorten(@RequestBody SubmitShortenDTO submitShortenDTO) {
        log.info("提交Shorten任务 提示词:{}",submitShortenDTO.getPrompt());
        MidjourneyUser user = getUser();
        Long num = checkUserLimit(user);
        if (num != null && num < 0){
            throw BusinessRuntimeException.getInstance("次数已用完");
        }
        MidjourneyUserConversation conversation;
        try {
            conversation = midjourneyService.submitShorten(user, submitShortenDTO.getPrompt());
        } catch (Exception e) {
            recoverUserLimit(user.getId(), user.getMode(),num);
            throw BusinessRuntimeException.getInstance(e.getMessage());
        }
        syncUser(user.getId(), user.getMode(), num);
        return GatewayResponse.SUCCESS.newBuilder().toResult(conversation);
    }

    /**
     * 执行动作
     */
    @PostMapping("/submit/action")
    public Result<MidjourneyUserConversation> action(@RequestBody SubmitActionDTO actionDTO) {
        log.info("任务id:{},执行动作：{}",actionDTO.getTaskId(),actionDTO.getCustomId());
        MidjourneyUser user = getUser();
        Long num = checkUserLimit(user);
        if (num != null && num < 0){
            throw BusinessRuntimeException.getInstance("次数已用完");
        }
        MidjourneyUserConversation conversation;
        try {
            conversation = midjourneyService.submitAction(user, actionDTO.getTaskId(), actionDTO.getCustomId());
        } catch (Exception e) {
            recoverUserLimit(user.getId(), user.getMode(),num);
            throw BusinessRuntimeException.getInstance(e.getMessage());
        }
        syncUser(user.getId(), user.getMode(), num);
        return GatewayResponse.SUCCESS.newBuilder().toResult(conversation);
    }


    /**
     * 查询会话列表
     */
    @GetMapping("/conversation/list")
    public Result<SearchResult<MidjourneyUserConversation>> conversationList(){
        MidjourneyUser user = getUser();
        log.info("查询会话列表 userId:{}",user.getId());
        MapBuilder builder = MapUtils.flatBuilder(request.getParameterMap());
        return GatewayResponse.SUCCESS.newBuilder().toResult(beanSearcher.search(MidjourneyUserConversation.class,
                builder.field(MidjourneyUserConversation::getUserId, user.getId())
                        .orderBy(MidjourneyUserConversation::getId).desc().build()));
    }

    /**
     * 根据ids查询会话
     */
    @GetMapping("/conversation/listByIds")
    public Result<List<MidjourneyUserConversation>> conversationListByIds(Integer mode,@RequestParam(value = "ids" ) List<Long> ids) throws Exception {
        MidjourneyUser user = getUser();
        log.info("根据ids查询会话 userId:{}",user.getId());
        return GatewayResponse.SUCCESS.newBuilder().toResult(midjourneyService.listConversationByIds(mode,ids));
    }

    /**
     * 取消任务
     */
    @PostMapping("/conversation/{id}/cancel")
    public Result<MidjourneyUserConversation> conversationCancel(@PathVariable("id") Long id){
        MidjourneyUser user = getUser();
        log.info("取消任务 id:{}",id);
        return GatewayResponse.SUCCESS.newBuilder().toResult(midjourneyService.cancelConversation(user,id));
    }


    /**
     * 修改midjourney账号模式
     */
    @PutMapping("/put/midjourney/user/{mode}")
    public Result<String> putMidjourneyUserMode(@PathVariable Integer mode, Long userId) {
        log.info("修改midjourney账号模式 userId:{},mode:{}", userId, mode);
        MidjourneyUser midjourneyUser = midjourneyUserMapper.selectById(userId);
        if (midjourneyUser == null) {
            throw BusinessRuntimeException.getInstance("参数错误");
        }
        midjourneyUser.setMode(mode);
        midjourneyUserMapper.updateById(midjourneyUser);
        redisService.del(RedisService.key.MIDJOURNEY_USER.getName() + midjourneyUser.getUserToken());
        return GatewayResponse.SUCCESS.newBuilder().toResult();
    }
}
