/*
 * 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.fluss.client.write;

import org.apache.fluss.annotation.Internal;
import org.apache.fluss.bucketing.BucketingFunction;
import org.apache.fluss.client.admin.Admin;
import org.apache.fluss.client.metadata.MetadataUpdater;
import org.apache.fluss.client.metrics.WriterMetricGroup;
import org.apache.fluss.client.write.RecordAccumulator.RecordAppendResult;
import org.apache.fluss.cluster.Cluster;
import org.apache.fluss.config.ConfigOptions;
import org.apache.fluss.config.Configuration;
import org.apache.fluss.exception.FlussRuntimeException;
import org.apache.fluss.exception.IllegalConfigurationException;
import org.apache.fluss.metadata.PhysicalTablePath;
import org.apache.fluss.metadata.TableInfo;
import org.apache.fluss.rpc.gateway.TabletServerGateway;
import org.apache.fluss.rpc.metrics.ClientMetricGroup;
import org.apache.fluss.utils.CopyOnWriteMap;
import org.apache.fluss.utils.clock.SystemClock;
import org.apache.fluss.utils.concurrent.ExecutorThreadFactory;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.concurrent.ThreadSafe;

import java.time.Duration;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;

import static org.apache.fluss.config.ConfigOptions.NoKeyAssigner.ROUND_ROBIN;
import static org.apache.fluss.config.ConfigOptions.NoKeyAssigner.STICKY;
import static org.apache.fluss.utils.ExceptionUtils.toException;

/**
 * A client that write records to server.
 *
 * <p>The writer consists of a pool of buffer space that holds records that haven't yet been
 * transmitted to the tablet server as well as a background I/O thread that is responsible for
 * turning these records into requests and transmitting them to the cluster. Failure to close the
 * {@link WriterClient} after use will leak these resources.
 *
 * <p>The send method is asynchronous. When called, it adds the log record to a buffer of pending
 * record sends and immediately returns. This allows the wrote record to batch together individual
 * records for efficiency.
 */
@ThreadSafe
@Internal
public class WriterClient {
    private static final Logger LOG = LoggerFactory.getLogger(WriterClient.class);

    public static final String SENDER_THREAD_PREFIX = "fluss-write-sender";
    /**
     * {@link ConfigOptions#CLIENT_WRITER_MAX_INFLIGHT_REQUESTS_PER_BUCKET} should be less than or
     * equal to this value when idempotence producer enabled to ensure message ordering.
     */
    private static final int MAX_IN_FLIGHT_REQUESTS_PER_BUCKET_FOR_IDEMPOTENCE = 5;

    private final Configuration conf;
    private final int maxRequestSize;
    private final RecordAccumulator accumulator;
    private final Sender sender;
    private final ExecutorService ioThreadPool;
    private final MetadataUpdater metadataUpdater;
    private final Map<PhysicalTablePath, BucketAssigner> bucketAssignerMap = new CopyOnWriteMap<>();
    private final IdempotenceManager idempotenceManager;
    private final WriterMetricGroup writerMetricGroup;
    private final DynamicPartitionCreator dynamicPartitionCreator;

    public WriterClient(
            Configuration conf,
            MetadataUpdater metadataUpdater,
            ClientMetricGroup clientMetricGroup,
            Admin admin) {
        int maxRequestSizeLocal = -1;
        IdempotenceManager idempotenceManagerLocal = null;
        try {
            this.conf = conf;
            this.metadataUpdater = metadataUpdater;
            maxRequestSizeLocal =
                    (int) conf.get(ConfigOptions.CLIENT_WRITER_REQUEST_MAX_SIZE).getBytes();
            this.maxRequestSize = maxRequestSizeLocal;
            idempotenceManagerLocal = buildIdempotenceManager();
            this.idempotenceManager = idempotenceManagerLocal;
            this.writerMetricGroup = new WriterMetricGroup(clientMetricGroup);

            short acks = configureAcks(idempotenceManager.idempotenceEnabled());
            int retries = configureRetries(idempotenceManager.idempotenceEnabled());
            this.accumulator =
                    new RecordAccumulator(
                            conf, idempotenceManager, writerMetricGroup, SystemClock.getInstance());
            this.sender = newSender(acks, retries);
            this.ioThreadPool = createThreadPool();
            ioThreadPool.submit(sender);

            this.dynamicPartitionCreator =
                    new DynamicPartitionCreator(
                            metadataUpdater,
                            admin,
                            conf.get(ConfigOptions.CLIENT_WRITER_DYNAMIC_CREATE_PARTITION_ENABLED),
                            this::maybeAbortBatches);
        } catch (Throwable t) {
            LOG.error("Failed to construct writer.", t);
            close(Duration.ofMillis(0));
            throw new FlussRuntimeException(
                    String.format(
                            "Failed to construct writer. Max request size: %d bytes, Idempotence enabled: %b",
                            maxRequestSizeLocal,
                            idempotenceManagerLocal != null
                                    && idempotenceManagerLocal.idempotenceEnabled()),
                    t);
        }
    }

