/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You 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
 *
 *      http://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 org.apache.camel.component.azure.eventhubs;

import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

import com.azure.messaging.eventhubs.EventProcessorClient;
import com.azure.messaging.eventhubs.models.ErrorContext;
import com.azure.messaging.eventhubs.models.EventContext;
import org.apache.camel.AsyncCallback;
import org.apache.camel.Exchange;
import org.apache.camel.Message;
import org.apache.camel.Processor;
import org.apache.camel.component.azure.eventhubs.client.EventHubsClientFactory;
import org.apache.camel.spi.Synchronization;
import org.apache.camel.support.DefaultConsumer;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import static org.apache.camel.component.azure.eventhubs.EventHubsConstants.COMPLETED_BY_SIZE;
import static org.apache.camel.component.azure.eventhubs.EventHubsConstants.COMPLETED_BY_TIMEOUT;

public class EventHubsConsumer extends DefaultConsumer {

    private static final Logger LOG = LoggerFactory.getLogger(EventHubsConsumer.class);

    // we use the EventProcessorClient as recommended by Azure docs to consume from all partitions
    private EventProcessorClient processorClient;

    private final AtomicInteger processedEvents;
    private ScheduledExecutorService scheduledExecutorService;
    private ScheduledFuture<?> lastScheduledTask;
    private EventHubsCheckpointUpdaterTask lastTask;

    public EventHubsConsumer(final EventHubsEndpoint endpoint, final Processor processor) {
        super(endpoint, processor);

        this.processedEvents = new AtomicInteger();
    }

    @Override
    protected void doStart() throws Exception {
        super.doStart();

        // create scheduled executor for checkpoint updates
        scheduledExecutorService = getEndpoint().getCamelContext().getExecutorServiceManager()
                .newScheduledThreadPool(this, "EventHubsCheckpoint", 1);

        // create the client
        processorClient = EventHubsClientFactory.createEventProcessorClient(getConfiguration(),
                this::onEventListener, this::onErrorListener);

        // start the client but we will rely on the Azure Client Scheduler for thread management
        processorClient.start();
    }

    @Override
    protected void doStop() throws Exception {
        if (processorClient != null) {
            // shutdown the client
            processorClient.stop();
            processorClient = null;
        }

        // shutdown scheduled executor
        if (scheduledExecutorService != null) {
            getEndpoint().getCamelContext().getExecutorServiceManager().shutdownGraceful(scheduledExecutorService);
            scheduledExecutorService = null;
        }

        // shutdown camel consumer
        super.doStop();
    }

    public EventHubsConfiguration getConfiguration() {
        return getEndpoint().getConfiguration();
    }

    @Override
    public EventHubsEndpoint getEndpoint() {
        return (EventHubsEndpoint) super.getEndpoint();
    }

    private Exchange createAzureEventHubExchange(final EventContext eventContext) {
        final Exchange exchange = createExchange(true);
        final Message message = exchange.getIn();

        // set body as byte[] and let camel typeConverters do the job to convert
        message.setBody(eventContext.getEventData().getBody());
        // set headers
        message.setHeader(EventHubsConstants.PARTITION_ID, eventContext.getPartitionContext().getPartitionId());
        message.setHeader(EventHubsConstants.PARTITION_KEY, eventContext.getEventData().getPartitionKey());
        message.setHeader(EventHubsConstants.OFFSET, eventContext.getEventData().getOffset());
        message.setHeader(EventHubsConstants.ENQUEUED_TIME, eventContext.getEventData().getEnqueuedTime());
        message.setHeader(EventHubsConstants.SEQUENCE_NUMBER, eventContext.getEventData().getSequenceNumber());
        if (eventContext.getEventData().getEnqueuedTime() != null) {
            long ts = eventContext.getEventData().getEnqueuedTime().getEpochSecond() * 1000;
            message.setHeader(EventHubsConstants.MESSAGE_TIMESTAMP, ts);
        }
        message.setHeader(EventHubsConstants.METADATA, eventContext.getEventData().getProperties());

        return exchange;
    }

