/*
 * 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.table.scanner.log;

import org.apache.fluss.annotation.Internal;
import org.apache.fluss.annotation.VisibleForTesting;
import org.apache.fluss.client.metrics.ScannerMetricGroup;
import org.apache.fluss.client.table.scanner.RemoteFileDownloader;
import org.apache.fluss.config.ConfigOptions;
import org.apache.fluss.config.Configuration;
import org.apache.fluss.fs.FsPath;
import org.apache.fluss.fs.FsPathAndFileName;
import org.apache.fluss.metadata.TablePath;
import org.apache.fluss.remote.RemoteLogSegment;
import org.apache.fluss.utils.ExceptionUtils;
import org.apache.fluss.utils.FlussPaths;
import org.apache.fluss.utils.concurrent.ShutdownableThread;

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

import javax.annotation.concurrent.ThreadSafe;

import java.io.Closeable;
import java.io.File;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.UUID;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.PriorityBlockingQueue;
import java.util.concurrent.Semaphore;
import java.util.concurrent.TimeUnit;

import static org.apache.fluss.utils.FileUtils.deleteDirectoryQuietly;
import static org.apache.fluss.utils.FlussPaths.LOG_FILE_SUFFIX;
import static org.apache.fluss.utils.FlussPaths.remoteLogSegmentDir;
import static org.apache.fluss.utils.FlussPaths.remoteLogSegmentFile;

/** Downloader to read remote log files to local disk. */
@ThreadSafe
@Internal
public class RemoteLogDownloader implements Closeable {
    private static final Logger LOG = LoggerFactory.getLogger(RemoteLogDownloader.class);

    private static final long POLL_TIMEOUT = 5000L;

    private final Path localLogDir;

    /**
     * A queue to hold the remote log segment files to be fetched. The queue is ordered by the
     * max_timestamp of the remote log segment. So we download the remote log segments from the
     * older to the newer.
     */
    private final PriorityBlockingQueue<RemoteLogDownloadRequest> segmentsToFetch;

    private final BlockingQueue<RemoteLogSegment> segmentsToRecycle;

    private final Semaphore prefetchSemaphore;

    private final DownloadRemoteLogThread downloadThread;

    private final RemoteFileDownloader remoteFileDownloader;

    private final ScannerMetricGroup scannerMetricGroup;

    private final long pollTimeout;

    public RemoteLogDownloader(
            TablePath tablePath,
            Configuration conf,
            RemoteFileDownloader remoteFileDownloader,
            ScannerMetricGroup scannerMetricGroup) {
        // default we give a 5s long interval to avoid frequent loop
        this(tablePath, conf, remoteFileDownloader, scannerMetricGroup, POLL_TIMEOUT);
    }

    @VisibleForTesting
    RemoteLogDownloader(
            TablePath tablePath,
            Configuration conf,
            RemoteFileDownloader remoteFileDownloader,
            ScannerMetricGroup scannerMetricGroup,
            long pollTimeout) {
        this.segmentsToFetch = new PriorityBlockingQueue<>();
        this.segmentsToRecycle = new LinkedBlockingQueue<>();
        this.remoteFileDownloader = remoteFileDownloader;
        this.scannerMetricGroup = scannerMetricGroup;
        this.pollTimeout = pollTimeout;
        this.prefetchSemaphore =
                new Semaphore(conf.getInt(ConfigOptions.CLIENT_SCANNER_REMOTE_LOG_PREFETCH_NUM));
        // The local tmp dir to store the fetched log segment files,
        // add UUID to avoid conflict between tasks.
        this.localLogDir =
                Paths.get(
                        conf.get(ConfigOptions.CLIENT_SCANNER_IO_TMP_DIR),
                        "remote-logs-" + UUID.randomUUID());
        this.downloadThread = new DownloadRemoteLogThread(tablePath);
    }

    public void start() {
        downloadThread.start();
    }

    /** Request to fetch remote log segment to local. This method is non-blocking. */
    public RemoteLogDownloadFuture requestRemoteLog(FsPath logTabletDir, RemoteLogSegment segment) {
        RemoteLogDownloadRequest request = new RemoteLogDownloadRequest(segment, logTabletDir);
        segmentsToFetch.add(request);
        return new RemoteLogDownloadFuture(request.future, () -> recycleRemoteLog(segment));
    }

    /**
     * Recycle the consumed remote log. The removal of the log file is async in the {@link
     * #downloadThread}.
     */
    void recycleRemoteLog(RemoteLogSegment segment) {
        segmentsToRecycle.add(segment);
        prefetchSemaphore.release();
    }