    /**
     * Asynchronously send a record to a table and invoke the provided callback when to send has
     * been acknowledged.
     */
    public void send(WriteRecord record, WriteCallback callback) {
        doSend(record, callback);
    }

    /**
     * Invoking this method makes all buffered records immediately available to send (even if <code>
     * linger.ms</code> is greater than 0) and blocks on the completion of the requests associated
     * with these records. The post-condition of <code>flush()</code> is that any previously sent
     * record will have completed (e.g. <code>Future.isDone() == true</code>). A request is
     * considered completed when it is successfully acknowledged according to the <code>acks</code>
     * configuration you have specified or else it results in an error.
     *
     * <p>Other threads can continue sending records while one thread is blocked waiting for a flush
     * call to complete, however no guarantee is made about the completion of records sent after the
     * flush call begins.
     */
    public void flush() {
        LOG.trace("Flushing accumulated records in writer.");
        long start = System.currentTimeMillis();
        accumulator.beginFlush();
        try {
            accumulator.awaitFlushCompletion();
        } catch (InterruptedException e) {
            throw new FlussRuntimeException(
                    String.format(
                            "Flush interrupted after %d ms. Writer may be in inconsistent state",
                            System.currentTimeMillis() - start),
                    e);
        }
        LOG.trace(
                "Flushed accumulated records in writer in {} ms.",
                System.currentTimeMillis() - start);
    }

    private void doSend(WriteRecord record, WriteCallback callback) {
        try {
            throwIfWriterClosed();

            TableInfo tableInfo = record.getTableInfo();
            PhysicalTablePath physicalTablePath = record.getPhysicalTablePath();
            dynamicPartitionCreator.checkAndCreatePartitionAsync(
                    physicalTablePath, tableInfo.getPartitionKeys());

            // maybe create bucket assigner.
            Cluster cluster = metadataUpdater.getCluster();
            BucketAssigner bucketAssigner =
                    bucketAssignerMap.computeIfAbsent(
                            physicalTablePath,
                            k -> createBucketAssigner(tableInfo, physicalTablePath, conf));

            // Append the record to the accumulator.
            int bucketId = bucketAssigner.assignBucket(record.getBucketKey(), cluster);

            RecordAppendResult result =
                    accumulator.append(
                            record, callback, cluster, bucketId, bucketAssigner.abortIfBatchFull());

            if (result.abortRecordForNewBatch) {
                int prevBucketId = bucketId;
                bucketAssigner.onNewBatch(cluster, prevBucketId);
                bucketId = bucketAssigner.assignBucket(record.getBucketKey(), cluster);
                LOG.trace(
                        "Retrying append due to new batch creation for table {} bucket {}, the old bucket was {}.",
                        physicalTablePath,
                        bucketId,
                        prevBucketId);
                result = accumulator.append(record, callback, cluster, bucketId, false);
            }

            if (result.batchIsFull || result.newBatchCreated) {
                LOG.trace(
                        "Waking up the sender since table {} bucket {} is either full or getting a new batch",
                        record.getPhysicalTablePath(),
                        bucketId);
                // TODO add the wakeup logic refer to Kafka.
            }
        } catch (Exception e) {
            throw new FlussRuntimeException(
                    String.format(
                            "Failed to send record to table %s. Writer state: %s",
                            record.getPhysicalTablePath(),
                            sender != null && sender.isRunning() ? "running" : "closed"),
                    e);
        }
    }

    private void maybeAbortBatches(Throwable t) {
        if (accumulator.hasIncomplete()) {
            LOG.error("Aborting all pending write batches due to fatal error", t);
            accumulator.abortBatches(toException(t));
        }
    }

    // Verify that writer instance has not been closed. This method throws IllegalStateException if
    // writer has already been closed.
    private void throwIfWriterClosed() {
        if (sender == null || !sender.isRunning()) {
            throw new IllegalStateException(
                    String.format(
                            "Cannot perform write operation after writer has been closed. Sender running: %b, Thread pool shutdown: %b",
                            sender != null && sender.isRunning(),
                            ioThreadPool == null || ioThreadPool.isShutdown()));
        }
    }