    private Exchange createAzureEventHubExchange(final ErrorContext errorContext) {
        final Exchange exchange = createExchange(true);
        final Message message = exchange.getIn();

        // set headers
        message.setHeader(EventHubsConstants.PARTITION_ID, errorContext.getPartitionContext().getPartitionId());

        // set exception
        exchange.setException(errorContext.getThrowable());

        return exchange;
    }

    private void onEventListener(final EventContext eventContext) {
        final Exchange exchange = createAzureEventHubExchange(eventContext);

        // add exchange callback
        exchange.getExchangeExtension().addOnCompletion(new Synchronization() {
            @Override
            public void onComplete(Exchange exchange) {
                // we update the consumer offsets
                processCommit(exchange, eventContext);
            }

            @Override
            public void onFailure(Exchange exchange) {
                // we do nothing here
                processRollback(exchange);
            }
        });
        // use default consumer callback
        AsyncCallback cb = defaultConsumerCallback(exchange, true);
        getAsyncProcessor().process(exchange, cb);
    }

    private void onErrorListener(final ErrorContext errorContext) {
        final Exchange exchange = createAzureEventHubExchange(errorContext);

        // log exception if an exception occurred and was not handled
        if (exchange.getException() != null) {
            getExceptionHandler().handleException("Error processing exchange", exchange,
                    exchange.getException());
        }
    }

    /**
     * Strategy to commit the offset after message being processed successfully.
     *
     * @param exchange the exchange
     */
    private void processCommit(final Exchange exchange, final EventContext eventContext) {
        if (lastTask == null || lastTask.isExpired()) {
            lastTask = new EventHubsCheckpointUpdaterTask(eventContext, processedEvents);
            // delegate the checkpoint update to a dedicated Thread
            long timeout = getConfiguration().getCheckpointBatchTimeout();
            lastTask.setScheduledTime(System.currentTimeMillis() + timeout);
            lastScheduledTask = scheduledExecutorService.schedule(lastTask, timeout, TimeUnit.MILLISECONDS);
        } else {
            // updates the eventContext to use for the offset to be the most accurate
            lastTask.setEventContext(eventContext);
        }

        try {
            var cnt = processedEvents.incrementAndGet();
            if (cnt == getConfiguration().getCheckpointBatchSize()) {
                processedEvents.set(0);
                exchange.getIn().setHeader(EventHubsConstants.CHECKPOINT_UPDATED_BY, COMPLETED_BY_SIZE);
                LOG.debug("eventhub consumer batch size of reached");
                if (lastScheduledTask != null) {
                    lastScheduledTask.cancel(false);
                }
                eventContext.updateCheckpointAsync()
                        .subscribe(unused -> LOG.debug("Processed one event..."),
                                error -> LOG.debug("Error when updating Checkpoint: {}", error.getMessage()),
                                () -> {
                                    LOG.debug("Checkpoint updated.");
                                });
            } else if (lastTask != null && lastTask.isExpired()) {
                exchange.getIn().setHeader(EventHubsConstants.CHECKPOINT_UPDATED_BY, COMPLETED_BY_TIMEOUT);
                LOG.debug("eventhub consumer batch timeout reached");
            } else {
                LOG.debug("neither eventhub consumer batch size of {}/{} nor batch timeout reached yet", cnt,
                        getConfiguration().getCheckpointBatchSize());
            }
            // we assume that the scheduled task has done the update by its side
        } catch (Exception ex) {
            getExceptionHandler().handleException("Error occurred during updating the checkpoint. This exception is ignored.",
                    exchange, ex);
        }
    }

    /**
     * Strategy when processing the exchange failed.
     *
     * @param exchange the exchange
     */
    private void processRollback(Exchange exchange) {
        final Exception cause = exchange.getException();
        if (cause != null) {
            getExceptionHandler().handleException("Error during processing exchange.", exchange, cause);
        }
    }
}
