package com.ai.da.common.RabbitMQ;

import com.ai.da.common.config.exception.BusinessException;
import com.ai.da.common.utils.RedisUtil;
import com.ai.da.model.dto.SuperResolutionDTO;
import com.ai.da.model.dto.TaskDTO;
import com.ai.da.service.SuperResolutionService;
import com.ai.da.service.TaskListService;
import com.alibaba.fastjson.JSONException;
import com.alibaba.fastjson.JSONObject;
import com.google.gson.Gson;
import com.google.gson.reflect.TypeToken;
import com.rabbitmq.client.Channel;
import lombok.extern.slf4j.Slf4j;
import org.springframework.amqp.core.Message;
import org.springframework.amqp.rabbit.annotation.RabbitHandler;
import org.springframework.amqp.rabbit.annotation.RabbitListener;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import java.io.IOException;
import java.lang.reflect.Type;
import java.util.HashMap;

@Slf4j
@Component
public class SRConsumer {

    @Autowired
    private RabbitMQProperties rabbitMQProperties;

    @Resource
    private RedisUtil redisUtil;

    @Resource
    private TaskListService taskListService;

    @Value("${redis.key.orderForSR}")
    private String consumptionOrderKey;

    @Value("${redis.key.SRCancelSet}")
    private String cancelSetKey;

    @Value("${redis.key.SRExceptionMap}")
    private String exceptionMapKey;

    @Value("${redis.key.taskList}")
    private String taskListKey;

    @Resource
    private SuperResolutionService superResolutionService;

    /**
     * 请求超分处理
     */
    public void superResolution(Message msg, Channel channel, String consumerName) {
        log.info("============SR start listening==========");
        long start = System.currentTimeMillis();

        SuperResolutionDTO superResolutionDTO;
        String uniqueId = null;

        try {
            superResolutionDTO = JSONObject.parseObject(msg.getBody(), SuperResolutionDTO.class);
            uniqueId = superResolutionDTO.getUniqueId();
            log.info("From " + consumerName + " : " + uniqueId);
            superResolutionService.updateSROutput(uniqueId, "Executing", null);
            taskListService.updateTaskStatusOrOutputRedis(uniqueId, "Executing", null);
            /*try {
                    Thread.sleep(2 * 60 * 1000);
                } catch (InterruptedException e) {
                    throw new RuntimeException(e);
                }*/
            // 2、判断当前消息是否在取消列表中
            Boolean isMember = redisUtil.isElementExistsInSet(cancelSetKey, uniqueId);
            if (isMember) {
                try {
                    // 2.1 手动确认该消息
                    channel.basicAck(msg.getMessageProperties().getDeliveryTag(), false);
                } catch (IOException ex) {
                    log.error("手动确认，不返回队列重新消费");
                }
            } else {
                // 请求python端进行超分
                superResolutionService.SR(superResolutionDTO);
            }
        } catch (BusinessException e) {
            log.error(e.getMsg());
            superResolutionDTO = JSONObject.parseObject(msg.getBody(), SuperResolutionDTO.class);
            // channel.basicNack() 为不确认deliveryTag对应的消息，第二个参数是否应用于多消息，第三个参数是否requeue
            setErrorMessage(msg, channel, e.getMsg(), superResolutionDTO);
            try {
                // 2.1 手动确认该消息
                channel.basicAck(msg.getMessageProperties().getDeliveryTag(), false);
            } catch (IOException ex) {
                log.error("手动确认，不返回队列重新消费, error message : " + e.getMessage());
            }
        } catch (JSONException e) {
            log.error(e.getMessage());
            setErrorMessage(msg, channel, e.getMessage(), null);
            try {
                // 2.1 手动确认该消息
                channel.basicAck(msg.getMessageProperties().getDeliveryTag(), false);
            } catch (IOException ex) {
                log.error("手动确认，不返回队列重新消费, error message : " + e.getMessage());
            }
        } catch (Exception e) {
            log.error(e.getMessage());
            superResolutionDTO = JSONObject.parseObject(msg.getBody(), SuperResolutionDTO.class);
            setErrorMessage(msg, channel, e.getMessage(), superResolutionDTO);
            try {
                // 2.1 手动确认该消息
                channel.basicAck(msg.getMessageProperties().getDeliveryTag(), false);
            } catch (IOException ex) {
                log.error("手动确认，不返回队列重新消费, error message : " + e.getMessage());
            }
        }

        long end = System.currentTimeMillis();

        log.info(" task_id： " + uniqueId + "----------" + consumerName + " 执行时长：" + (end - start) + "毫秒");
        log.info("=============SR end listening===========");
    }

