fluss-rs 0.1.0

// 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.

use crate::client::credentials::CredentialsReceiver;
use crate::error::{Error, Result};
use crate::io::{FileIO, Storage};
use crate::metadata::TableBucket;
use crate::proto::{PbRemoteLogFetchInfo, PbRemoteLogSegment};
use futures::TryStreamExt;
use parking_lot::Mutex;
use std::{
    cmp::{Ordering, Reverse},
    collections::{BinaryHeap, HashMap},
    future::Future,
    io, mem,
    path::{Path, PathBuf},
    pin::Pin,
    sync::Arc,
    time::Duration,
};

#[cfg(test)]
use std::{
    env,
    time::{SystemTime, UNIX_EPOCH},
};
use tempfile::TempDir;
use tokio::io::AsyncWriteExt;
use tokio::sync::{Notify, OwnedSemaphorePermit, Semaphore, mpsc, oneshot};
use tokio::task::JoinSet;

/// Default maximum number of remote log segments to prefetch
/// Matches Java's CLIENT_SCANNER_REMOTE_LOG_PREFETCH_NUM (default: 4)
pub const DEFAULT_SCANNER_REMOTE_LOG_PREFETCH_NUM: usize = 4;

/// Default maximum concurrent remote log downloads
/// Matches Java's REMOTE_FILE_DOWNLOAD_THREAD_NUM (default: 3)
pub const DEFAULT_REMOTE_FILE_DOWNLOAD_THREAD_NUM: usize = 3;

/// Initial retry backoff delay (milliseconds)
/// Prevents hot-spin retry loops on persistent failures
const RETRY_BACKOFF_BASE_MS: u64 = 100;

/// Maximum retry backoff delay (milliseconds)
/// Caps exponential backoff to avoid excessive delays
const RETRY_BACKOFF_MAX_MS: u64 = 5_000;

/// Maximum number of retries before giving up
/// After this many retries, the download will fail permanently
const MAX_RETRY_COUNT: u32 = 10;

/// Calculate exponential backoff delay with jitter for retries
fn calculate_backoff_delay(retry_count: u32) -> tokio::time::Duration {
    use rand::Rng;

    // Exponential backoff: base * 2^retry_count
    let exponential_ms = RETRY_BACKOFF_BASE_MS.saturating_mul(1 << retry_count.min(10)); // Cap exponent to prevent overflow

    // Cap at maximum
    let capped_ms = exponential_ms.min(RETRY_BACKOFF_MAX_MS);

    // Add jitter (±25% randomness) to avoid thundering herd
    let mut rng = rand::rng();
    let jitter = rng.random_range(0.75..=1.25);
    let final_ms = ((capped_ms as f64) * jitter) as u64;

    tokio::time::Duration::from_millis(final_ms)
}

/// Result of a fetch operation containing file path and size
#[derive(Debug)]
pub struct FetchResult {
    pub file_path: PathBuf,
    pub file_size: usize,
}

/// Trait for fetching remote log segments (allows dependency injection for testing)
pub trait RemoteLogFetcher: Send + Sync {
    fn fetch(
        &self,
        request: &RemoteLogDownloadRequest,
    ) -> Pin<Box<dyn Future<Output = Result<FetchResult>> + Send>>;
}

/// Represents a remote log segment that needs to be downloaded
#[derive(Debug, Clone)]
pub struct RemoteLogSegment {
    pub segment_id: String,
    pub start_offset: i64,
    #[allow(dead_code)]
    pub end_offset: i64,
    #[allow(dead_code)]
    pub size_in_bytes: i32,
    pub table_bucket: TableBucket,
    pub max_timestamp: i64,
}

impl RemoteLogSegment {
    pub fn from_proto(segment: &PbRemoteLogSegment, table_bucket: TableBucket) -> Self {
        Self {
            segment_id: segment.remote_log_segment_id.clone(),
            start_offset: segment.remote_log_start_offset,
            end_offset: segment.remote_log_end_offset,
            size_in_bytes: segment.segment_size_in_bytes,
            table_bucket,
            // Match Java's behavior: use -1 for missing timestamp
            // (Java: CommonRpcMessageUtils.java:171-174)
            max_timestamp: segment.max_timestamp.unwrap_or(-1),
        }
    }

    /// Get the local file name for this remote log segment
    pub fn local_file_name(&self) -> String {
        // Format: ${remote_segment_id}_${offset_prefix}.log
        let offset_prefix = format!("{:020}", self.start_offset);
        format!("{}_{}.log", self.segment_id, offset_prefix)
    }
}

/// Represents remote log fetch information
#[derive(Debug, Clone)]
pub struct RemoteLogFetchInfo {
    pub remote_log_tablet_dir: String,
    #[allow(dead_code)]
    pub partition_name: Option<String>,
    pub remote_log_segments: Vec<RemoteLogSegment>,
    pub first_start_pos: i32,
}

impl RemoteLogFetchInfo {
    pub fn from_proto(info: &PbRemoteLogFetchInfo, table_bucket: TableBucket) -> Self {
        let segments = info
            .remote_log_segments
            .iter()
            .map(|s| RemoteLogSegment::from_proto(s, table_bucket.clone()))
            .collect();

        Self {
            remote_log_tablet_dir: info.remote_log_tablet_dir.clone(),
            partition_name: info.partition_name.clone(),
            remote_log_segments: segments,
            first_start_pos: info.first_start_pos.unwrap_or(0),
        }
    }
}

