/*
 * 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.kafka.streams.processor.internals;

import org.apache.kafka.clients.producer.Producer;
import org.apache.kafka.clients.producer.ProducerRecord;
import org.apache.kafka.common.KafkaException;
import org.apache.kafka.common.PartitionInfo;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.errors.AuthenticationException;
import org.apache.kafka.common.errors.AuthorizationException;
import org.apache.kafka.common.errors.InvalidPidMappingException;
import org.apache.kafka.common.errors.InvalidProducerEpochException;
import org.apache.kafka.common.errors.InvalidTopicException;
import org.apache.kafka.common.errors.OffsetMetadataTooLarge;
import org.apache.kafka.common.errors.OutOfOrderSequenceException;
import org.apache.kafka.common.errors.ProducerFencedException;
import org.apache.kafka.common.errors.RetriableException;
import org.apache.kafka.common.errors.SecurityDisabledException;
import org.apache.kafka.common.errors.SerializationException;
import org.apache.kafka.common.errors.TimeoutException;
import org.apache.kafka.common.errors.TransactionAbortedException;
import org.apache.kafka.common.errors.UnknownServerException;
import org.apache.kafka.common.header.Headers;
import org.apache.kafka.common.header.internals.RecordHeaders;
import org.apache.kafka.common.metrics.Sensor;
import org.apache.kafka.common.serialization.Serializer;
import org.apache.kafka.common.utils.LogContext;
import org.apache.kafka.streams.errors.ProductionExceptionHandler;
import org.apache.kafka.streams.errors.StreamsException;
import org.apache.kafka.streams.errors.TaskCorruptedException;
import org.apache.kafka.streams.errors.TaskMigratedException;
import org.apache.kafka.streams.errors.internals.DefaultErrorHandlerContext;
import org.apache.kafka.streams.errors.internals.FailedProcessingException;
import org.apache.kafka.streams.processor.RecordContext;
import org.apache.kafka.streams.processor.StreamPartitioner;
import org.apache.kafka.streams.processor.TaskId;
import org.apache.kafka.streams.processor.internals.metrics.StreamsMetricsImpl;
import org.apache.kafka.streams.processor.internals.metrics.TaskMetrics;
import org.apache.kafka.streams.processor.internals.metrics.TopicMetrics;

import org.slf4j.Logger;

import java.text.MessageFormat;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.atomic.AtomicReference;

import static org.apache.kafka.streams.processor.internals.ClientUtils.producerRecordSizeInBytes;

public class RecordCollectorImpl implements RecordCollector {
    private static final String SEND_EXCEPTION_MESSAGE = "Error encountered sending record to topic %s for task %s due to:%n%s";

    private final Logger log;
    private final TaskId taskId;
    private final StreamsProducer streamsProducer;
    private final ProductionExceptionHandler productionExceptionHandler;
    private final Map<TopicPartition, Long> offsets;

    private final StreamsMetricsImpl streamsMetrics;
    private final Sensor droppedRecordsSensor;
    private final Map<String, Sensor> producedSensorByTopic = new HashMap<>();

    // we get `sendException` from "singleton" `StreamsProducer` to share it across all instances of `RecordCollectorImpl`
    private final AtomicReference<KafkaException> sendException;

    /**
     * @throws StreamsException fatal error that should cause the thread to die (from producer.initTxn)
     */
    public RecordCollectorImpl(final LogContext logContext,
                               final TaskId taskId,
                               final StreamsProducer streamsProducer,
                               final ProductionExceptionHandler productionExceptionHandler,
                               final StreamsMetricsImpl streamsMetrics,
                               final ProcessorTopology topology) {
        this.log = logContext.logger(getClass());
        this.taskId = taskId;
        this.streamsProducer = streamsProducer;
        this.sendException = streamsProducer.sendException();
        this.productionExceptionHandler = productionExceptionHandler;
        this.streamsMetrics = streamsMetrics;

        final String threadId = Thread.currentThread().getName();
        this.droppedRecordsSensor = TaskMetrics.droppedRecordsSensor(threadId, taskId.toString(), streamsMetrics);
        for (final String topic : topology.sinkTopics()) {
            final String processorNodeId = topology.sink(topic).name();
            producedSensorByTopic.put(
                topic,
                TopicMetrics.producedSensor(
                    threadId,
                    taskId.toString(),
                    processorNodeId,
                    topic,
                    streamsMetrics
                ));
        }

        this.offsets = new HashMap<>();
    }

    @Override
    public void initialize() {
        if (streamsProducer.eosEnabled()) {
            streamsProducer.initTransaction();
        }
    }

    /**
     * @throws StreamsException fatal error that should cause the thread to die
     * @throws TaskMigratedException recoverable error that would cause the task to be removed
     */
    @Override
    public <K, V> void send(final String topic,
                            final K key,
                            final V value,
                            final Headers headers,
                            final Long timestamp,
                            final Serializer<K> keySerializer,
                            final Serializer<V> valueSerializer,
                            final String processorNodeId,
                            final InternalProcessorContext<Void, Void> context,
                            final StreamPartitioner<? super K, ? super V> partitioner) {

        if (partitioner != null) {
            final List<PartitionInfo> partitions;
            try {
                partitions = streamsProducer.partitionsFor(topic);
            } catch (final TimeoutException timeoutException) {
                log.warn("Could not get partitions for topic {}, will retry", topic);

                // re-throw to trigger `task.timeout.ms`
                throw timeoutException;
            } catch (final KafkaException fatal) {
                // here we cannot drop the message on the floor even if it is a transient timeout exception,
                // so we treat everything the same as a fatal exception
                throw new StreamsException("Could not determine the number of partitions for topic '" + topic +
                    "' for task " + taskId + " due to " + fatal,
                    fatal
                );
            }
            if (!partitions.isEmpty()) {
                final Optional<Set<Integer>> maybeMulticastPartitions = partitioner.partitions(topic, key, value, partitions.size());
                if (maybeMulticastPartitions.isEmpty()) {
                    // A null//empty partition indicates we should use the default partitioner
                    send(topic, key, value, headers, null, timestamp, keySerializer, valueSerializer, processorNodeId, context);
                } else {
                    final Set<Integer> multicastPartitions = maybeMulticastPartitions.get();
                    if (multicastPartitions.isEmpty()) {
                        // If a record is not to be sent to any partition, mark it as a dropped record.
                        log.warn("Skipping record as partitioner returned empty partitions. "
                                + "topic=[{}]", topic);
                        droppedRecordsSensor.record();
                    } else {
                        for (final int multicastPartition: multicastPartitions) {
                            send(topic, key, value, headers, multicastPartition, timestamp, keySerializer, valueSerializer, processorNodeId, context);
                        }
                    }
                }
            } else {
                throw new StreamsException("Could not get partition information for topic " + topic + " for task " + taskId +
                    ". This can happen if the topic does not exist.");
            }
        } else {
            send(topic, key, value, headers, null, timestamp, keySerializer, valueSerializer, processorNodeId, context);
        }

    }

    @Override
    public <K, V> void send(final String topic,
                            final K key,
                            final V value,
                            final Headers headers,
                            final Integer partition,
                            final Long timestamp,
                            final Serializer<K> keySerializer,
                            final Serializer<V> valueSerializer,
                            final String processorNodeId,
                            final InternalProcessorContext<Void, Void> context) {
        checkForException();

        final byte[] keyBytes;
        final byte[] valBytes;
        try {
            keyBytes = keySerializer.serialize(topic, headers, key);
        } catch (final ClassCastException exception) {
            throw createStreamsExceptionForClassCastException(
                ProductionExceptionHandler.SerializationExceptionOrigin.KEY,
                topic,
                key,
                keySerializer,
                exception);
        } catch (final Exception serializationException) {
            // while Java distinguishes checked vs unchecked exceptions, other languages
            // like Scala or Kotlin do not, and thus we need to catch `Exception`
            // (instead of `RuntimeException`) to work well with those languages
            handleException(
                ProductionExceptionHandler.SerializationExceptionOrigin.KEY,
                topic,
                key,
                value,
                headers,
                partition,
                timestamp,
                processorNodeId,
                context,
                serializationException
            );
            return;
        }

        try {
            valBytes = valueSerializer.serialize(topic, headers, value);
        } catch (final ClassCastException exception) {
            throw createStreamsExceptionForClassCastException(
                ProductionExceptionHandler.SerializationExceptionOrigin.VALUE,
                topic,
                value,
                valueSerializer,
                exception);
        } catch (final Exception serializationException) {
            // while Java distinguishes checked vs unchecked exceptions, other languages
            // like Scala or Kotlin do not, and thus we need to catch `Exception`
            // (instead of `RuntimeException`) to work well with those languages
            handleException(
                ProductionExceptionHandler.SerializationExceptionOrigin.VALUE,
                topic,
                key,
                value,
                headers,
                partition,
                timestamp,
                processorNodeId,
                context,
                serializationException);
            return;
        }

        final ProducerRecord<byte[], byte[]> serializedRecord = new ProducerRecord<>(topic, partition, timestamp, keyBytes, valBytes, headers);

        // As many records could be in-flight,
        // freeing raw records in the context to reduce memory pressure
        freeRawInputRecordFromContext(context);

        send(key, value, processorNodeId, context, serializedRecord);
    }

    public <K, V> void send(final K key,
                            final V value,
                            final String processorNodeId,
                            final InternalProcessorContext<?, ?> context,
                            final ProducerRecord<byte[], byte[]> serializedRecord) {

        streamsProducer.send(serializedRecord, (metadata, exception) -> {
            try {
                // if there's already an exception record, skip logging offsets or new exceptions
                if (sendException.get() != null) {
                    return;
                }

                if (exception == null) {
                    final TopicPartition tp = new TopicPartition(metadata.topic(), metadata.partition());
                    if (metadata.offset() >= 0L) {
                        offsets.put(tp, metadata.offset());
                    } else {
                        log.warn("Received offset={} in produce response for {}", metadata.offset(), tp);
                    }

                    if (!serializedRecord.topic().endsWith("-changelog")) {
                        // we may not have created a sensor during initialization if the node uses dynamic topic routing,
                        // as all topics are not known up front, so create the sensor for this topic if absent
                        final Sensor topicProducedSensor = producedSensorByTopic.computeIfAbsent(
                            serializedRecord.topic(),
                            t -> TopicMetrics.producedSensor(
                                Thread.currentThread().getName(),
                                taskId.toString(),
                                processorNodeId,
                                serializedRecord.topic(),
                                context.metrics()
                            )
                        );
                        final long bytesProduced = producerRecordSizeInBytes(serializedRecord);
                        topicProducedSensor.record(
                            bytesProduced,
                            context.currentSystemTimeMs()
                        );
                    }
                } else {
                    recordSendError(
                        serializedRecord.topic(),
                        exception,
                        serializedRecord,
                        context,
                        processorNodeId
                    );

                    // KAFKA-7510 only put message key and value in TRACE level log so we don't leak data by default
                    log.trace("Failed record: (key {} value {} timestamp {}) topic=[{}] partition=[{}]", key, value, serializedRecord.timestamp(), serializedRecord.topic(), serializedRecord.partition());
                }
            } catch (final RuntimeException fatal) {
                sendException.set(new StreamsException("Producer.send `Callback` failed", fatal));
            }
        });
    }

    private static void freeRawInputRecordFromContext(final InternalProcessorContext<Void, Void> context) {
        if (context != null && context.recordContext() != null) {
            context.recordContext().freeRawRecord();
        }
    }

    private <K, V> void handleException(final ProductionExceptionHandler.SerializationExceptionOrigin origin,
                                        final String topic,
                                        final K key,
                                        final V value,
                                        final Headers headers,
                                        final Integer partition,
                                        final Long timestamp,
                                        final String processorNodeId,
                                        final InternalProcessorContext<?, ?> context,
                                        final Exception serializationException) {
        log.debug(String.format("Error serializing record for topic %s", topic), serializationException);

        final ProducerRecord<K, V> record = new ProducerRecord<>(topic, partition, timestamp, key, value, headers);

        final ProductionExceptionHandler.Response response;
        try {
            response = Objects.requireNonNull(
                productionExceptionHandler.handleSerializationError(
                    errorHandlerContext(context, processorNodeId),
                    record,
                    serializationException,
                    origin
                ),
                "Invalid ProductionExceptionHandler response."
            );
        } catch (final Exception fatalUserException) {
            // while Java distinguishes checked vs unchecked exceptions, other languages
            // like Scala or Kotlin do not, and thus we need to catch `Exception`
            // (instead of `RuntimeException`) to work well with those languages
            log.error(
                String.format(
                    "Production error callback failed after serialization error for record %s: %s",
                    origin.toString().toLowerCase(Locale.ROOT),
                    errorHandlerContext(context, processorNodeId)
                ),
                serializationException
            );
            throw new FailedProcessingException(
                "Fatal user code error in production error callback",
                processorNodeId,
                fatalUserException
            );
        }

        final List<ProducerRecord<byte[], byte[]>> deadLetterQueueRecords = response.deadLetterQueueRecords();
        if (!deadLetterQueueRecords.isEmpty()) {
            for (final ProducerRecord<byte[], byte[]> deadLetterQueueRecord : deadLetterQueueRecords) {
                this.send(
                        deadLetterQueueRecord.key(),
                        deadLetterQueueRecord.value(),
                        processorNodeId,
                        context,
                        deadLetterQueueRecord
                );
            }
        }

        if (maybeFailResponse(response.result()) == ProductionExceptionHandler.Result.FAIL) {
            throw new StreamsException(
                String.format(
                    "Unable to serialize record. ProducerRecord(topic=[%s], partition=[%d], timestamp=[%d]",
                    topic,
                    partition,
                    timestamp),
                serializationException
            );
        }

        log.warn("Unable to serialize record, continue processing. " +
                    "ProducerRecord(topic=[{}], partition=[{}], timestamp=[{}])",
                topic,
                partition,
                timestamp);

        droppedRecordsSensor.record();
    }

    private DefaultErrorHandlerContext errorHandlerContext(final InternalProcessorContext<?, ?> context,
                                                           final String processorNodeId) {
        final RecordContext recordContext = context != null ? context.recordContext() : null;

        return recordContext != null ?
            new DefaultErrorHandlerContext(
                context,
                recordContext.topic(),
                recordContext.partition(),
                recordContext.offset(),
                recordContext.headers(),
                processorNodeId,
                taskId,
                recordContext.timestamp(),
                context.recordContext().sourceRawKey(),
                context.recordContext().sourceRawValue()
            ) :
            new DefaultErrorHandlerContext(
                context,
                null,
                -1,
                -1,
                new RecordHeaders(),
                processorNodeId,
                taskId,
                -1L,
                null,
                null
            );
    }

    private <KV> StreamsException createStreamsExceptionForClassCastException(final ProductionExceptionHandler.SerializationExceptionOrigin origin,
                                                                              final String topic,
                                                                              final KV keyOrValue,
                                                                              final Serializer<KV> keyOrValueSerializer,
                                                                              final ClassCastException exception) {
        final String keyOrValueClass = keyOrValue == null
            ? String.format("unknown because %s is null", origin.toString().toLowerCase(Locale.ROOT)) : keyOrValue.getClass().getName();

        return new StreamsException(
            MessageFormat.format(
                String.format(
                        "ClassCastException while producing data to topic %s. " +
                            "The {0} serializer %s is not compatible to the actual {0} type: %s. " +
                            "Change the default {0} serde in StreamConfig or provide the correct {0} serde via method parameters " +
                            "(for example if using the DSL, `#to(String topic, Produced<K, V> produced)` with " +
                            "`Produced.{0}Serde(WindowedSerdes.timeWindowedSerdeFrom(String.class))`).",
                        topic,
                        keyOrValueSerializer.getClass().getName(),
                        keyOrValueClass),
                origin.toString().toLowerCase(Locale.ROOT)),
                exception);
    }

    private void recordSendError(final String topic,
                                 final Exception productionException,
                                 final ProducerRecord<byte[], byte[]> serializedRecord,
                                 final InternalProcessorContext<?, ?> context,
                                 final String processorNodeId) {
        String errorMessage = String.format(SEND_EXCEPTION_MESSAGE, topic, taskId, productionException.toString());

        if (isFatalException(productionException)) {
            errorMessage += "\nWritten offsets would not be recorded and no more records would be sent since this is a fatal error.";
            sendException.set(new StreamsException(errorMessage, productionException));
        } else if (productionException instanceof ProducerFencedException ||
                productionException instanceof InvalidPidMappingException ||
                productionException instanceof InvalidProducerEpochException ||
                productionException instanceof OutOfOrderSequenceException) {
            errorMessage += "\nWritten offsets would not be recorded and no more records would be sent since the producer is fenced, " +
                "indicating the task may be migrated out";
            sendException.set(new TaskMigratedException(errorMessage, productionException));
        } else if (productionException instanceof TransactionAbortedException) {
            // swallow silently
            //
            // TransactionAbortedException is only thrown after `abortTransaction()` was called,
            // so it's only a followup error, and Kafka Streams is already handling the original error
        } else {
            final ProductionExceptionHandler.Response response;
            try {
                response = Objects.requireNonNull(
                    productionExceptionHandler.handleError(
                        errorHandlerContext(context, processorNodeId),
                        serializedRecord,
                        productionException
                    ),
                    "Invalid ProductionExceptionHandler response."
                );
            } catch (final Exception fatalUserException) {
                // while Java distinguishes checked vs unchecked exceptions, other languages
                // like Scala or Kotlin do not, and thus we need to catch `Exception`
                // (instead of `RuntimeException`) to work well with those languages
                log.error(
                    "Production error callback failed after production error for record {}",
                    serializedRecord,
                    productionException
                );
                sendException.set(new FailedProcessingException(
                    "Fatal user code error in production error callback",
                    processorNodeId,
                    fatalUserException
                    )
                );
                return;
            }

            final List<ProducerRecord<byte[], byte[]>> deadLetterQueueRecords = response.deadLetterQueueRecords();
            if (!deadLetterQueueRecords.isEmpty()) {
                for (final ProducerRecord<byte[], byte[]> deadLetterQueueRecord : deadLetterQueueRecords) {
                    this.send(
                            deadLetterQueueRecord.key(),
                            deadLetterQueueRecord.value(),
                            processorNodeId,
                            context,
                            deadLetterQueueRecord
                    );
                }
            }

            if (productionException instanceof RetriableException && response.result() == ProductionExceptionHandler.Result.RETRY) {
                errorMessage += "\nThe broker is either slow or in bad state (like not having enough replicas) in responding the request, " +
                    "or the connection to broker was interrupted sending the request or receiving the response. " +
                    "\nConsider overwriting `max.block.ms` and /or " +
                    "`delivery.timeout.ms` to a larger value to wait longer for such scenarios and avoid timeout errors";
                sendException.set(new TaskCorruptedException(Collections.singleton(taskId)));
            } else {
                if (maybeFailResponse(response.result()) == ProductionExceptionHandler.Result.FAIL) {
                    errorMessage += "\nException handler choose to FAIL the processing, no more records would be sent.";
                    sendException.set(new StreamsException(errorMessage, productionException));
                } else {
                    errorMessage += "\nException handler choose to CONTINUE processing in spite of this error but written offsets would not be recorded.";
                    droppedRecordsSensor.record();
                }
            }
        }

        log.error(errorMessage, productionException);
    }

    private ProductionExceptionHandler.Result maybeFailResponse(final ProductionExceptionHandler.Result result) {
        if (result == ProductionExceptionHandler.Result.RETRY) {
            log.warn("ProductionExceptionHandler returned RETRY for a non-retriable exception. Will treat it as FAIL.");
            return ProductionExceptionHandler.Result.FAIL;
        } else {
            return result;
        }
    }

    private boolean isFatalException(final Exception exception) {
        final boolean securityException = exception instanceof AuthenticationException ||
            exception instanceof AuthorizationException ||
            exception instanceof SecurityDisabledException;

        final boolean communicationException = exception instanceof InvalidTopicException ||
            exception instanceof UnknownServerException ||
            exception instanceof SerializationException ||
            exception instanceof OffsetMetadataTooLarge ||
            exception instanceof IllegalStateException;

        return securityException || communicationException;
    }

    /**
     * @throws StreamsException fatal error that should cause the thread to die
     * @throws TaskMigratedException recoverable error that would cause the task to be removed
     */
    @Override
    public void flush() {
        log.debug("Flushing record collector");
        streamsProducer.flush();
        checkForException();
    }

    /**
     * @throws StreamsException fatal error that should cause the thread to die
     * @throws TaskMigratedException recoverable error that would cause the task to be removed
     */
    @Override
    public void closeClean() {
        log.info("Closing record collector clean");

        removeAllProducedSensors();

        // No need to abort transaction during a clean close: either we have successfully committed the ongoing
        // transaction during handleRevocation and thus there is no transaction in flight, or else none of the revoked
        // tasks had any data in the current transaction and therefore there is no need to commit or abort it.

        close();
    }

    /**
     * @throws StreamsException fatal error that should cause the thread to die
     * @throws TaskMigratedException recoverable error that would cause the task to be removed
     */
    @Override
    public void closeDirty() {
        log.info("Closing record collector dirty");

        if (streamsProducer.eosEnabled()) {
            // We may be closing dirty because the commit failed, so we must abort the transaction to be safe
            streamsProducer.abortTransaction();
        }

        close();
    }

    private void close() {
        offsets.clear();
        checkForException();
    }

    private void removeAllProducedSensors() {
        for (final Sensor sensor : producedSensorByTopic.values()) {
            streamsMetrics.removeSensor(sensor);
        }
    }

    @Override
    public Map<TopicPartition, Long> offsets() {
        return Map.copyOf(offsets);
    }

    private void checkForException() {
        final KafkaException exception = sendException.get();

        if (exception != null) {
            sendException.compareAndSet(exception, null);
            throw exception;
        }
    }

    // for testing only
    Producer<byte[], byte[]> producer() {
        return streamsProducer.kafkaProducer();
    }
}