    /**
     * 获取超分结果
     */
    public void getSRResult(Message msg, Channel channel, String consumerName) {
        log.info("============SRResult start listening==========");
        long start = System.currentTimeMillis();

        JSONObject result = null;
        String taskId = null;

        try {
            result = JSONObject.parseObject(msg.getBody(), JSONObject.class);
            log.info("SR response : {}", result);
            taskId = result.get("tasks_id").toString();
        } catch (JSONException e) {
            log.error("SRResult 返回数据格式不合规范");
            log.error(e.getMessage());
            setErrorMessage(msg, channel, e.getMessage(), null);
            try {
                // 第二个参数，是否批量确认消息，当传false时，只确认当前 deliveryTag对应的消息;当传true时，会确认当前及之前所有未确认的消息。
                channel.basicAck(msg.getMessageProperties().getDeliveryTag(), false);
            } catch (IOException exception) {
                log.error("手动确认，取消返回队列，不再重新消费");
            }
        } catch (Exception e){
            log.error(e.getMessage());
            setErrorMessage(msg, channel, e.getMessage(), null);
            try {
                // 第二个参数，是否批量确认消息，当传false时，只确认当前 deliveryTag对应的消息;当传true时，会确认当前及之前所有未确认的消息。
                channel.basicAck(msg.getMessageProperties().getDeliveryTag(), false);
            } catch (IOException exception) {
                log.error("手动确认，取消返回队列，不再重新消费");
            }
        }


        try {
            // 2、判断状态是否成功
            if ("SUCCESS".equals(result.get("status").toString())) {
                String output = result.get("data").toString();
                superResolutionService.setSRResult(taskId, output, "success");
                taskListService.updateTaskStatusOrOutputRedis(taskId, "success", output);
            } else {
                superResolutionService.setSRResult(taskId, null, "fail");
                taskListService.updateTaskStatusOrOutputRedis(taskId, "fail", null);
                HashMap<String, String> exceptionInfo = new HashMap<>();
                // 获取输入信息
                String task = redisUtil.getFromString(taskListKey + taskId + taskId.substring(taskId.lastIndexOf("-") + 1));
                Gson gson = new Gson();
                Type type = new TypeToken<TaskDTO<SuperResolutionDTO>>() {
                }.getType();
                TaskDTO<SuperResolutionDTO> taskDTO = gson.fromJson(task, type);
                // 将输入信息和报错信息均存入redis todo 加判空
                exceptionInfo.put(taskId, "Input ==> " + taskDTO.getInputParam() + "Fail Message ==> " + result.get("message").toString());
                // 将报错信息存入redis
                redisUtil.addToMap(exceptionMapKey, exceptionInfo);
            }
        } catch (Exception e) {
            log.error(e.getMessage());
            // channel.basicNack() 为不确认deliveryTag对应的消息，第二个参数是否应用于多消息，第三个参数是否requeue
            try {
                // 第二个参数，是否批量确认消息，当传false时，只确认当前 deliveryTag对应的消息;当传true时，会确认当前及之前所有未确认的消息。
                channel.basicAck(msg.getMessageProperties().getDeliveryTag(), false);
            } catch (IOException exception) {
                log.error("手动确认，取消返回队列，不再重新消费");
            }
        }

        long end = System.currentTimeMillis();

        log.info(" task_id： " + taskId + "----------" + consumerName + " 执行时长：" + (end - start) + "毫秒");
        log.info("=============SRResult end listening===========");
    }

    private void setErrorMessage(Message msg, Channel channel, String message, SuperResolutionDTO superResolutionDTO) {
        String uniqueId;
        try {
            // 第二个参数，是否批量确认消息，当传false时，只确认当前 deliveryTag对应的消息;当传true时，会确认当前及之前所有未确认的消息。
            channel.basicAck(msg.getMessageProperties().getDeliveryTag(), false);
            uniqueId = superResolutionDTO.getUniqueId();
            // 将消息从redis排队队列中删除,需保证被消费的消息存储到db之后再从redis删除
            redisUtil.removeFromZSet(consumptionOrderKey, uniqueId);
        } catch (Exception exception) {
            log.error("手动确认，取消返回队列，不再重新消费");
            throw new BusinessException("message.confirm.fail");
        }
        // 将入参和错误信息存入redis
        String exceptionMessage = JSONObject.toJSONString(superResolutionDTO) +
                " Exception message ： " + message;
//                    " Exception message ： " + e.getMessage();
        HashMap<String, String> exceptionInfo = new HashMap<>();
        uniqueId = superResolutionDTO.getUniqueId();
        exceptionInfo.put(String.valueOf(uniqueId), exceptionMessage);
        // 存redis
        redisUtil.addToMap(exceptionMapKey, exceptionInfo);
        taskListService.updateTaskStatusOrOutputRedis(uniqueId, "fail", null);
    }

    @RabbitListener(queues = "#{rabbitMQProperties.queues.sr}")
    @RabbitHandler
    public void SRConsumer1(Message msg, Channel channel) {
        superResolution(msg, channel, "consumer 1");
    }

    @RabbitListener(queues = "#{rabbitMQProperties.queues.srResult}")
    @RabbitHandler
    public void SRResultConsumer1(Message msg, Channel channel) {
        getSRResult(msg, channel, "consumer 1");
    }

}