/// RAII guard for prefetch permit that notifies coordinator on drop
///
/// NOTE: File deletion is now handled by FileSource::drop(), not here.
/// This ensures the file is closed before deletion
#[derive(Debug)]
pub struct PrefetchPermit {
    permit: Option<OwnedSemaphorePermit>,
    recycle_notify: Arc<Notify>,
}

impl PrefetchPermit {
    fn new(permit: OwnedSemaphorePermit, recycle_notify: Arc<Notify>) -> Self {
        Self {
            permit: Some(permit),
            recycle_notify,
        }
    }
}

impl Drop for PrefetchPermit {
    fn drop(&mut self) {
        // Release capacity (critical: permit must be dropped before notify)
        let _ = self.permit.take(); // drops permit here

        // Then wake coordinator so it can acquire the now-available permit
        self.recycle_notify.notify_one();
    }
}

/// Downloaded remote log file with prefetch permit
/// File remains on disk for memory efficiency; file deletion is handled by FileCleanupGuard in FileSource
#[derive(Debug)]
pub struct RemoteLogFile {
    /// Path to the downloaded file on local disk
    pub file_path: PathBuf,
    /// Size of the file in bytes
    /// Currently unused but kept for potential future use (logging, metrics, etc.)
    #[allow(dead_code)]
    pub file_size: usize,
    /// RAII permit that releases prefetch semaphore slot and notifies coordinator when dropped
    pub permit: PrefetchPermit,
}

/// Represents a request to download a remote log segment with priority ordering
#[derive(Debug)]
pub struct RemoteLogDownloadRequest {
    segment: RemoteLogSegment,
    remote_log_tablet_dir: String,
    result_sender: oneshot::Sender<Result<RemoteLogFile>>,
    retry_count: u32,
    next_retry_at: Option<tokio::time::Instant>,
}

impl RemoteLogDownloadRequest {
    /// Get the segment (used by test fetcher implementations)
    #[cfg(test)]
    pub fn segment(&self) -> &RemoteLogSegment {
        &self.segment
    }
}

// Total ordering for priority queue (Rust requirement: cmp==Equal implies Eq)
// Primary: Java semantics (timestamp cross-bucket, offset within-bucket)
// Tie-breakers: table_bucket fields (table_id, partition_id, bucket_id), then segment_id
impl Ord for RemoteLogDownloadRequest {
    fn cmp(&self, other: &Self) -> Ordering {
        if self.segment.table_bucket == other.segment.table_bucket {
            // Same bucket: order by start_offset (ascending - earlier segments first)
            self.segment
                .start_offset
                .cmp(&other.segment.start_offset)
                .then_with(|| self.segment.segment_id.cmp(&other.segment.segment_id))
        } else {
            // Different buckets: order by max_timestamp (ascending - older segments first)
            // Then by table_bucket fields for true total ordering
            self.segment
                .max_timestamp
                .cmp(&other.segment.max_timestamp)
                .then_with(|| {
                    self.segment
                        .table_bucket
                        .table_id()
                        .cmp(&other.segment.table_bucket.table_id())
                })
                .then_with(|| {
                    self.segment
                        .table_bucket
                        .partition_id()
                        .cmp(&other.segment.table_bucket.partition_id())
                })
                .then_with(|| {
                    self.segment
                        .table_bucket
                        .bucket_id()
                        .cmp(&other.segment.table_bucket.bucket_id())
                })
                .then_with(|| self.segment.segment_id.cmp(&other.segment.segment_id))
        }
    }
}