    /**
     * Fetch a remote log segment file to local. This method will block until there is a log segment
     * to fetch.
     */
    void fetchOnce() throws Exception {
        // blocks until there is capacity (the fetched file is consumed)
        prefetchSemaphore.acquire();

        // wait until there is a remote fetch request
        RemoteLogDownloadRequest request = segmentsToFetch.poll(pollTimeout, TimeUnit.MILLISECONDS);
        if (request == null) {
            prefetchSemaphore.release();
            return;
        }

        try {
            // 1. cleanup the finished logs first to free up disk space
            cleanupRemoteLogs();

            // 2. do the actual download work
            FsPathAndFileName fsPathAndFileName = request.getFsPathAndFileName();
            scannerMetricGroup.remoteFetchRequestCount().inc();

            long startTime = System.currentTimeMillis();
            // download the remote file to local
            remoteFileDownloader
                    .downloadFileAsync(fsPathAndFileName, localLogDir)
                    .whenComplete(
                            (bytes, throwable) -> {
                                if (throwable != null) {
                                    LOG.error(
                                            "Failed to download remote log segment file {}.",
                                            fsPathAndFileName.getFileName(),
                                            ExceptionUtils.stripExecutionException(throwable));
                                    // release the semaphore for the failed request
                                    prefetchSemaphore.release();
                                    // add back the request to the queue,
                                    // so we do not complete the request.future here
                                    segmentsToFetch.add(request);
                                    scannerMetricGroup.remoteFetchErrorCount().inc();
                                } else {
                                    LOG.info(
                                            "Successfully downloaded remote log segment file {} to local cost {} ms.",
                                            fsPathAndFileName.getFileName(),
                                            System.currentTimeMillis() - startTime);
                                    File localFile =
                                            new File(
                                                    localLogDir.toFile(),
                                                    fsPathAndFileName.getFileName());
                                    scannerMetricGroup.remoteFetchBytes().inc(bytes);
                                    request.future.complete(localFile);
                                }
                            });
        } catch (Throwable t) {
            prefetchSemaphore.release();
            // add back the request to the queue
            segmentsToFetch.add(request);
            scannerMetricGroup.remoteFetchErrorCount().inc();
            // log the error and continue instead of shutdown the download thread
            LOG.error("Failed to download remote log segment.", t);
        }
    }

    private void cleanupRemoteLogs() {
        RemoteLogSegment segment;
        while ((segment = segmentsToRecycle.poll()) != null) {
            cleanupFinishedRemoteLog(segment);
        }
    }

    private void cleanupFinishedRemoteLog(RemoteLogSegment segment) {
        try {
            Path logFile = localLogDir.resolve(getLocalFileNameOfRemoteSegment(segment));
            Files.deleteIfExists(logFile);
            LOG.info(
                    "Consumed and deleted the fetched log segment file {} for bucket {}.",
                    logFile.getFileName(),
                    segment.tableBucket());
        } catch (IOException e) {
            LOG.warn("Failed to delete the local fetch segment file {}.", localLogDir, e);
        }
    }

    @Override
    public void close() throws IOException {
        try {
            downloadThread.shutdown();
        } catch (InterruptedException e) {
            // ignore
        }

        deleteDirectoryQuietly(localLogDir.toFile());
    }

    @VisibleForTesting
    Semaphore getPrefetchSemaphore() {
        return prefetchSemaphore;
    }

    @VisibleForTesting
    Path getLocalLogDir() {
        return localLogDir;
    }

    @VisibleForTesting
    int getSizeOfSegmentsToFetch() {
        return segmentsToFetch.size();
    }

    protected static FsPathAndFileName getFsPathAndFileName(
            FsPath remoteLogTabletDir, RemoteLogSegment segment) {
        FsPath remotePath =
                remoteLogSegmentFile(
                        remoteLogSegmentDir(remoteLogTabletDir, segment.remoteLogSegmentId()),
                        segment.remoteLogStartOffset());
        return new FsPathAndFileName(remotePath, getLocalFileNameOfRemoteSegment(segment));
    }

    /**
     * Get the local file name of the remote log segment.
     *
     * <p>The file name is in pattern:
     *
     * <pre>
     *     ${remote_segment_id}_${offset_prefix}.log
     * </pre>
     */
    private static String getLocalFileNameOfRemoteSegment(RemoteLogSegment segment) {
        return segment.remoteLogSegmentId()
                + "_"
                + FlussPaths.filenamePrefixFromOffset(segment.remoteLogStartOffset())
                + LOG_FILE_SUFFIX;
    }

    /**
     * Thread to download remote log files to local. The thread will keep fetching remote log files
     * until it is interrupted.
     */
    private class DownloadRemoteLogThread extends ShutdownableThread {
        public DownloadRemoteLogThread(TablePath tablePath) {
            super(String.format("DownloadRemoteLog-[%s]", tablePath.toString()), true);
        }

        @Override
        public void doWork() throws Exception {
            fetchOnce();
            cleanupRemoteLogs();
        }
    }

    /** Represents a request to download a remote log segment file to local. */
    static class RemoteLogDownloadRequest implements Comparable<RemoteLogDownloadRequest> {
        final RemoteLogSegment segment;
        final FsPath remoteLogTabletDir;
        final CompletableFuture<File> future = new CompletableFuture<>();

        public RemoteLogDownloadRequest(RemoteLogSegment segment, FsPath remoteLogTabletDir) {
            this.segment = segment;
            this.remoteLogTabletDir = remoteLogTabletDir;
        }

        public FsPathAndFileName getFsPathAndFileName() {
            return RemoteLogDownloader.getFsPathAndFileName(remoteLogTabletDir, segment);
        }

        @Override
        public int compareTo(RemoteLogDownloadRequest o) {
            if (segment.tableBucket().equals(o.segment.tableBucket())) {
                // strictly download in the offset order if they belong to the same bucket
                return Long.compare(
                        segment.remoteLogStartOffset(), o.segment.remoteLogStartOffset());
            } else {
                // download segment from old to new across buckets
                return Long.compare(segment.maxTimestamp(), o.segment.maxTimestamp());
            }
        }
    }
}
