package com.lanyuanxiaoyao.service.sync.functions;

import cn.hutool.core.util.ObjectUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.http.HttpUtil;
import com.lanyuanxiaoyao.service.common.Constants;
import com.lanyuanxiaoyao.service.common.entity.FlinkJob;
import com.lanyuanxiaoyao.service.common.entity.TableMeta;
import com.lanyuanxiaoyao.service.common.utils.LogHelper;
import com.lanyuanxiaoyao.service.common.utils.NameHelper;
import com.lanyuanxiaoyao.service.common.utils.RecordHelper;
import com.lanyuanxiaoyao.service.sync.configuration.GlobalConfiguration;
import com.lanyuanxiaoyao.service.sync.metrics.MessageSizeSizeMetric;
import com.lanyuanxiaoyao.service.sync.metrics.RateMetric;
import com.lanyuanxiaoyao.service.sync.utils.LoadBalance;
import com.lanyuanxiaoyao.service.sync.utils.MetricsUtils;
import com.lanyuanxiaoyao.service.sync.utils.StatusUtils;
import dev.failsafe.Failsafe;
import dev.failsafe.RetryPolicy;
import java.time.Duration;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.AtomicReference;
import org.apache.flink.api.common.state.CheckpointListener;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.runtime.state.FunctionInitializationContext;
import org.apache.flink.runtime.state.FunctionSnapshotContext;
import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction;
import org.apache.flink.streaming.api.functions.source.RichParallelSourceFunction;
import org.apache.pulsar.client.api.*;
import org.apache.pulsar.client.impl.schema.StringSchema;
import org.apache.pulsar.client.internal.DefaultImplementation;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import static com.lanyuanxiaoyao.service.common.utils.LogHelper.LogPoint.*;

/**
 * Pulsar Reader Source
 *
 * @author ZhangJiacheng
 * @date 2022-06-11
 */
public class PulsarMessageSourceReader extends RichParallelSourceFunction<String> implements CheckpointedFunction, CheckpointListener {
    private static final Logger logger = LoggerFactory.getLogger(PulsarMessageSourceReader.class);
    private static final RetryPolicy<String> MESSAGE_ID_RETRY = RetryPolicy.<String>builder()
            .handle(Exception.class)
            .withDelay(Duration.ofSeconds(1))
            .withMaxAttempts(10)
            .build();
    private final String topic;
    private final GlobalConfiguration globalConfiguration;
    private final FlinkJob flinkJob;
    private final TableMeta tableMeta;
    private final AtomicReference<MessageId> lastMessageId = new AtomicReference<>();
    private final AtomicLong lastPublishTime = new AtomicLong(0);
    private final RateMetric messageReceiveMetric;
    private final MessageSizeSizeMetric messageSizeReceiveMetric;
    private final Map<Long, MessageId> messageIdMap = new ConcurrentHashMap<>();
    private boolean running = true;
    private PulsarClient client = null;
    private Reader<String> reader = null;

    public PulsarMessageSourceReader(GlobalConfiguration globalConfiguration, FlinkJob flinkJob, TableMeta tableMeta) {
        logger.info("Use PulsarMessageSourceReader");
        this.globalConfiguration = globalConfiguration;
        this.flinkJob = flinkJob;
        this.topic = tableMeta.getTopic();
        this.tableMeta = tableMeta;

        String messageId = tableMeta.getConfig().getMessageId();
        logger.info("{} {}", Constants.LOG_POINT_PULSAR_SOURCE_BOOTSTRAP_MESSAGE_ID, messageId);
        if (StrUtil.isNotBlank(messageId)) {
            lastMessageId.set(parseMessageId(messageId));
        } else {
            logger.warn("Message id is empty");
            lastMessageId.set(MessageId.earliest);
        }

        messageReceiveMetric = new RateMetric(
                globalConfiguration,
                Constants.METRICS_SYNC_SOURCE_MESSAGE_RECEIVE,
                flinkJob,
                tableMeta
        );
        messageSizeReceiveMetric = new MessageSizeSizeMetric(
                globalConfiguration,
                Constants.METRICS_SYNC_SOURCE_MESSAGE_SIZE_RECEIVE_BYTES,
                flinkJob, tableMeta
        );
    }

    private static MessageId parseMessageId(String messageIdText) {
        return DefaultImplementation.newMessageId(Long.parseLong(messageIdText.split(":")[0]), Long.parseLong(messageIdText.split(":")[1]), -1);
    }

