/*
 * Copyright 2013-present the original author or authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.alibaba.cloud.stream.binder.rocketmq.integration.inbound.pull;

import java.util.Collection;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

import com.alibaba.cloud.stream.binder.rocketmq.integration.inbound.RocketMQConsumerFactory;
import com.alibaba.cloud.stream.binder.rocketmq.metrics.Instrumentation;
import com.alibaba.cloud.stream.binder.rocketmq.metrics.InstrumentationManager;
import com.alibaba.cloud.stream.binder.rocketmq.properties.RocketMQConsumerProperties;
import com.alibaba.cloud.stream.binder.rocketmq.support.RocketMQMessageConverterSupport;
import com.alibaba.cloud.stream.binder.rocketmq.utils.RocketMQUtils;
import org.apache.rocketmq.client.consumer.DefaultLitePullConsumer;
import org.apache.rocketmq.client.consumer.MessageSelector;
import org.apache.rocketmq.client.exception.MQClientException;
import org.apache.rocketmq.common.message.MessageExt;
import org.apache.rocketmq.common.message.MessageQueue;
import org.jspecify.annotations.Nullable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import org.springframework.beans.factory.DisposableBean;
import org.springframework.cloud.stream.binder.ExtendedConsumerProperties;
import org.springframework.context.Lifecycle;
import org.springframework.integration.IntegrationMessageHeaderAccessor;
import org.springframework.integration.endpoint.AbstractMessageSource;
import org.springframework.integration.support.MessageBuilder;
import org.springframework.messaging.Message;
import org.springframework.util.CollectionUtils;
import org.springframework.util.ObjectUtils;

/**
 * @author <a href="mailto:fangjian0423@gmail.com">Jim</a>
 */
public class RocketMQMessageSource extends AbstractMessageSource<Object>
		implements DisposableBean, Lifecycle {

	private final static Logger log = LoggerFactory
			.getLogger(RocketMQMessageSource.class);

	private @Nullable DefaultLitePullConsumer consumer;

	private final Map<String, Collection<MessageQueue>> messageQueuesForTopic = new ConcurrentHashMap<>();

	private volatile boolean running;

	private final String topic;

	private final MessageSelector messageSelector;

	private final ExtendedConsumerProperties<RocketMQConsumerProperties> extendedConsumerProperties;

	private volatile @Nullable Iterator<MessageExt> messageExtIterator = null;

	public RocketMQMessageSource(String name,
			ExtendedConsumerProperties<RocketMQConsumerProperties> extendedConsumerProperties) {
		this.topic = name;
		String subscription = extendedConsumerProperties.getExtension()
				.getSubscription();
		this.messageSelector = RocketMQUtils.getMessageSelector(
				subscription != null ? subscription : "");
		this.extendedConsumerProperties = extendedConsumerProperties;

	}

	@Override
	public synchronized void start() {
		Instrumentation instrumentation = new Instrumentation(topic, this);
		try {
			if (this.isRunning()) {
				throw new IllegalStateException(
						"pull consumer already running. " + this.toString());
			}
			this.consumer = RocketMQConsumerFactory
					.initPullConsumer(topic, extendedConsumerProperties);
			// This parameter must be 1, otherwise doReceive cannot be handled singly.
			// this.consumer.setPullBatchSize(1);
			this.consumer.subscribe(topic, messageSelector);
			this.consumer.setAutoCommit(false);
			// register TopicMessageQueueChangeListener for messageQueuesForTopic
			consumer.registerTopicMessageQueueChangeListener(topic,
					messageQueuesForTopic::put);
			this.consumer.start();
			// Initialize messageQueuesForTopic immediately
			messageQueuesForTopic.put(topic, consumer.fetchMessageQueues(topic));
			instrumentation.markStartedSuccessfully();
		}
		catch (MQClientException e) {
			instrumentation.markStartFailed(e);
			log.error("DefaultMQPullConsumer startup error: " + e.getMessage(), e);
		}
		finally {
			InstrumentationManager.addHealthInstrumentation(instrumentation);
		}
		this.running = this.consumer != null;
	}

	private @Nullable MessageQueue acquireCurrentMessageQueue(String topic, int queueId,
			String brokerName) {
		Collection<MessageQueue> messageQueueSet = messageQueuesForTopic.get(topic);
		if (CollectionUtils.isEmpty(messageQueueSet)) {
			return null;
		}
		for (MessageQueue messageQueue : messageQueueSet) {
			if (messageQueue.getQueueId() == queueId && ObjectUtils
					.nullSafeEquals(brokerName, messageQueue.getBrokerName())) {
				return messageQueue;
			}
		}
		return null;
	}

	@Override
	public synchronized void stop() {
		if (this.isRunning() && null != consumer) {
			consumer.unsubscribe(topic);
			consumer.shutdown();
			this.running = false;
		}
	}

	@Override
	public synchronized boolean isRunning() {
		return running;
	}

	@Override
	protected synchronized @Nullable Object doReceive() {
		if (consumer == null) {
			return null;
		}
		if (messageExtIterator == null) {
			List<MessageExt> messageExtList = consumer.poll();
			if (CollectionUtils.isEmpty(messageExtList)) {
				return null;
			}
			messageExtIterator = messageExtList.iterator();
		}
		MessageExt messageExt = messageExtIterator.next();
		if (!messageExtIterator.hasNext()) {
			messageExtIterator = null;
		}
		if (null == messageExt) {
			return null;
		}
		MessageQueue messageQueue = this.acquireCurrentMessageQueue(messageExt.getTopic(),
				messageExt.getQueueId(), messageExt.getBrokerName());
		if (messageQueue == null) {
			throw new IllegalArgumentException(
					"The message queue is not in assigned list");
		}
		Message message = RocketMQMessageConverterSupport
				.convertMessage2Spring(messageExt);
		return MessageBuilder.fromMessage(message)
				.setHeader(IntegrationMessageHeaderAccessor.ACKNOWLEDGMENT_CALLBACK,
						new RocketMQAckCallback(this.consumer, messageQueue, messageExt))
				.build();
	}

	@Override
	public String getComponentType() {
		return "rocketmq:message-source";
	}

}
