package com.bonus.mq.rocket;

import cn.hutool.core.util.StrUtil;
import com.bonus.mq.*;
import com.bonus.mq.tx.*;
import com.bonus.mq.rocket.tx.RocketTxHolder;
import org.apache.rocketmq.client.Validators;
import org.apache.rocketmq.client.exception.MQBrokerException;
import org.apache.rocketmq.client.exception.MQClientException;
import org.apache.rocketmq.client.impl.producer.DefaultMQProducerImpl;
import org.apache.rocketmq.client.producer.LocalTransactionState;
import org.apache.rocketmq.client.producer.SendResult;
import org.apache.rocketmq.client.producer.SendStatus;
import org.apache.rocketmq.common.message.MessageAccessor;
import org.apache.rocketmq.remoting.exception.RemotingException;
import org.apache.rocketmq.spring.core.RocketMQTemplate;
import org.apache.rocketmq.spring.support.RocketMQUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.messaging.Message;
import org.springframework.messaging.MessageHeaders;
import org.springframework.messaging.MessagingException;
import org.springframework.messaging.converter.MessageConversionException;
import org.springframework.messaging.converter.MessageConverter;
import org.springframework.messaging.converter.SmartMessageConverter;
import org.springframework.messaging.support.MessageBuilder;
import org.springframework.util.MimeTypeUtils;

public class RocketTemplate extends MQTemplate {
   private static final Logger log = LoggerFactory.getLogger(RocketTemplate.class);
   private final RocketMQTemplate rocketMQTemplate;
   private final String parentTopic;
   private final String charset = "UTF-8";
   private final MessageConverter converter;
   private final DefaultMQProducerImpl defaultMQProducerImpl;

   public RocketTemplate(RocketMQTemplate rocketMQTemplate, String parentTopic, MessageConverter messageConverter, MQProperties mqProperties) {
      super(mqProperties);
      this.rocketMQTemplate = rocketMQTemplate;
      this.defaultMQProducerImpl = rocketMQTemplate.getProducer().getDefaultMQProducerImpl();
      this.parentTopic = parentTopic;
      this.converter = messageConverter;
   }

   public void sendDelayImpl(String destination, Object payload, int delayTime) {
      Message<?> message = MessageBuilder.withPayload(payload).setHeader("X-Trace-Id", this.getTraceId()).build();
      SendResult sendResult = this.rocketMQTemplate.syncSendDelayTimeMills(destination, message, delayTime);
      if (sendResult.getSendStatus() != SendStatus.SEND_OK) {
         throw new MessagingException("Message send failed: " + String.valueOf(sendResult.getSendStatus()));
      }
   }

   public void sendImpl(String destination, Object payload) {
      Message<?> message = MessageBuilder.withPayload(payload).setHeader("X-Trace-Id", this.getTraceId()).build();
      SendResult sendResult = this.rocketMQTemplate.syncSend(destination, message);
      if (sendResult.getSendStatus() != SendStatus.SEND_OK) {
         throw new MessagingException("Message send failed: " + String.valueOf(sendResult.getSendStatus()));
      }
   }

   private org.apache.rocketmq.common.message.Message createRocketMqMessage(String destination, Object payload) {
      Message<?> msg = this.doConvert(payload);
      return RocketMQUtil.convertToRocketMessage(this.converter, "UTF-8", destination, msg);
   }

   private org.apache.rocketmq.common.message.Message createRocketMqttMessage(String destination, Object payload) {
      Message<?> msg = this.doConvert(payload);
      org.apache.rocketmq.common.message.Message toRocketMessage = RocketMQUtil.convertToRocketMessage(this.converter, "UTF-8", this.parentTopic, msg);
      toRocketMessage.putUserProperty("__SUB_TOPIC", destination);
      return toRocketMessage;
   }

   private Message<?> doConvert(Object payload) {
      Message<?> message = this.converter instanceof SmartMessageConverter ? ((SmartMessageConverter)this.converter).toMessage(payload, (MessageHeaders)null, (Object)null) : this.converter.toMessage(payload, (MessageHeaders)null);
      if (message == null) {
         String var10002 = payload.getClass().getName();
         throw new MessageConversionException("Unable to convert payload with type='" + var10002 + "', converter=[" + String.valueOf(this.converter) + "]");
      } else {
         MessageBuilder<?> builder = MessageBuilder.fromMessage(message);
         builder.setHeaderIfAbsent("contentType", MimeTypeUtils.TEXT_PLAIN);
         builder.setHeader("X-Trace-Id", this.getTraceId());
         return builder.build();
      }
   }