    @Override
    public void initializeState(FunctionInitializationContext context) throws Exception {
        LogHelper.info(logger, CHECKPOINT_INITIAL);
        String queryUrl = StrUtil.format(
                "{}/api/message_id?flink_job_id={}&alias={}",
                LoadBalance.getCustomPublishUrl(globalConfiguration),
                flinkJob.getId(),
                tableMeta.getAlias()
        );
        logger.info("Query url: {}", queryUrl);
        String messageId = Failsafe.with(MESSAGE_ID_RETRY)
                .onFailure(event -> {
                    if (ObjectUtil.isNotNull(event.getException())) {
                        logger.error(StrUtil.format("{} Get message id error", Constants.LOG_POINT_PULSAR_SOURCE_GET_MESSAGE_ID_ERROR), event.getException());
                    }
                })
                .get(() ->
                        HttpUtil.createGet(queryUrl)
                                .header(Constants.API_HEADER_NAME, Constants.API_VERSION)
                                .execute()
                                .body()
                );
        LogHelper.info(logger, CHECKPOINT_INITIAL_MESSAGE_ID, "Get message id: {}", messageId);
        if (StrUtil.isNotBlank(messageId)) {
            lastMessageId.set(parseMessageId(messageId));
        } else {
            logger.warn(StrUtil.format("{} Message id is empty, now message id is {}", Constants.LOG_POINT_MESSAGE_ID_EMPTY, lastMessageId.get()));
        }
    }

    @Override
    public void run(SourceContext<String> context) throws Exception {
        String currentValue = null;
        while (running) {
            Message<String> message;
            try {
                message = reader.readNext();
                if (ObjectUtil.isNotNull(message)) {
                    String value = message.getValue();
                    currentValue = value;
                    if (ObjectUtil.isEmpty(value)) {
                        logger.warn("{} {}", message.getValue(), message.getMessageId());
                    }
                    synchronized (context.getCheckpointLock()) {
                        context.collect(value);
                    }

                    if (RecordHelper.isNotVersionUpdateRecord(value)) {
                        lastPublishTime.set(message.getPublishTime());
                    }
                    lastMessageId.set(message.getMessageId());

                    messageReceiveMetric.increment();
                    try {
                        messageSizeReceiveMetric.increment(message.getValue().getBytes().length);
                    } catch (Throwable t) {
                        logger.warn("Parse message size failure", t);
                    }
                }
            } catch (Throwable t) {
                logger.error("Read message failure, current value: " + currentValue, t);
            }
        }
    }

    @Override
    public void open(Configuration configuration) throws Exception {
        super.open(configuration);

        // 初始化指标
        MetricsUtils.createMakePointTimer(globalConfiguration, messageReceiveMetric);
        MetricsUtils.createMakePointTimer(globalConfiguration, messageSizeReceiveMetric);

        try {
            client = PulsarClient.builder()
                    .serviceUrl(tableMeta.getPulsarAddress())
                    .build();
            reader = client.newReader(new StringSchema())
                    .topic(topic)
                    .receiverQueueSize(10000)
                    .subscriptionName(NameHelper.pulsarSubscriptionName(flinkJob.getId(), tableMeta.getAlias()))
                    .startMessageId(lastMessageId.get())
                    .startMessageIdInclusive()
                    .create();
        } catch (Exception exception) {
            logger.error(StrUtil.format("Connect pulsar error ({} {})", tableMeta.getPulsarAddress(), topic), exception);
            throw exception;
        }
        logger.info("Message id set to {}", lastMessageId.get());
    }

    @Override
    public void cancel() {
        running = false;
    }

    @Override
    public void close() throws Exception {
        super.close();
        if (reader != null) {
            try {
                reader.close();
            } catch (PulsarClientException e) {
                logger.error("Pulsar reader close error", e);
            }
        }
        if (client != null) {
            try {
                client.close();
            } catch (PulsarClientException e) {
                logger.error("Pulsar client close error", e);
            }
        }
    }

    @Override
    public void snapshotState(FunctionSnapshotContext context) {
        MessageId messageId = lastMessageId.get();
        messageIdMap.put(context.getCheckpointId(), messageId);
        LogHelper.info(logger, CHECKPOINT_START, "Checkpoint start message id: {}, checkpoint id: {}", messageId, context.getCheckpointId());

        messageReceiveMetric.publish();
        messageSizeReceiveMetric.publish();
    }

    @Override
    public void notifyCheckpointComplete(long checkpointId) {
        MessageId messageId = messageIdMap.getOrDefault(checkpointId, MessageId.earliest);
        LogHelper.info(logger, CHECKPOINT_COMPLETE, "Checkpoint complete message id: {}, checkpoint id: {}", messageId, checkpointId);
        StatusUtils.syncCheckpoint(globalConfiguration, flinkJob, tableMeta, messageId.toString(), lastPublishTime.get());
        messageIdMap.remove(checkpointId);
    }

    @Override
    public void notifyCheckpointAborted(long checkpointId) throws Exception {
        CheckpointListener.super.notifyCheckpointAborted(checkpointId);
    }
}