impl PartialOrd for RemoteLogDownloadRequest {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl PartialEq for RemoteLogDownloadRequest {
    fn eq(&self, other: &Self) -> bool {
        self.cmp(other) == Ordering::Equal
    }
}

impl Eq for RemoteLogDownloadRequest {}

/// Result of a download task
enum DownloadResult {
    /// Successful download - deliver result to future
    Success {
        result: RemoteLogFile,
        result_sender: oneshot::Sender<Result<RemoteLogFile>>,
    },
    /// Download failed - re-queue request for retry (Java pattern)
    FailedRetry { request: RemoteLogDownloadRequest },
    /// Download failed permanently after max retries - fail the future
    FailedPermanently {
        error: Error,
        result_sender: oneshot::Sender<Result<RemoteLogFile>>,
    },
    /// Cancelled - don't deliver, don't re-queue
    Cancelled,
}

/// Production implementation of RemoteLogFetcher that downloads from actual storage
struct ProductionFetcher {
    credentials_rx: CredentialsReceiver,
    local_log_dir: Arc<TempDir>,
    remote_log_read_concurrency: usize,
}

impl RemoteLogFetcher for ProductionFetcher {
    fn fetch(
        &self,
        request: &RemoteLogDownloadRequest,
    ) -> Pin<Box<dyn Future<Output = Result<FetchResult>> + Send>> {
        let mut credentials_rx = self.credentials_rx.clone();
        let local_log_dir = self.local_log_dir.clone();
        let remote_log_read_concurrency = self.remote_log_read_concurrency;

        // Clone data needed for async operation to avoid lifetime issues
        let segment = request.segment.clone();
        let remote_log_tablet_dir = request.remote_log_tablet_dir.to_string();

        Box::pin(async move {
            let local_file_name = segment.local_file_name();
            let local_file_path = local_log_dir.path().join(&local_file_name);

            // Build remote path
            let offset_prefix = format!("{:020}", segment.start_offset);
            let remote_path = format!(
                "{}/{}/{}.log",
                remote_log_tablet_dir, segment.segment_id, offset_prefix
            );

            // Get credentials from watch channel, waiting if not yet fetched
            // - None = not yet fetched, wait
            // - Some(props) = fetched (may be empty if no auth needed)
            let remote_fs_props = {
                let maybe_props = credentials_rx.borrow().clone();
                match maybe_props {
                    Some(props) => props,
                    None => {
                        // Credentials not yet fetched, wait for first update
                        log::info!("Waiting for credentials to be available...");
                        // If the sender side has been dropped (e.g. during shutdown),
                        // this will return an error. Surface that as a proper error
                        // instead of silently falling back to empty credentials.
                        if let Err(e) = credentials_rx.changed().await {
                            let io_err = io::Error::new(
                                io::ErrorKind::BrokenPipe,
                                format!(
                                    "credentials manager shut down before credentials were obtained: {e}"
                                ),
                            );
                            return Err(io_err.into());
                        }
                        // After a successful change notification, credentials should be set.
                        // If they are still missing, treat this as an error instead of
                        // defaulting to an empty map (which could break auth flows).
                        credentials_rx
                            .borrow()
                            .clone()
                            .ok_or_else(|| Error::UnexpectedError {
                                message: "credentials not available after watch notification"
                                    .to_string(),
                                source: None,
                            })?
                    }
                }
            };

            // Download file to disk (streaming, no memory spike)
            let file_path = RemoteLogDownloader::download_file(
                &remote_log_tablet_dir,
                &remote_path,
                &local_file_path,
                &remote_fs_props,
                remote_log_read_concurrency,
            )
            .await?;

            // Get file size
            let metadata = tokio::fs::metadata(&file_path).await?;
            let file_size = metadata.len() as usize;

            // Return file path - file stays on disk until PrefetchPermit is dropped
            Ok(FetchResult {
                file_path,
                file_size,
            })
        })
    }
}

/// Coordinator that owns all download state and orchestrates downloads
struct DownloadCoordinator {
    download_queue: BinaryHeap<Reverse<RemoteLogDownloadRequest>>,
    active_downloads: JoinSet<DownloadResult>,
    in_flight: usize,
    prefetch_semaphore: Arc<Semaphore>,
    max_concurrent_downloads: usize,
    recycle_notify: Arc<Notify>,
    fetcher: Arc<dyn RemoteLogFetcher>,
}

impl DownloadCoordinator {
    /// Check if we should wait for recycle notification
    /// Only wait if we're blocked on permits AND have pending work
    fn should_wait_for_recycle(&self) -> bool {
        !self.download_queue.is_empty()
            && self.in_flight < self.max_concurrent_downloads
            && self.prefetch_semaphore.available_permits() == 0
    }