    private IdempotenceManager buildIdempotenceManager() {
        boolean idempotenceEnabled =
                conf.getBoolean(ConfigOptions.CLIENT_WRITER_ENABLE_IDEMPOTENCE);
        int maxInflightRequestPerBucket =
                conf.getInt(ConfigOptions.CLIENT_WRITER_MAX_INFLIGHT_REQUESTS_PER_BUCKET);
        if (idempotenceEnabled
                && maxInflightRequestPerBucket
                        > MAX_IN_FLIGHT_REQUESTS_PER_BUCKET_FOR_IDEMPOTENCE) {
            throw new IllegalConfigurationException(
                    String.format(
                            "Invalid configuration for idempotent writer. The value of %s (%d) should be less than or equal to %d when idempotence is enabled to ensure message ordering",
                            ConfigOptions.CLIENT_WRITER_MAX_INFLIGHT_REQUESTS_PER_BUCKET.key(),
                            maxInflightRequestPerBucket,
                            MAX_IN_FLIGHT_REQUESTS_PER_BUCKET_FOR_IDEMPOTENCE));
        }

        TabletServerGateway tabletServerGateway = metadataUpdater.newRandomTabletServerClient();
        return idempotenceEnabled
                ? new IdempotenceManager(
                        true, maxInflightRequestPerBucket, tabletServerGateway, metadataUpdater)
                : new IdempotenceManager(
                        false, maxInflightRequestPerBucket, tabletServerGateway, metadataUpdater);
    }

    private short configureAcks(boolean idempotenceEnabled) {
        String acks = conf.get(ConfigOptions.CLIENT_WRITER_ACKS);
        short ack;
        if (acks.equals("all")) {
            ack = Short.parseShort("-1");
        } else {
            ack = Short.parseShort(acks);
        }

        if (idempotenceEnabled && ack != -1) {
            throw new IllegalConfigurationException(
                    String.format(
                            "Invalid acks configuration for idempotent writer. Must set %s to 'all' (current value: '%s') in order to use the idempotent writer. Otherwise we cannot guarantee idempotence",
                            ConfigOptions.CLIENT_WRITER_ACKS.key(), acks));
        }

        return ack;
    }

    private int configureRetries(boolean idempotenceEnabled) {
        int retries = conf.getInt(ConfigOptions.CLIENT_WRITER_RETRIES);
        if (idempotenceEnabled && retries == 0) {
            throw new IllegalConfigurationException(
                    String.format(
                            "Invalid retries configuration for idempotent writer. Must set %s to non-zero (current value: %d) when using the idempotent writer. Otherwise we cannot guarantee idempotence",
                            ConfigOptions.CLIENT_WRITER_RETRIES.key(), retries));
        }
        return retries;
    }

    private Sender newSender(short acks, int retries) {
        return new Sender(
                accumulator,
                (int) conf.get(ConfigOptions.CLIENT_REQUEST_TIMEOUT).toMillis(),
                maxRequestSize,
                acks,
                retries,
                metadataUpdater,
                idempotenceManager,
                writerMetricGroup);
    }

    public void close(Duration timeout) {
        LOG.info("Closing writer.");

        writerMetricGroup.close();

        if (sender != null) {
            sender.initiateClose();
        }

        if (ioThreadPool != null) {
            ioThreadPool.shutdown();

            try {
                if (!ioThreadPool.awaitTermination(timeout.toMillis(), TimeUnit.MILLISECONDS)) {
                    ioThreadPool.shutdownNow();

                    if (!ioThreadPool.awaitTermination(timeout.toMillis(), TimeUnit.MILLISECONDS)) {
                        LOG.error("Failed to shutdown writer.");
                    }
                }
            } catch (InterruptedException e) {
                ioThreadPool.shutdownNow();
                Thread.currentThread().interrupt();
            }
        }

        if (sender != null) {
            sender.forceClose();
        }

        LOG.info("Writer closed.");
    }

    private ExecutorService createThreadPool() {
        return Executors.newFixedThreadPool(1, new ExecutorThreadFactory(SENDER_THREAD_PREFIX));
    }

    private BucketAssigner createBucketAssigner(
            TableInfo tableInfo, PhysicalTablePath physicalTablePath, Configuration conf) {
        int bucketNumber = tableInfo.getNumBuckets();
        List<String> bucketKeys = tableInfo.getBucketKeys();
        if (!bucketKeys.isEmpty()) {
            BucketingFunction function =
                    BucketingFunction.of(
                            tableInfo.getTableConfig().getDataLakeFormat().orElse(null));
            return new HashBucketAssigner(bucketNumber, function);
        } else {
            ConfigOptions.NoKeyAssigner noKeyAssigner =
                    conf.get(ConfigOptions.CLIENT_WRITER_BUCKET_NO_KEY_ASSIGNER);
            if (noKeyAssigner == ROUND_ROBIN) {
                return new RoundRobinBucketAssigner(physicalTablePath, bucketNumber);
            } else if (noKeyAssigner == STICKY) {
                return new StickyBucketAssigner(physicalTablePath, bucketNumber);
            } else {
                throw new IllegalArgumentException(
                        "Unsupported append only row bucket assigner: " + noKeyAssigner);
            }
        }
    }
}