   public TxHolder txBeginImpl(String destination, Object payload) throws MQBrokerException, RemotingException, InterruptedException, MQClientException {
      return this.txBegin(this.createRocketMqMessage(destination, payload));
   }

   public TxHolder mqttTxBegin(String destination, Object payload) throws MQBrokerException, RemotingException, InterruptedException, MQClientException {
      return this.txBegin(this.createRocketMqttMessage(destination, payload));
   }

   public void txCommit(TxHolder txHolder) {
      try {
         RocketTxHolder rocketTxHolder = (RocketTxHolder)txHolder;
         this.defaultMQProducerImpl.endTransaction(rocketTxHolder.getMessage(), rocketTxHolder.getSendResult(), LocalTransactionState.COMMIT_MESSAGE, (Throwable)null);
      } catch (Throwable var3) {
         log.error("Commit transaction message failed, transaction state will be checked later", var3);
      }

   }

   public void txRollback(TxHolder txHolder, Throwable exception) {
      try {
         RocketTxHolder rocketTxHolder = (RocketTxHolder)txHolder;
         this.defaultMQProducerImpl.endTransaction(rocketTxHolder.getMessage(), rocketTxHolder.getSendResult(), LocalTransactionState.ROLLBACK_MESSAGE, exception);
      } catch (Throwable var4) {
         log.error("Rollback transaction message failed, transaction state will be checked later", var4);
      }

   }

   public void txUnknown(TxHolder txHolder, Throwable exception) {
      try {
         RocketTxHolder rocketTxHolder = (RocketTxHolder)txHolder;
         this.defaultMQProducerImpl.endTransaction(rocketTxHolder.getMessage(), rocketTxHolder.getSendResult(), LocalTransactionState.UNKNOW, exception);
      } catch (Throwable var4) {
         log.error("Rollback transaction message failed, transaction state will be checked later", var4);
      }

   }

   private RocketTxHolder txBegin(org.apache.rocketmq.common.message.Message msg) throws MQBrokerException, RemotingException, InterruptedException, MQClientException {
      try {
         if (msg.getDelayTimeLevel() != 0) {
            MessageAccessor.clearProperty(msg, "DELAY");
         }

         Validators.checkMessage(msg, this.rocketMQTemplate.getProducer());
         MessageAccessor.putProperty(msg, "TRAN_MSG", "true");
         MessageAccessor.putProperty(msg, "PGROUP", this.rocketMQTemplate.getProducer().getProducerGroup());
         SendResult sendResult = this.defaultMQProducerImpl.send(msg);
         if (sendResult.getSendStatus() != SendStatus.SEND_OK) {
            throw new MessagingException("Half transaction message send failed");
         } else {
            if (sendResult.getTransactionId() != null) {
               msg.putUserProperty("__transactionId__", sendResult.getTransactionId());
            }

            String transactionId = msg.getProperty("UNIQ_KEY");
            if (StrUtil.isBlank(transactionId)) {
               msg.setTransactionId(transactionId);
            }

            return new RocketTxHolder(msg, this, sendResult);
         }
      } catch (Throwable var4) {
         throw var4;
      }
   }

   public void sendMqtt(String destination, Object payload) throws MQBrokerException, RemotingException, InterruptedException, MQClientException {
      try {
         org.apache.rocketmq.common.message.Message rocketMsg = this.createRocketMqttMessage(destination, payload);
         SendResult sendResult = this.rocketMQTemplate.getProducer().send(rocketMsg, (long)this.rocketMQTemplate.getProducer().getSendMsgTimeout());
         if (sendResult.getSendStatus() != SendStatus.SEND_OK) {
            throw new MessagingException("Message send failed: " + String.valueOf(sendResult.getSendStatus()));
         }
      } catch (Throwable var5) {
         throw var5;
      }
   }
}