    /// Find the earliest retry deadline among pending requests
    fn next_retry_deadline(&self) -> Option<tokio::time::Instant> {
        self.download_queue
            .iter()
            .filter_map(|Reverse(req)| req.next_retry_at)
            .min()
    }
}

impl DownloadCoordinator {
    /// Try to start as many downloads as possible (event-driven drain)
    fn drain(&mut self) {
        // Collect deferred requests (backoff not ready) to push back later
        let mut deferred = Vec::new();
        // Scan entire queue once to find ready requests (prevents head-of-line blocking)
        // Bound to reasonable max to avoid excessive work if queue is huge
        let max_scan = self.download_queue.len().min(100);
        let mut scanned = 0;

        while !self.download_queue.is_empty()
            && self.in_flight < self.max_concurrent_downloads
            && scanned < max_scan
        {
            // Try acquire prefetch permit (non-blocking)
            let permit = match self.prefetch_semaphore.clone().try_acquire_owned() {
                Ok(p) => p,
                Err(_) => break, // No permits available
            };

            // Pop highest priority request
            let Some(Reverse(request)) = self.download_queue.pop() else {
                drop(permit);
                break;
            };

            scanned += 1;

            // Retry backoff check: defer if retry time hasn't arrived yet
            if let Some(next_retry_at) = request.next_retry_at {
                let now = tokio::time::Instant::now();
                if next_retry_at > now {
                    // Not ready for retry yet - defer and continue looking for ready requests
                    drop(permit);
                    deferred.push(request);
                    continue; // Don't block - keep looking for ready requests
                }
            }

            // Cancellation check: skip if sender closed
            if request.result_sender.is_closed() {
                drop(permit);
                continue; // Try next request
            }

            // Clone data for the spawned task
            let fetcher = self.fetcher.clone();
            let recycle_notify = self.recycle_notify.clone();

            // Spawn download task
            self.active_downloads.spawn(async move {
                spawn_download_task(request, permit, fetcher, recycle_notify).await
            });
            self.in_flight += 1;
        }

        // Push deferred requests back to queue (maintains priority order)
        if !deferred.is_empty() {
            for req in deferred {
                self.download_queue.push(Reverse(req));
            }
        }
    }
}

/// Spawn a download task that attempts download once
/// Matches Java's RemoteLogDownloader.java
///
/// Benefits over infinite in-place retry:
/// - Failed downloads don't block prefetch slots
/// - Other segments can make progress while one is failing
/// - Natural retry through coordinator re-picking from queue
async fn spawn_download_task(
    request: RemoteLogDownloadRequest,
    permit: tokio::sync::OwnedSemaphorePermit,
    fetcher: Arc<dyn RemoteLogFetcher>,
    recycle_notify: Arc<Notify>,
) -> DownloadResult {
    // Check if receiver still alive (early cancellation check)
    if request.result_sender.is_closed() {
        drop(permit);
        return DownloadResult::Cancelled;
    }

    // Try download ONCE
    let download_result = fetcher.fetch(&request).await;

    match download_result {
        Ok(fetch_result) => {
            // Success - permit will be released on drop (FileSource handles file deletion)
            DownloadResult::Success {
                result: RemoteLogFile {
                    file_path: fetch_result.file_path,
                    file_size: fetch_result.file_size,
                    permit: PrefetchPermit::new(permit, recycle_notify.clone()),
                },
                result_sender: request.result_sender,
            }
        }
        Err(_e) if request.result_sender.is_closed() => {
            // Receiver dropped (cancelled) - release permit, don't re-queue
            drop(permit);
            DownloadResult::Cancelled
        }
        Err(e) => {
            // Download failed - check if we should retry or give up
            let retry_count = request.retry_count + 1;

            if retry_count > MAX_RETRY_COUNT {
                // Too many retries - give up and fail the future
                log::error!(
                    "Failed to download remote log segment {} after {} retries: {}. Giving up.",
                    request.segment.segment_id,
                    retry_count,
                    e
                );
                drop(permit); // Release immediately

                DownloadResult::FailedPermanently {
                    error: Error::UnexpectedError {
                        message: format!(
                            "Failed to download remote log segment after {retry_count} retries: {e}"
                        ),
                        source: Some(Box::new(e)),
                    },
                    result_sender: request.result_sender,
                }
            } else {
                // Retry with exponential backoff
                let backoff_delay = calculate_backoff_delay(retry_count);
                let next_retry_at = tokio::time::Instant::now() + backoff_delay;

                log::warn!(
                    "Failed to download remote log segment {}: {}. Retry {}/{} after {:?}",
                    request.segment.segment_id,
                    e,
                    retry_count,
                    MAX_RETRY_COUNT,
                    backoff_delay
                );
                drop(permit); // Release immediately - critical!

                // Update retry state
                let mut retry_request = request;
                retry_request.retry_count = retry_count;
                retry_request.next_retry_at = Some(next_retry_at);

                // Re-queue request to same priority queue
                // Future stays with request, NOT completed - will complete on successful retry
                DownloadResult::FailedRetry {
                    request: retry_request,
                }
            }
        }
    }
}

/// Coordinator event loop - owns all download state and reacts to events
async fn coordinator_loop(
    mut coordinator: DownloadCoordinator,
    mut request_receiver: mpsc::UnboundedReceiver<RemoteLogDownloadRequest>,
) {
    loop {
        // Drain once at start of iteration to process ready work
        coordinator.drain();

        // Calculate sleep duration until next retry (if any deferred requests)
        let next_retry_sleep = coordinator.next_retry_deadline().map(|deadline| {
            let now = tokio::time::Instant::now();
            if deadline > now {
                deadline - now
            } else {
                tokio::time::Duration::from_millis(0) // Ready now
            }
        });

        tokio::select! {
            // Event 1: NewRequest
            Some(request) = request_receiver.recv() => {
                coordinator.download_queue.push(Reverse(request));
                // Immediately try to start this download
                continue;
            }

            // Event 2: DownloadFinished
            Some(result) = coordinator.active_downloads.join_next() => {
                coordinator.in_flight -= 1;

                match result {
                    Ok(DownloadResult::Success { result, result_sender }) => {
                        // Success - deliver result to future
                        if !result_sender.is_closed() {
                            let _ = result_sender.send(Ok(result));
                        }
                        // Permit held in RemoteLogFile until consumed
                    }
                    Ok(DownloadResult::FailedRetry { request }) => {
                        // Re-queue immediately (don't block coordinator with sleep)
                        // The retry time will be checked in drain() before processing
                        // (Java line 177: segmentsToFetch.add(request))
                        // Permit already released (Java line 174)
                        coordinator.download_queue.push(Reverse(request));
                    }
                    Ok(DownloadResult::FailedPermanently { error, result_sender }) => {
                        // Permanent failure - deliver error to future
                        if !result_sender.is_closed() {
                            let _ = result_sender.send(Err(error));
                        }
                        // Permit already released
                    }
                    Ok(DownloadResult::Cancelled) => {
                        // Cancelled - permit already released, nothing to do
                    }
                    Err(e) => {
                        log::error!("Download task panicked: {e:?}");
                        // Permit already released via RAII
                    }
                }
                // Immediately try to start another download
                continue;
            }

            // Event 3: Recycled (only wait when blocked on permits with pending work)
            _ = coordinator.recycle_notify.notified(),
                if coordinator.should_wait_for_recycle() => {
                // Wake up to try draining
                continue;
            }

            // Event 4: Retry timer - wake up when next retry is ready
            _ = tokio::time::sleep(next_retry_sleep.unwrap_or(tokio::time::Duration::from_secs(3600))),
                if next_retry_sleep.is_some() => {
                // Wake up to retry deferred requests
                continue;
            }

            else => break,  // All channels closed AND no work pending
        }
    }
}

type CompletionCallback = Box<dyn Fn() + Send + Sync>;

/// Future for a remote log download request
pub struct RemoteLogDownloadFuture {
    result: Arc<Mutex<Option<Result<RemoteLogFile>>>>,
    completion_callbacks: Arc<Mutex<Vec<CompletionCallback>>>,
}

impl RemoteLogDownloadFuture {
    pub fn new(receiver: oneshot::Receiver<Result<RemoteLogFile>>) -> Self {
        let result = Arc::new(Mutex::new(None));
        let result_clone = Arc::clone(&result);
        let completion_callbacks: Arc<Mutex<Vec<CompletionCallback>>> =
            Arc::new(Mutex::new(Vec::new()));
        let callbacks_clone = Arc::clone(&completion_callbacks);

        // Spawn a task to wait for the download and update result, then call callbacks
        tokio::spawn(async move {
            let download_result = match receiver.await {
                Ok(Ok(path)) => Ok(path),
                Ok(Err(e)) => Err(e),
                Err(e) => Err(Error::UnexpectedError {
                    message: format!("Download & Read future cancelled: {e:?}"),
                    source: None,
                }),
            };

            *result_clone.lock() = Some(download_result);

            // Call all registered callbacks
            // We need to take the callbacks to avoid holding the lock while calling them
            // This also ensures that any callbacks registered after this point will be called immediately
            let callbacks: Vec<CompletionCallback> = {
                let mut callbacks_guard = callbacks_clone.lock();
                mem::take(&mut *callbacks_guard)
            };
            for callback in callbacks {
                callback();
            }

            // After calling callbacks, any new callbacks registered will see is_done() == true
            // and will be called immediately in on_complete()
        });

        Self {
            result,
            completion_callbacks,
        }
    }

