package cn.axzo.workflow.core.mq;

import cn.axzo.framework.rocketmq.Event;
import cn.axzo.framework.rocketmq.RocketMQEventProducer;
import cn.axzo.framework.rocketmq.utils.TraceUtils;
import com.google.common.base.Preconditions;
import com.google.common.base.Strings;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Maps;
import lombok.NonNull;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections4.ListUtils;
import org.apache.rocketmq.spring.core.RocketMQTemplate;
import org.springframework.beans.BeanUtils;

import java.util.HashMap;
import java.util.List;
import java.util.Optional;
import java.util.function.BiConsumer;

import static cn.axzo.workflow.common.constant.BpmnConstants.MQ_OWNERSHIP_APPLICATION;

/**
 * 默认的 RocketMQ 事件生产者的装饰器
 *
 * @author wangli
 * @since 2024/5/14 17:20
 */
@Slf4j
public class CustomRocketMQEventProducer extends RocketMQEventProducer {

    private BiConsumer<Event, Context<RocketMQMessageMeta>> sendBeforeCallback;
    private BiConsumer<Event, Context<RocketMQMessageMeta>> rollbackHandler;
    private final String applicationName;

    public CustomRocketMQEventProducer(RocketMQTemplate rocketMQTemplate, String defaultModule,
                                       String appName, Context<RocketMQMessageMeta> defaultContext,
                                       BiConsumer<Event, Context<RocketMQMessageMeta>> sendBeforeCallback,
                                       BiConsumer<Event, Context<RocketMQMessageMeta>> sendAfterCallback,
                                       BiConsumer<Event, Context<RocketMQMessageMeta>> rollbackHandler) {
        super(rocketMQTemplate, defaultModule, appName, defaultContext, sendAfterCallback);
        this.applicationName = appName;
        this.sendBeforeCallback = sendBeforeCallback;
        this.rollbackHandler = rollbackHandler;
    }

    @Override
    public void send(@NonNull Event event, @NonNull Context context) {
        log.info("发送事件内容:{}", event.toPrettyJsonString());
        if (sendBeforeCallback != null) {
            sendBeforeCallback.accept(event, context);
        }
        // XXX：不要在send的时候修改event的值，有副作用。
        // 例如：当将同一个event发送到不同的topic的时候，buildSchemaHash会用不同的topic赋值两次，导致一些异常case
        Event copiedEvent = Event.builder().build();
        BeanUtils.copyProperties(event, copiedEvent);
        if (Strings.isNullOrEmpty(copiedEvent.getTargetId())) {
            log.warn("targetId of event is black, best practice of targetId is present, event = {}", event.toJsonString());
        }

        if (copiedEvent.getData() == null) {
            log.warn("data of event is empty, best practice of data must present, event = {}", event.toJsonString());
        }
        Preconditions.checkArgument(!Strings.isNullOrEmpty(copiedEvent.getEventModule()), "eventModule不能为空");
        Preconditions.checkArgument(!Strings.isNullOrEmpty(copiedEvent.getEventName()), "eventName不能为空");

        // 复制一份 context，并加入链路跟踪信息 traceId
        HashMap newHeaders = Maps.newHashMap(Optional.ofNullable(context.getHeaders()).orElse(ImmutableMap.of()));
        newHeaders.put(TraceUtils.TRACE_ID, TraceUtils.getOrCreateTraceId());
        newHeaders.put(TraceUtils.CTX_LOG_ID, TraceUtils.getOrCreateTraceId());
        newHeaders.put(TraceUtils.TRACE_ID_IN_MDC, TraceUtils.getOrCreateTraceId());
        newHeaders.put(MQ_OWNERSHIP_APPLICATION, applicationName);
        final Context copiedContext = context.toBuilder().headers(newHeaders).build();

        Runnable runnable = () -> {
            try {
                getSender().accept(copiedEvent, copiedContext);
            } catch (Exception e) {
                log.error("====MQ PRODUCER ====, context={}, message = {}", copiedContext, copiedEvent.toPrettyJsonString(), e);
                throw e;
            }
        };

        Runnable rollbackRunnable = () -> {
            try {
                getRollbackHandler().accept(copiedEvent, copiedContext);
            } catch (Exception e) {
                // ignore
            }
        };
        if (copiedContext.isTransactional()) {
            //  https://www.jianshu.com/p/59891ede5f90
            log.info("runnable is transaction event={}", copiedEvent.toJsonString());
            runnable.run();
        } else {
            // 并发会导致事件时序出现问题. 所以串行执行
            log.info("runnable not transaction event={}", copiedEvent.toJsonString());
            getAfterCommitExecutor().executeAndRollback(runnable, rollbackRunnable);
        }

        List<Runnable> runnables = getAfterCommitExecutor().getRunnables();
        log.info("runnables.size(): {}", ListUtils.emptyIfNull(runnables).size());
    }

    @Override
    public BiConsumer<Event, Context<RocketMQMessageMeta>> getRollbackHandler() {
        return rollbackHandler;
    }
}