    /// Register a callback to be called when download completes (similar to Java's onComplete)
    pub fn on_complete<F>(&self, callback: F)
    where
        F: Fn() + Send + Sync + 'static,
    {
        // Acquire callbacks lock first to ensure atomicity of the check-and-register operation
        let mut callbacks_guard = self.completion_callbacks.lock();

        // Check completion status while holding the callbacks lock.
        // This ensures that:
        // 1. If the task completes between checking is_done() and registering the callback,
        //    we'll see the completion state correctly
        // 2. The background task cannot clear the callbacks list while we're checking/registering
        let is_done = self.result.lock().is_some();

        if is_done {
            // If already completed, call immediately (drop lock first to avoid deadlock)
            drop(callbacks_guard);
            callback();
        } else {
            // Register the callback while holding the callbacks lock.
            // This ensures that even if the background task completes right after we check
            // is_done(), it will wait for us to release the lock before taking callbacks.
            // When it does take callbacks, it will see our callback in the list and execute it.
            callbacks_guard.push(Box::new(callback));
            // Lock is automatically released here
        }
    }

    pub fn is_done(&self) -> bool {
        self.result.lock().is_some()
    }

    /// Take the RemoteLogFile (including the permit) from this future
    /// This should only be called when the download is complete
    /// This is the correct way to consume the download - it transfers permit ownership
    pub fn take_remote_log_file(&self) -> Result<RemoteLogFile> {
        let mut guard = self.result.lock();
        match guard.take() {
            Some(Ok(remote_log_file)) => Ok(remote_log_file),
            Some(Err(e)) => {
                let error_msg = format!("{e}");
                Err(Error::IoUnexpectedError {
                    message: format!("Fail to get remote log file: {error_msg}"),
                    source: io::Error::other(error_msg),
                })
            }
            None => Err(Error::IoUnexpectedError {
                message: "Remote log file already taken or not ready".to_string(),
                source: io::Error::other("Remote log file already taken or not ready"),
            }),
        }
    }
}

/// Downloader for remote log segment files.
///
/// # Shutdown behavior
///
/// When the downloader is dropped, the request channel closes, signaling the coordinator
/// to stop accepting new work. The coordinator will finish any in-flight downloads but
/// won't wait for completion. Pending futures will fail.
pub struct RemoteLogDownloader {
    request_sender: Option<mpsc::UnboundedSender<RemoteLogDownloadRequest>>,
}

impl RemoteLogDownloader {
    pub fn new(
        local_log_dir: TempDir,
        max_prefetch_segments: usize,
        max_concurrent_downloads: usize,
        remote_log_read_concurrency: usize,
        credentials_rx: CredentialsReceiver,
    ) -> Result<Self> {
        let fetcher = Arc::new(ProductionFetcher {
            credentials_rx,
            local_log_dir: Arc::new(local_log_dir),
            remote_log_read_concurrency: remote_log_read_concurrency.max(1),
        });

        Self::new_with_fetcher(fetcher, max_prefetch_segments, max_concurrent_downloads)
    }

    /// Create a RemoteLogDownloader with a custom fetcher (for testing).
    pub fn new_with_fetcher(
        fetcher: Arc<dyn RemoteLogFetcher>,
        max_prefetch_segments: usize,
        max_concurrent_downloads: usize,
    ) -> Result<Self> {
        let (request_sender, request_receiver) = mpsc::unbounded_channel();

        let coordinator = DownloadCoordinator {
            download_queue: BinaryHeap::new(),
            active_downloads: JoinSet::new(),
            in_flight: 0,
            prefetch_semaphore: Arc::new(Semaphore::new(max_prefetch_segments)),
            max_concurrent_downloads,
            recycle_notify: Arc::new(Notify::new()),
            fetcher,
        };

        // Spawn coordinator task - it will exit when request_sender is dropped
        tokio::spawn(coordinator_loop(coordinator, request_receiver));

        Ok(Self {
            request_sender: Some(request_sender),
        })
    }

    /// Request to fetch a remote log segment to local. This method is non-blocking.
    pub fn request_remote_log(
        &self,
        remote_log_tablet_dir: &str,
        segment: &RemoteLogSegment,
    ) -> RemoteLogDownloadFuture {
        let (result_sender, result_receiver) = oneshot::channel();

        let request = RemoteLogDownloadRequest {
            segment: segment.clone(),
            remote_log_tablet_dir: remote_log_tablet_dir.to_string(),
            result_sender,
            retry_count: 0,
            next_retry_at: None,
        };

        // Send to coordinator (non-blocking)
        if let Some(ref sender) = self.request_sender {
            if sender.send(request).is_err() {
                // Coordinator is gone - immediately fail the future
                let (error_sender, error_receiver) = oneshot::channel();
                let _ = error_sender.send(Err(Error::UnexpectedError {
                    message: "RemoteLogDownloader coordinator has shut down".to_string(),
                    source: None,
                }));
                return RemoteLogDownloadFuture::new(error_receiver);
            }
        }

        RemoteLogDownloadFuture::new(result_receiver)
    }
}

impl Drop for RemoteLogDownloader {
    fn drop(&mut self) {
        // Drop the request sender to signal coordinator shutdown.
        // This causes request_receiver.recv() to return None, allowing the
        // coordinator to exit gracefully after processing pending work.
        // The coordinator task will finish on its own when it sees the channel closed.
        drop(self.request_sender.take());
    }
}

impl RemoteLogDownloader {
    /// Download a file from remote storage to local using streaming read/write.
    async fn download_file(
        remote_log_tablet_dir: &str,
        remote_path: &str,
        local_path: &Path,
        remote_fs_props: &HashMap<String, String>,
        remote_log_read_concurrency: usize,
    ) -> Result<PathBuf> {
        // Handle both URL (e.g., "s3://bucket/path") and local file paths
        // If the path doesn't contain "://", treat it as a local file path
        let remote_log_tablet_dir_url = if remote_log_tablet_dir.contains("://") {
            remote_log_tablet_dir.to_string()
        } else {
            format!("file://{remote_log_tablet_dir}")
        };

        // Create FileIO from the remote log tablet dir URL to get the storage
        let file_io_builder = FileIO::from_url(&remote_log_tablet_dir_url)?;

        // For S3/S3A URLs, inject S3 credentials from props
        let file_io_builder = if remote_log_tablet_dir.starts_with("s3://")
            || remote_log_tablet_dir.starts_with("s3a://")
            || remote_log_tablet_dir.starts_with("oss://")
        {
            file_io_builder.with_props(
                remote_fs_props
                    .iter()
                    .map(|(k, v)| (k.as_str(), v.as_str())),
            )
        } else {
            file_io_builder
        };

        // Build storage and create operator directly
        let storage = Storage::build(file_io_builder)?;
        let (op, relative_path) = storage.create(remote_path)?;

        // Timeout for remote storage operations (30 seconds)
        const REMOTE_OP_TIMEOUT: Duration = Duration::from_secs(30);
        const CHUNK_SIZE: usize = 8 * 1024 * 1024; // 8MiB

        Self::download_file_streaming(
            &op,
            relative_path,
            remote_path,
            local_path,
            CHUNK_SIZE,
            remote_log_read_concurrency,
            REMOTE_OP_TIMEOUT,
        )
        .await?;

        Ok(local_path.to_path_buf())
    }

    async fn download_file_streaming(
        op: &opendal::Operator,
        relative_path: &str,
        remote_path: &str,
        local_path: &Path,
        chunk_size: usize,
        streaming_read_concurrency: usize,
        remote_op_timeout: Duration,
    ) -> Result<()> {
        let mut local_file = tokio::fs::File::create(local_path).await?;

        let reader_future = op
            .reader_with(relative_path)
            .chunk(chunk_size)
            .concurrent(streaming_read_concurrency);
        let reader = tokio::time::timeout(remote_op_timeout, reader_future)
            .await
            .map_err(|e| Error::IoUnexpectedError {
                message: format!("Timeout creating streaming reader for {remote_path}: {e}."),
                source: io::ErrorKind::TimedOut.into(),
            })??;

        let mut stream = tokio::time::timeout(remote_op_timeout, reader.into_bytes_stream(..))
            .await
            .map_err(|e| Error::IoUnexpectedError {
                message: format!("Timeout creating streaming bytes stream for {remote_path}: {e}."),
                source: io::ErrorKind::TimedOut.into(),
            })??;

        let mut chunk_count = 0u64;
        while let Some(chunk) = tokio::time::timeout(remote_op_timeout, stream.try_next())
            .await
            .map_err(|e| Error::IoUnexpectedError {
                message: format!(
                    "Timeout streaming chunk from remote storage: {remote_path}, exception: {e}."
                ),
                source: io::ErrorKind::TimedOut.into(),
            })??
        {
            chunk_count += 1;
            if chunk_count <= 3 || chunk_count % 10 == 0 {
                log::debug!("Remote log streaming download: chunk #{chunk_count} ({remote_path})");
            }
            local_file.write_all(&chunk).await?;
        }

        local_file.sync_all().await?;
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::atomic::{AtomicUsize, Ordering};

    /// Helper function to create a TableBucket for testing
    fn create_table_bucket(table_id: i64, bucket_id: i32) -> TableBucket {
        TableBucket::new(table_id, bucket_id)
    }

    /// Simplified fake fetcher for testing
    struct FakeFetcher {
        completion_gate: Arc<Notify>,
        in_flight: Arc<AtomicUsize>,
        max_seen_in_flight: Arc<AtomicUsize>,
        fail_count: Arc<Mutex<usize>>,
        auto_complete: bool,
    }

    impl FakeFetcher {
        fn new(fail_count: usize, auto_complete: bool) -> Self {
            Self {
                completion_gate: Arc::new(Notify::new()),
                in_flight: Arc::new(AtomicUsize::new(0)),
                max_seen_in_flight: Arc::new(AtomicUsize::new(0)),
                fail_count: Arc::new(Mutex::new(fail_count)),
                auto_complete,
            }
        }

        fn max_seen_in_flight(&self) -> usize {
            self.max_seen_in_flight.load(Ordering::SeqCst)
        }

        fn in_flight(&self) -> usize {
            self.in_flight.load(Ordering::SeqCst)
        }

        fn release_one(&self) {
            self.completion_gate.notify_one();
        }

        fn release_all(&self) {
            self.completion_gate.notify_waiters();
        }
    }

    impl RemoteLogFetcher for FakeFetcher {
        fn fetch(
            &self,
            request: &RemoteLogDownloadRequest,
        ) -> Pin<Box<dyn Future<Output = Result<FetchResult>> + Send>> {
            let gate = self.completion_gate.clone();
            let in_flight = self.in_flight.clone();
            let max_seen = self.max_seen_in_flight.clone();
            let fail_count = self.fail_count.clone();
            let segment_id = request.segment().segment_id.clone();
            let auto_complete = self.auto_complete;

            Box::pin(async move {
                // Track in-flight
                let current = in_flight.fetch_add(1, Ordering::SeqCst) + 1;
                max_seen.fetch_max(current, Ordering::SeqCst);

                // Wait for gate (or auto-complete)
                if !auto_complete {
                    gate.notified().await;
                } else {
                    tokio::task::yield_now().await;
                }

                // Check if should fail
                let should_fail = {
                    let mut count = fail_count.lock();
                    if *count > 0 {
                        *count -= 1;
                        true
                    } else {
                        false
                    }
                };

                in_flight.fetch_sub(1, Ordering::SeqCst);

                if should_fail {
                    Err(Error::UnexpectedError {
                        message: format!("Fake fetch failed for {segment_id}"),
                        source: None,
                    })
                } else {
                    let fake_data = vec![1, 2, 3, 4];
                    let temp_dir = env::temp_dir();
                    let timestamp = SystemTime::now()
                        .duration_since(UNIX_EPOCH)
                        .unwrap()
                        .as_nanos();
                    let file_path =
                        temp_dir.join(format!("fake_segment_{segment_id}_{timestamp}.log"));
                    tokio::fs::write(&file_path, &fake_data).await?;

                    Ok(FetchResult {
                        file_path,
                        file_size: fake_data.len(),
                    })
                }
            })
        }
    }

    /// Helper function to create a RemoteLogSegment for testing
    fn create_segment(
        segment_id: &str,
        start_offset: i64,
        max_timestamp: i64,
        table_bucket: TableBucket,
    ) -> RemoteLogSegment {
        RemoteLogSegment {
            segment_id: segment_id.to_string(),
            start_offset,
            end_offset: start_offset + 1000,
            size_in_bytes: 1024,
            table_bucket,
            max_timestamp,
        }
    }

    /// Helper function to create a RemoteLogDownloadRequest for testing
    fn create_request(segment: RemoteLogSegment) -> RemoteLogDownloadRequest {
        let (result_sender, _) = oneshot::channel();
        RemoteLogDownloadRequest {
            remote_log_tablet_dir: "test_dir".to_string(),
            segment,
            result_sender,
            retry_count: 0,
            next_retry_at: None,
        }
    }

    #[test]
    fn test_priority_ordering_matching_java_test_case() {
        // Test priority ordering: timestamp across buckets, offset within bucket
        // Does NOT test tie-breakers (segment_id) - those are implementation details

        let bucket1 = create_table_bucket(1, 0);
        let bucket2 = create_table_bucket(1, 1);
        let bucket3 = create_table_bucket(1, 2);
        let bucket4 = create_table_bucket(1, 3);

        // Create segments with distinct timestamps/offsets (no ties)
        let seg_negative = create_segment("seg_neg", 0, -1, bucket1.clone());
        let seg_zero = create_segment("seg_zero", 0, 0, bucket2.clone());
        let seg_1000 = create_segment("seg_1000", 0, 1000, bucket3.clone());
        let seg_2000 = create_segment("seg_2000", 0, 2000, bucket4.clone());
        let seg_same_bucket_100 = create_segment("seg_sb_100", 100, 5000, bucket1.clone());
        let seg_same_bucket_50 = create_segment("seg_sb_50", 50, 5000, bucket1.clone());

        let mut heap = BinaryHeap::new();
        heap.push(Reverse(create_request(seg_2000)));
        heap.push(Reverse(create_request(seg_same_bucket_100)));
        heap.push(Reverse(create_request(seg_1000)));
        heap.push(Reverse(create_request(seg_zero)));
        heap.push(Reverse(create_request(seg_negative)));
        heap.push(Reverse(create_request(seg_same_bucket_50)));

        // Verify ordering by timestamp/offset, not segment_id
        let first = heap.pop().unwrap().0;
        assert_eq!(first.segment.max_timestamp, -1, "Lowest timestamp first");

        let second = heap.pop().unwrap().0;
        assert_eq!(second.segment.max_timestamp, 0);

        let third = heap.pop().unwrap().0;
        assert_eq!(third.segment.max_timestamp, 1000);

        let fourth = heap.pop().unwrap().0;
        assert_eq!(fourth.segment.max_timestamp, 2000);

        // Last two are same bucket (ts=5000), ordered by offset
        let fifth = heap.pop().unwrap().0;
        assert_eq!(fifth.segment.max_timestamp, 5000);
        assert_eq!(
            fifth.segment.start_offset, 50,
            "Lower offset first within bucket"
        );

        let sixth = heap.pop().unwrap().0;
        assert_eq!(sixth.segment.max_timestamp, 5000);
        assert_eq!(sixth.segment.start_offset, 100);
    }

    #[tokio::test]
    async fn test_concurrency_and_priority() {
        // Test concurrency limiting and priority-based scheduling together
        let fake_fetcher = Arc::new(FakeFetcher::new(0, false)); // Manual control

        let downloader = RemoteLogDownloader::new_with_fetcher(
            fake_fetcher.clone(),
            10, // High prefetch limit
            2,  // Max concurrent downloads = 2
        )
        .unwrap();

        let bucket = create_table_bucket(1, 0);

        // Request 4 segments with same priority (to isolate concurrency limiting from priority)
        let segs: Vec<_> = (0..4)
            .map(|i| create_segment(&format!("seg{i}"), i * 100, 1000, bucket.clone()))
            .collect();

        let _futures: Vec<_> = segs
            .iter()
            .map(|seg| downloader.request_remote_log("dir", seg))
            .collect();

        // Wait for exactly 2 to start
        tokio::time::sleep(Duration::from_millis(50)).await;
        assert_eq!(
            fake_fetcher.in_flight(),
            2,
            "Concurrency limit: exactly 2 should be in-flight"
        );

        // Release one
        fake_fetcher.release_one();
        tokio::time::sleep(Duration::from_millis(50)).await;

        // Max should never exceed 2
        assert_eq!(
            fake_fetcher.max_seen_in_flight(),
            2,
            "Max concurrent should not exceed 2"
        );

        // Release all
        fake_fetcher.release_all();
    }

    #[tokio::test]
    async fn test_prefetch_limit() {
        // Test that prefetch semaphore limits outstanding downloads
        let fake_fetcher = Arc::new(FakeFetcher::new(0, true)); // Auto-complete

        let downloader = RemoteLogDownloader::new_with_fetcher(
            fake_fetcher,
            2,  // Max prefetch = 2
            10, // High concurrent limit
        )
        .unwrap();

        let bucket = create_table_bucket(1, 0);

        // Request 4 downloads
        let segs: Vec<_> = (0..4)
            .map(|i| create_segment(&format!("seg{i}"), i * 100, 1000, bucket.clone()))
            .collect();

        let mut futures: Vec<_> = segs
            .iter()
            .map(|seg| downloader.request_remote_log("dir", seg))
            .collect();

        // Wait for first 2 to complete
        let deadline = tokio::time::Instant::now() + Duration::from_secs(2);
        loop {
            if futures.iter().filter(|f| f.is_done()).count() >= 2 {
                break;
            }
            if tokio::time::Instant::now() > deadline {
                panic!("Timeout waiting for first 2 downloads");
            }
            tokio::time::sleep(Duration::from_millis(10)).await;
        }

        // Verify 3rd and 4th are blocked (prefetch limit)
        tokio::time::sleep(Duration::from_millis(50)).await;
        assert_eq!(
            futures.iter().filter(|f| f.is_done()).count(),
            2,
            "Prefetch limit: only 2 should complete"
        );

        // Drop first 2 (releases permits)
        let f4 = futures.pop().unwrap();
        let f3 = futures.pop().unwrap();
        drop(futures);

        // 3rd and 4th should now complete
        let deadline = tokio::time::Instant::now() + Duration::from_secs(2);
        loop {
            if f3.is_done() && f4.is_done() {
                break;
            }
            if tokio::time::Instant::now() > deadline {
                panic!("Timeout after permit release");
            }
            tokio::time::sleep(Duration::from_millis(10)).await;
        }
    }

    #[tokio::test]
    async fn test_retry_and_cancellation() {
        // Test retry with exponential backoff
        let fake_fetcher = Arc::new(FakeFetcher::new(2, true)); // Fail twice, succeed third time

        let downloader =
            RemoteLogDownloader::new_with_fetcher(fake_fetcher.clone(), 10, 1).unwrap();

        let bucket = create_table_bucket(1, 0);
        let seg = create_segment("seg1", 0, 1000, bucket);

        let future = downloader.request_remote_log("dir", &seg);

        // Should succeed after retries
        let deadline = tokio::time::Instant::now() + Duration::from_secs(5);
        loop {
            if future.is_done() {
                break;
            }
            if tokio::time::Instant::now() > deadline {
                panic!("Timeout waiting for retry to succeed");
            }
            tokio::time::sleep(Duration::from_millis(50)).await;
        }

        assert!(future.is_done(), "Should succeed after retries");

        // Test cancellation
        let seg2 = create_segment("seg2", 100, 1000, create_table_bucket(1, 0));
        let fake_fetcher2 = Arc::new(FakeFetcher::new(100, true)); // Fail forever
        let downloader2 =
            RemoteLogDownloader::new_with_fetcher(fake_fetcher2.clone(), 10, 1).unwrap();

        let future2 = downloader2.request_remote_log("dir", &seg2);
        tokio::time::sleep(Duration::from_millis(50)).await;

        // Drop to cancel
        drop(future2);
        tokio::time::sleep(Duration::from_millis(50)).await;

        assert_eq!(
            fake_fetcher2.in_flight(),
            0,
            "Cancellation should release resources"
        );
    }
}