use std::{
    collections::{HashMap, hash_map::Entry},
    io::ErrorKind,
    path::Path,
    sync::{
        Arc,
        atomic::{AtomicBool, Ordering},
    },
    time::{Duration, Instant},
};

use anyhow::{Context, Result, anyhow, bail, ensure};
use log::{info, warn};
use storage_backend::storage::BackendStorage;
use tici_grpc::FragMeta;
use tici_system::{disk::get_disk_space_stats, readable::ReadableSize};
use tokio_util::sync::CancellationToken;

use crate::utils::{
    metrics::{READER_FRAG_GAUGE, READER_QUERY_LATENCY},
    s3_frag,
};

#[derive(Debug)]
pub struct LocalFragDataMeta {
    local_frag_path: String,
    ref_count: u64,
    update_time: Instant,
}

impl LocalFragDataMeta {
    fn new(local_frag_path: String, ref_count: u64) -> Self {
        Self {
            local_frag_path,
            ref_count,
            update_time: Instant::now(),
        }
    }

    fn inc_ref_count(&mut self) -> u64 {
        self.ref_count += 1;
        self.update_time = Instant::now();
        self.ref_count
    }

    fn dec_ref_count(&mut self) -> u64 {
        if self.ref_count == 0 {
            panic!("ref_count is 0, cannot be decrease");
        }

        self.ref_count -= 1;
        self.update_time = Instant::now();
        self.ref_count
    }
}

// One S3 object (identified by s3_frag_path) maps to one local directory.
// Format: {frag_base_path}/{s3_frag_path}
#[derive(Eq, Hash, PartialEq, Debug, Clone)]
struct FragCacheKey {
    s3_frag_path: String,
}

impl std::fmt::Display for FragCacheKey {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.s3_frag_path)
    }
}

impl FragCacheKey {
    fn new(s3_frag_path: String) -> Self {
        Self { s3_frag_path }
    }

    fn from_local_path(local_frag_path: &str, local_data_path: &str) -> Result<Self> {
        let relative_path = local_frag_path
            .strip_prefix(local_data_path)
            .ok_or_else(|| {
                anyhow!(
                    "local_frag_path {} is not in local_data_path {}",
                    local_frag_path,
                    local_data_path
                )
            })?
            .trim_start_matches('/')
            .trim_end_matches('/');

        if relative_path.is_empty() {
            bail!(
                "empty relative path from local_frag_path={}",
                local_frag_path
            );
        }

        Ok(Self::new(relative_path.to_string()))
    }
}

type DownloadingFragCell = Arc<tokio::sync::OnceCell<String>>;

#[derive(Debug)]
enum FragCacheValue {
    Downloading(DownloadingFragCell),
    Ready(LocalFragDataMeta),
}

enum AcquireFragState {
    Downloading { download_cell: DownloadingFragCell },
    Ready(String),
}

// FragLease must be cheap to clone because FragReaderMeta is cloned in the
// reader build path. We therefore model it as a shareable handle around one
// acquired cache reference instead of storing the release state inline.
//
// If FragLease directly owned these fields and simply derived Clone, each clone
// would run Drop and we would release the same acquired ref multiple times. If
// FragLease were non-cloneable, it would no longer fit the existing
// FragReaderMeta clone flow. Sharing one inner lease keeps the runtime behavior
// aligned with the cache semantics: one acquire maps to exactly one release
// when the last cloned handle drops.
#[derive(Clone)]
pub struct FragLease {
    inner: Arc<FragLeaseInner>,
}

struct FragLeaseInner {
    frag_data_cache: Arc<FragDataCache>,
    frag_cache_key: FragCacheKey,
    local_frag_path: String,
}

impl FragLease {
    fn new(
        frag_data_cache: Arc<FragDataCache>,
        frag_cache_key: FragCacheKey,
        local_frag_path: String,
    ) -> Self {
        Self {
            inner: Arc::new(FragLeaseInner {
                frag_data_cache,
                frag_cache_key,
                local_frag_path,
            }),
        }
    }

    pub fn local_frag_path(&self) -> &str {
        &self.inner.local_frag_path
    }
}

impl Drop for FragLeaseInner {
    fn drop(&mut self) {
        // The cache ref is released exactly once for the acquire that created
        // this lease, even if the outer FragLease handle was cloned.
        self.frag_data_cache
            .release_frag(&self.frag_cache_key)
            .unwrap();
    }
}

pub struct FragDataCache {
    backend_storage: Arc<dyn BackendStorage>,
    local_data_path: String,
    temporary_path: String,
    frag_states: parking_lot::Mutex<HashMap<FragCacheKey, FragCacheValue>>,
    remove_zero_ref_frag_delay: Duration,
    disk_reserve_size: ReadableSize,
}

// FragDataCache is a basic local disk file cache. It downloads files from S3
// and stores them in a local directory.
impl FragDataCache {
    pub fn new(
        backend_storage: Arc<dyn BackendStorage>,
        local_data_path: String,
        temporary_path: String,
        remove_zero_ref_frag_delay: Duration,
        disk_reserve_size: ReadableSize,
    ) -> Result<Self> {
        std::fs::create_dir_all(&local_data_path)
            .with_context(|| format!("Failed to create directory {}", local_data_path))?;
        Self::clear_temporary(&local_data_path)?;
        let frag_states = Self::load_cached_frags(&local_data_path)?;
        Self::update_cached_frag_gauge(&frag_states);
        Ok(FragDataCache {
            backend_storage,
            local_data_path,
            temporary_path,
            frag_states: parking_lot::Mutex::new(frag_states),
            remove_zero_ref_frag_delay,
            disk_reserve_size,
        })
    }

    fn clear_temporary(local_data_path: &str) -> Result<()> {
        let tmp_download_path = s3_frag::temporary_download_path(local_data_path);
        if let Err(e) = std::fs::remove_dir_all(&tmp_download_path) {
            if e.kind() != ErrorKind::NotFound {
                return Err(e).with_context(|| format!("Remove {}", tmp_download_path));
            }
        }
        Ok(())
    }

    // Traverse the local directory to collect all downloaded fragments. The
    // directory format pattern likes local_data_path/xxx/yyy/zzz/frag/index.
    // Therefore, we only need to traverse the deepest index directories, whose
    // parent directories are the fragment directories.
    fn load_cached_frags(local_data_path: &str) -> Result<HashMap<FragCacheKey, FragCacheValue>> {
        let dir = Path::new(local_data_path);
        let mut frag_states = HashMap::new();
        let mut directories_to_check = vec![dir.to_path_buf()];
        while let Some(current_dir) = directories_to_check.pop() {
            let read_dir = std::fs::read_dir(&current_dir)
                .with_context(|| format!("Failed to read directory {:?}", current_dir))?;

            let mut has_subdirectories = false;
            for entry in read_dir {
                let entry = entry.with_context(|| {
                    format!("Failed to read entry in directory {:?}", current_dir)
                })?;
                if entry.path().is_dir() {
                    has_subdirectories = true;
                    directories_to_check.push(entry.path());
                }
            }

            // It is the deepest directory.
            if !has_subdirectories {
                if current_dir.file_name().unwrap().to_str().unwrap() != "index" {
                    info!("current_dir {:?} is not index", current_dir);
                    continue;
                }
                let local_frag_path = current_dir.parent().unwrap().display().to_string();
                let frag_cache_key =
                    FragCacheKey::from_local_path(&local_frag_path, local_data_path)?;
                frag_states.insert(
                    frag_cache_key,
                    FragCacheValue::Ready(LocalFragDataMeta::new(local_frag_path, 0)),
                );
            }
        }
        Ok(frag_states)
    }

    fn cached_frag_count(frag_states: &HashMap<FragCacheKey, FragCacheValue>) -> usize {
        frag_states
            .values()
            .filter(|frag_state| matches!(frag_state, FragCacheValue::Ready(_)))
            .count()
    }

    fn update_cached_frag_gauge(frag_states: &HashMap<FragCacheKey, FragCacheValue>) {
        READER_FRAG_GAUGE.set(Self::cached_frag_count(frag_states) as i64);
    }

    pub fn local_frag_path(&self, s3_frag_path: &str) -> String {
        format!("{}/{}", self.local_data_path, s3_frag_path)
    }

    pub async fn acquire_frag(
        self: &Arc<Self>,
        s3_bucket: &str,
        frag_meta: &FragMeta,
    ) -> Result<FragLease> {
        let start = Instant::now();
        let frag_cache_key = FragCacheKey::new(frag_meta.frag_path.clone());
        if let Some(local_frag_path) = self.acquire_ready_frag(&frag_cache_key) {
            return Ok(FragLease::new(
                self.clone(),
                frag_cache_key,
                local_frag_path,
            ));
        }
        let frag_size = frag_meta
            .property
            .as_ref()
            .ok_or_else(|| {
                anyhow!(
                    "frag_meta.property is None, frag_cache_key={}",
                    frag_cache_key
                )
            })?
            .size;
        let download_cell = match self.get_or_insert_frag_state(&frag_cache_key) {
            AcquireFragState::Ready(local_frag_path) => {
                return Ok(FragLease::new(
                    self.clone(),
                    frag_cache_key,
                    local_frag_path,
                ));
            }
            AcquireFragState::Downloading { download_cell } => download_cell,
        };
        let did_download = AtomicBool::new(false);
        let local_frag_path_result = {
            let frag_cache_key_for_download = frag_cache_key.clone();
            let download_cell_for_download = download_cell.clone();
            download_cell
                .get_or_try_init(|| async {
                    // `get_or_try_init` leaves the cell uninitialized on
                    // error, so another waiter may run this initializer on the
                    // same cell. Keep one download cell as one shared attempt:
                    // failed cells are removed before returning the error, and
                    // waiters on the old cell fail the current-state check
                    // instead of turning waiter count into implicit retries.
                    let download_result = async {
                        self.ensure_current_downloading_frag_state(
                            &frag_cache_key_for_download,
                            &download_cell_for_download,
                        )?;
                        self.check_disk_space(frag_size)?;
                        let download_start = Instant::now();
                        let local_frag_path = s3_frag::download_frag(
                            self.backend_storage.as_ref(),
                            s3_bucket,
                            frag_meta,
                            &self.local_data_path,
                        )
                        .await?;
                        did_download.store(true, Ordering::Relaxed);

                        let download_elapsed = download_start.elapsed();
                        READER_QUERY_LATENCY
                            .with_label_values(&["download_frag"])
                            .observe(download_elapsed.as_secs_f64());
                        let total_elapsed = start.elapsed().as_millis();
                        info!(
                            "Frag downloaded: {} -> {}, download time={}ms, total time={}ms",
                            frag_cache_key_for_download,
                            local_frag_path,
                            download_elapsed.as_millis(),
                            total_elapsed,
                        );
                        let marked_ready = self.mark_frag_ready(
                            &frag_cache_key_for_download,
                            &download_cell_for_download,
                            &local_frag_path,
                            1,
                        );
                        ensure!(
                            marked_ready,
                            "frag_cache_key {} should still be downloading when download finishes",
                            frag_cache_key_for_download
                        );
                        fail::fail_point!("frag_data_cache_after_mark_frag_ready");
                        Ok(local_frag_path)
                    }
                    .await;
                    if download_result.is_err() {
                        self.remove_downloading_frag_state_if_same_cell(
                            &frag_cache_key_for_download,
                            &download_cell_for_download,
                        );
                    }
                    download_result
                })
                .await
        };
        // `get_or_try_init` only propagates the `Err` from this task's own init
        // closure (other waiters' failures cause a re-run rather than
        // forwarding their `Err`), and the closure already cleans up the failed
        // cell before returning. Centralizing cleanup inside the initializer
        // keeps the removal responsibility in a single place.
        let local_frag_path = local_frag_path_result?;
        if did_download.load(Ordering::Relaxed) {
            return Ok(FragLease::new(
                self.clone(),
                frag_cache_key,
                local_frag_path.clone(),
            ));
        }

        // Waiters on the same download cell still need to take their own
        // references after the downloader publishes the frag as `Ready`.
        let local_frag_path = self.acquire_ready_frag(&frag_cache_key).ok_or_else(|| {
            anyhow!(
                "frag_cache_key {} should exist after acquire succeeds, local_frag_path={}",
                frag_cache_key,
                local_frag_path
            )
        })?;
        Ok(FragLease::new(
            self.clone(),
            frag_cache_key,
            local_frag_path,
        ))
    }

    fn release_frag(&self, frag_cache_key: &FragCacheKey) -> Result<()> {
        let mut frag_states = self.frag_states.lock();
        let frag_state = frag_states
            .get_mut(frag_cache_key)
            .ok_or_else(|| anyhow!("frag_cache_key {} not found", frag_cache_key))?;
        let ref_count = match frag_state {
            FragCacheValue::Ready(frag) => frag.dec_ref_count(),
            FragCacheValue::Downloading(_) => {
                bail!("frag_cache_key {} is still downloading", frag_cache_key)
            }
        };
        info!(
            "release_frag: frag_cache_key={}, ref_count={}",
            frag_cache_key, ref_count
        );
        Ok(())
    }

    fn acquire_ready_frag(&self, frag_cache_key: &FragCacheKey) -> Option<String> {
        let mut frag_states = self.frag_states.lock();
        let frag_state = frag_states.get_mut(frag_cache_key)?;
        match frag_state {
            FragCacheValue::Ready(frag) => {
                let ref_count = frag.inc_ref_count();
                info!(
                    "get_cached_frag: frag_cache_key={}, ref_count={}",
                    frag_cache_key, ref_count
                );
                Some(frag.local_frag_path.clone())
            }
            FragCacheValue::Downloading(_) => None,
        }
    }

    fn get_or_insert_frag_state(&self, frag_cache_key: &FragCacheKey) -> AcquireFragState {
        let mut frag_states = self.frag_states.lock();
        match frag_states.entry(frag_cache_key.clone()) {
            Entry::Occupied(mut entry) => match entry.get_mut() {
                FragCacheValue::Ready(frag) => {
                    let ref_count = frag.inc_ref_count();
                    info!(
                        "get_cached_frag: frag_cache_key={}, ref_count={}",
                        frag_cache_key, ref_count
                    );
                    AcquireFragState::Ready(frag.local_frag_path.clone())
                }
                FragCacheValue::Downloading(download_cell) => AcquireFragState::Downloading {
                    download_cell: download_cell.clone(),
                },
            },
            Entry::Vacant(entry) => {
                let download_cell = Arc::new(tokio::sync::OnceCell::new());
                entry.insert(FragCacheValue::Downloading(download_cell.clone()));
                AcquireFragState::Downloading { download_cell }
            }
        }
    }

    fn ensure_current_downloading_frag_state(
        &self,
        frag_cache_key: &FragCacheKey,
        download_cell: &DownloadingFragCell,
    ) -> Result<()> {
        // A failed download removes its cell from `frag_states` before the
        // OnceCell wakes other waiters. Those waiters may still own the old
        // Arc, but they must not initialize it after a fresh cell has become
        // the current download state for the same fragment.
        let frag_states = self.frag_states.lock();
        let is_current = matches!(
            frag_states.get(frag_cache_key),
            Some(FragCacheValue::Downloading(current_cell))
                if Arc::ptr_eq(current_cell, download_cell)
        );
        ensure!(
            is_current,
            "frag_cache_key {} download state is no longer current",
            frag_cache_key
        );
        Ok(())
    }

    fn remove_downloading_frag_state_if_same_cell(
        &self,
        frag_cache_key: &FragCacheKey,
        download_cell: &DownloadingFragCell,
    ) -> bool {
        let mut frag_states = self.frag_states.lock();
        match frag_states.entry(frag_cache_key.clone()) {
            Entry::Occupied(entry) => {
                // Remove only if `download_cell` is still the map entry. A newer
                // downloader may already have installed another cell.
                let should_remove = matches!(
                    entry.get(),
                    FragCacheValue::Downloading(current_cell)
                        if Arc::ptr_eq(current_cell, download_cell)
                );
                if should_remove {
                    entry.remove();
                    return true;
                }
                false
            }
            Entry::Vacant(_) => false,
        }
    }

    fn mark_frag_ready(
        &self,
        frag_cache_key: &FragCacheKey,
        download_cell: &DownloadingFragCell,
        local_frag_path: &str,
        initial_ref_count: u64,
    ) -> bool {
        // The downloader publishes the frag together with the references it has
        // already acquired locally. Waiters still go through
        // `acquire_ready_frag()` afterwards to take their own refs.
        let mut frag_states = self.frag_states.lock();
        let mut should_update_gauge = false;
        let mut marked_ready = false;
        match frag_states.entry(frag_cache_key.clone()) {
            Entry::Occupied(mut entry) => {
                let should_promote_to_ready = matches!(
                    entry.get(),
                    FragCacheValue::Downloading(current_cell)
                        if Arc::ptr_eq(current_cell, download_cell)
                );
                if should_promote_to_ready {
                    entry.insert(FragCacheValue::Ready(LocalFragDataMeta::new(
                        local_frag_path.to_string(),
                        initial_ref_count,
                    )));
                    should_update_gauge = true;
                    marked_ready = true;
                } else if matches!(entry.get(), FragCacheValue::Downloading(_)) {
                    warn!(
                        "Skip marking frag ready because download state changed, frag_cache_key={}",
                        frag_cache_key
                    );
                }
            }
            Entry::Vacant(_) => {
                warn!(
                    "Skip marking frag ready because cache state is missing, frag_cache_key={}",
                    frag_cache_key
                );
            }
        }
        if should_update_gauge {
            Self::update_cached_frag_gauge(&frag_states);
        }
        marked_ready
    }

    pub fn remove_zero_ref_frags(&self) {
        let start = Instant::now();
        let mut unreferenced_count = 0;
        let mut expired_frag_cache_keys = Vec::new();
        let mut frag_states = self.frag_states.lock();
        for (frag_cache_key, frag_state) in frag_states.iter() {
            let FragCacheValue::Ready(frag) = frag_state else {
                continue;
            };
            if frag.ref_count > 0 {
                continue;
            }
            unreferenced_count += 1;
            if frag.update_time.elapsed() > self.remove_zero_ref_frag_delay {
                info!("Remove cached frag: {} -> {:?}", frag_cache_key, frag);
                expired_frag_cache_keys.push(frag_cache_key.clone());
                if let Err(e) = std::fs::remove_dir_all(&frag.local_frag_path) {
                    warn!("Remove {} err={:?}", frag.local_frag_path, e);
                }
            }
        }

        let cached_frag_count = Self::cached_frag_count(&frag_states);
        for frag_cache_key in expired_frag_cache_keys.iter() {
            frag_states.remove(frag_cache_key);
        }
        Self::update_cached_frag_gauge(&frag_states);

        info!(
            "remove_stale_frags: cached_frag_count={}, unreferenced_count={}, remove_count={}, elapsed={}us",
            cached_frag_count,
            unreferenced_count,
            expired_frag_cache_keys.len(),
            start.elapsed().as_micros()
        );
    }

    pub fn spawn_gc_task(
        frag_data_cache: Arc<FragDataCache>,
        cancel_token: CancellationToken,
    ) -> tokio::task::JoinHandle<()> {
        tokio::spawn(async move {
            loop {
                tokio::select! {
                    _ = cancel_token.cancelled() => {
                        info!("GC task received cancellation signal");
                        break;
                    }
                    _ = tokio::time::sleep(frag_data_cache.remove_zero_ref_frag_delay) => {
                        frag_data_cache.remove_zero_ref_frags();
                    }
                }
            }
        })
    }

    fn check_disk_space(&self, frag_size: u64) -> Result<()> {
        let (disk_capacity, disk_available) = get_disk_space_stats(&self.local_data_path)
            .with_context(|| {
                format!(
                    "get_disk_space_stats: local_data_path={}",
                    self.local_data_path
                )
            })?;

        let required_size = frag_size + self.disk_reserve_size.0;
        ensure!(
            disk_available >= required_size,
            "Insufficient disk space for writing fragment at '{}': capacity={} bytes, available={} bytes, required={} bytes (= fragment={} bytes + reserve={} bytes)",
            self.local_data_path,
            disk_capacity,
            disk_available,
            required_size,
            frag_size,
            self.disk_reserve_size
        );
        Ok(())
    }

    pub fn get_temporary_path(&self) -> String {
        self.temporary_path.clone()
    }
}

#[cfg(test)]
mod tests {
    use std::{sync::Arc, time::Duration};

    use bytes::Bytes;
    use futures::StreamExt;
    use storage_backend::s3_multi_file::{FileBundle, FileMeta, MultiFileStream};
    use tempfile::TempDir;
    use tici_grpc::{FragMeta, FragProperty, FragStorageRef, KeyRange};
    use tici_system::readable::ReadableSize;

    use crate::{
        config::{FragReaderConfig, WorkerNodeConfig},
        reader::{frag_data_cache::FragDataCache, frag_reader::FragReaderMeta},
        test_utils::{
            MockBackendStorage, TEST_HDFSLOGS_INDEX_META, TEST_S3_BUCKET, is_sorted_frag,
            test_create_backend_storage, test_get_shard_writer_meta, write_test_data,
        },
        utils::{
            meta_client_pool::{MetaClientPool, MockMetaClientPool},
            s3_frag,
        },
    };

    async fn create_multi_file_object_data() -> Vec<u8> {
        let temp_dir = TempDir::new().unwrap();
        let local_frag_path = temp_dir.path().join("frag");
        std::fs::create_dir_all(local_frag_path.join("index")).unwrap();
        let meta_path = local_frag_path.join("index/meta.json");
        std::fs::write(&meta_path, b"{}").unwrap();

        let file_bundle = FileBundle::new(
            local_frag_path.display().to_string(),
            1,
            vec![FileMeta {
                relative_path: "index/meta.json".to_string(),
                size: std::fs::metadata(&meta_path).unwrap().len(),
            }],
        );
        let stream = MultiFileStream::new(&file_bundle).unwrap().into_stream();
        let chunks = stream.map(|chunk| chunk.unwrap()).collect::<Vec<_>>().await;
        let total_size = chunks.iter().map(Bytes::len).sum();
        let mut data = Vec::with_capacity(total_size);
        for chunk in chunks {
            data.extend_from_slice(&chunk);
        }
        data
    }

    fn ready_frag_count(frag_data_cache: &FragDataCache) -> usize {
        frag_data_cache
            .frag_states
            .lock()
            .values()
            .filter(|frag_state| matches!(frag_state, super::FragCacheValue::Ready(_)))
            .count()
    }

    fn contains_frag_state(
        frag_data_cache: &FragDataCache,
        frag_cache_key: &super::FragCacheKey,
    ) -> bool {
        frag_data_cache
            .frag_states
            .lock()
            .contains_key(frag_cache_key)
    }

    const TEST_FRAG_DATA_CACHE_FRAG_PATH: &str = "frag-data-cache-prefix/shard_1/f_1";
    const TEST_FRAG_DATA_CACHE_STORAGE_REF_KEY: &str = "multi.0";

    #[test]
    fn test_get_or_insert_frag_state_transitions_from_downloading_to_ready() {
        let backend_storage =
            Arc::new(MockBackendStorage::new().with_default_get_object(vec![], Duration::ZERO));
        let frag_data_cache = FragDataCache::new(
            backend_storage,
            TempDir::new().unwrap().path().to_str().unwrap().to_string(),
            TempDir::new().unwrap().path().to_str().unwrap().to_string(),
            Duration::from_secs(60),
            ReadableSize::mb(64),
        )
        .unwrap();
        let frag_cache_key = super::FragCacheKey::new("shard_1/frag_1".to_string());

        let download_cell = match frag_data_cache.get_or_insert_frag_state(&frag_cache_key) {
            super::AcquireFragState::Downloading { download_cell } => download_cell,
            super::AcquireFragState::Ready(_) => panic!("frag should not be ready before download"),
        };
        match frag_data_cache.get_or_insert_frag_state(&frag_cache_key) {
            super::AcquireFragState::Downloading {
                download_cell: reused_cell,
            } => {
                assert!(Arc::ptr_eq(&download_cell, &reused_cell));
            }
            super::AcquireFragState::Ready(_) => {
                panic!("frag should still be in downloading state")
            }
        }

        let local_frag_path = frag_data_cache.local_frag_path(&frag_cache_key.s3_frag_path);
        assert!(frag_data_cache.mark_frag_ready(
            &frag_cache_key,
            &download_cell,
            &local_frag_path,
            0,
        ));

        assert_eq!(
            frag_data_cache.acquire_ready_frag(&frag_cache_key),
            Some(local_frag_path.clone())
        );

        let frag_states = frag_data_cache.frag_states.lock();
        match frag_states.get(&frag_cache_key).unwrap() {
            super::FragCacheValue::Ready(frag) => {
                assert_eq!(frag.local_frag_path, local_frag_path);
                assert_eq!(frag.ref_count, 1);
            }
            super::FragCacheValue::Downloading(_) => {
                panic!("frag should have been promoted to ready state")
            }
        }
    }

    #[test]
    fn test_remove_downloading_frag_state_if_same_cell_removes_only_when_cell_matches() {
        let backend_storage = Arc::new(MockBackendStorage::new());
        let frag_data_cache = FragDataCache::new(
            backend_storage,
            TempDir::new().unwrap().path().to_str().unwrap().to_string(),
            TempDir::new().unwrap().path().to_str().unwrap().to_string(),
            Duration::from_secs(60),
            ReadableSize::mb(64),
        )
        .unwrap();
        let frag_cache_key = super::FragCacheKey::new("shard_1/frag_1".to_string());
        let download_cell = match frag_data_cache.get_or_insert_frag_state(&frag_cache_key) {
            super::AcquireFragState::Downloading { download_cell } => download_cell,
            super::AcquireFragState::Ready(_) => panic!("frag should not be ready before download"),
        };
        let extra_holder = download_cell.clone();

        assert!(
            frag_data_cache
                .remove_downloading_frag_state_if_same_cell(&frag_cache_key, &download_cell)
        );
        assert!(!contains_frag_state(&frag_data_cache, &frag_cache_key));

        let new_download_cell = match frag_data_cache.get_or_insert_frag_state(&frag_cache_key) {
            super::AcquireFragState::Downloading { download_cell } => download_cell,
            super::AcquireFragState::Ready(_) => panic!("frag should not be ready before download"),
        };
        assert!(!Arc::ptr_eq(&extra_holder, &new_download_cell));
        assert!(
            !frag_data_cache
                .remove_downloading_frag_state_if_same_cell(&frag_cache_key, &download_cell)
        );
        assert!(contains_frag_state(&frag_data_cache, &frag_cache_key));
    }

    #[tokio::test]
    async fn test_load_cached_frags() {
        let shard_meta = test_get_shard_writer_meta();
        let index_meta = TEST_HDFSLOGS_INDEX_META.clone();
        let temp_dir = TempDir::new().unwrap();
        let config = WorkerNodeConfig::new_for_test(temp_dir.path().to_str().unwrap());

        let backend_storage = test_create_backend_storage().await.unwrap();
        let mock_client_pool = Arc::new(MockMetaClientPool::new());
        write_test_data(
            &config,
            &[shard_meta.clone()],
            &index_meta,
            backend_storage.clone(),
            mock_client_pool.clone(),
        )
        .await
        .unwrap();

        let default_reader_cfg = FragReaderConfig {
            disk_reserve_size: ReadableSize::mb(64),
            ..FragReaderConfig::default()
        };
        let frag_data_cache_path = TempDir::new().unwrap().path().to_str().unwrap().to_string();
        let frag_data_cache = Arc::new(
            FragDataCache::new(
                backend_storage.clone(),
                frag_data_cache_path.clone(),
                config.storage.temporary_path(),
                default_reader_cfg.remove_zero_ref_frag_delay.0,
                default_reader_cfg.disk_reserve_size,
            )
            .unwrap(),
        );
        assert_eq!(ready_frag_count(&frag_data_cache), 0);

        let mut mock_client = mock_client_pool.get_meta_client().await.unwrap();
        let (frag_metas, _seq) = mock_client
            .get_frags_by_keyrange(
                shard_meta.shard_header.shard_id,
                &shard_meta.index_key,
                KeyRange::default(),
            )
            .await
            .unwrap();
        let mut _frag_leases = Vec::with_capacity(frag_metas.len());
        for frag_meta in &frag_metas {
            let frag_lease = frag_data_cache
                .acquire_frag(TEST_S3_BUCKET, frag_meta)
                .await
                .unwrap();
            assert!(is_sorted_frag(frag_lease.local_frag_path(), &index_meta));
            _frag_leases.push(frag_lease);
        }
        assert_eq!(ready_frag_count(&frag_data_cache), frag_metas.len());

        let frag_data_cache2 = Arc::new(
            FragDataCache::new(
                backend_storage.clone(),
                frag_data_cache_path.clone(),
                config.storage.temporary_path(),
                default_reader_cfg.remove_zero_ref_frag_delay.0,
                default_reader_cfg.disk_reserve_size,
            )
            .unwrap(),
        );
        let cached_frags = frag_data_cache.frag_states.lock();
        for frag_state in cached_frags.values() {
            match frag_state {
                super::FragCacheValue::Ready(frag) => assert_eq!(frag.ref_count, 1),
                super::FragCacheValue::Downloading(_) => {
                    panic!("download should have completed before assertion")
                }
            }
        }
        let cached_frags2 = frag_data_cache2.frag_states.lock();
        for frag_state in cached_frags2.values() {
            match frag_state {
                super::FragCacheValue::Ready(frag) => assert_eq!(frag.ref_count, 0),
                super::FragCacheValue::Downloading(_) => {
                    panic!("reloaded cache should only contain ready frags")
                }
            }
        }

        let mut s3_frag_paths = cached_frags
            .iter()
            .filter_map(|(frag_cache_key, frag_state)| match frag_state {
                super::FragCacheValue::Ready(_) => Some(frag_cache_key),
                super::FragCacheValue::Downloading(_) => None,
            })
            .collect::<Vec<_>>();
        s3_frag_paths.sort_by(|a, b| a.s3_frag_path.cmp(&b.s3_frag_path));
        let mut s3_frag_paths2 = cached_frags2
            .iter()
            .filter_map(|(frag_cache_key, frag_state)| match frag_state {
                super::FragCacheValue::Ready(_) => Some(frag_cache_key),
                super::FragCacheValue::Downloading(_) => None,
            })
            .collect::<Vec<_>>();
        s3_frag_paths2.sort_by(|a, b| a.s3_frag_path.cmp(&b.s3_frag_path));
        assert_eq!(s3_frag_paths, s3_frag_paths2);
    }

    #[tokio::test]
    async fn test_remove_zero_ref_frags() {
        let shard_meta = test_get_shard_writer_meta();
        let index_meta = TEST_HDFSLOGS_INDEX_META.clone();
        let temp_dir = TempDir::new().unwrap();
        let config = WorkerNodeConfig::new_for_test(temp_dir.path().to_str().unwrap());

        let backend_storage = test_create_backend_storage().await.unwrap();
        let meta_client_pool = Arc::new(MockMetaClientPool::new());

        write_test_data(
            &config,
            &[shard_meta.clone()],
            &index_meta,
            backend_storage.clone(),
            meta_client_pool.clone(),
        )
        .await
        .unwrap();

        let frag_data_cache_path = TempDir::new().unwrap().path().to_str().unwrap().to_string();

        let frag_data_cache = Arc::new(
            FragDataCache::new(
                backend_storage.clone(),
                frag_data_cache_path.clone(),
                config.storage.temporary_path(),
                Duration::from_millis(0),
                ReadableSize::mb(64),
            )
            .unwrap(),
        );
        assert_eq!(ready_frag_count(&frag_data_cache), 0);
        let mut mock_client = meta_client_pool.get_meta_client().await.unwrap();
        let (frag_metas, _seq) = mock_client
            .get_frags_by_keyrange(
                shard_meta.shard_header.shard_id,
                &shard_meta.index_key,
                KeyRange::default(),
            )
            .await
            .unwrap();
        let mut frag_leases = Vec::with_capacity(frag_metas.len());
        for frag_meta in &frag_metas {
            let frag_lease = frag_data_cache
                .acquire_frag(TEST_S3_BUCKET, frag_meta)
                .await
                .unwrap();
            assert!(is_sorted_frag(frag_lease.local_frag_path(), &index_meta));
            frag_leases.push(frag_lease);
        }
        assert_eq!(ready_frag_count(&frag_data_cache), frag_metas.len());
        for frag_state in frag_data_cache.frag_states.lock().values() {
            match frag_state {
                super::FragCacheValue::Ready(frag) => assert_eq!(frag.ref_count, 1),
                super::FragCacheValue::Downloading(_) => {
                    panic!("download should have completed before assertion")
                }
            }
        }
        frag_data_cache.remove_zero_ref_frags();
        assert_eq!(ready_frag_count(&frag_data_cache), frag_metas.len());

        drop(frag_leases);
        frag_data_cache.remove_zero_ref_frags();
        assert_eq!(ready_frag_count(&frag_data_cache), 0);
        assert!(
            s3_frag::list_all_files_in_absolute(&frag_data_cache_path)
                .unwrap()
                .is_empty()
        );
    }

    #[tokio::test]
    async fn test_acquire_frag_avoids_duplicate_downloads() {
        let backend_storage = Arc::new(MockBackendStorage::new().with_default_get_object(
            create_multi_file_object_data().await,
            Duration::from_millis(100),
        ));
        let frag_data_cache_path = TempDir::new().unwrap().path().to_str().unwrap().to_string();
        let temporary_path = TempDir::new().unwrap().path().to_str().unwrap().to_string();
        let frag_data_cache = Arc::new(
            FragDataCache::new(
                backend_storage.clone(),
                frag_data_cache_path,
                temporary_path,
                Duration::from_secs(60),
                ReadableSize::mb(64),
            )
            .unwrap(),
        );
        let frag_meta = FragMeta::new(
            0,
            TEST_FRAG_DATA_CACHE_FRAG_PATH,
            FragProperty {
                size: backend_storage.object_len() as u64,
                split_handles: vec![vec![]],
                ..Default::default()
            },
            vec![FragStorageRef {
                sub_key: TEST_FRAG_DATA_CACHE_STORAGE_REF_KEY.to_string(),
                offset: 0,
                size: backend_storage.object_len() as u64,
            }],
            1,
        )
        .unwrap();

        let frag_data_cache_2 = frag_data_cache.clone();
        let frag_meta_2 = frag_meta.clone();
        let (frag_path_1, frag_path_2) = tokio::join!(
            frag_data_cache.acquire_frag(TEST_S3_BUCKET, &frag_meta),
            frag_data_cache_2.acquire_frag(TEST_S3_BUCKET, &frag_meta_2)
        );
        let frag_lease_1 = frag_path_1.unwrap();
        let frag_lease_2 = frag_path_2.unwrap();

        assert_eq!(
            frag_lease_1.local_frag_path(),
            frag_lease_2.local_frag_path()
        );
        assert_eq!(backend_storage.get_object_count(), 1);
        assert!(
            std::path::Path::new(frag_lease_1.local_frag_path())
                .join("index/meta.json")
                .exists()
        );
        let frag_states = frag_data_cache.frag_states.lock();
        match frag_states.get(&super::FragCacheKey::new(frag_meta.frag_path.clone())) {
            Some(super::FragCacheValue::Ready(frag)) => assert_eq!(frag.ref_count, 2),
            Some(super::FragCacheValue::Downloading(_)) => {
                panic!("frag should be ready after acquire succeeds")
            }
            None => panic!("frag should exist in cache after acquire succeeds"),
        }
    }

    #[tokio::test]
    async fn test_acquire_frag_removes_failed_download_cell_and_retries() {
        let backend_storage = Arc::new(MockBackendStorage::new().with_failing_default_get_object(
            create_multi_file_object_data().await,
            Duration::from_millis(100),
            1,
        ));
        let frag_data_cache_path = TempDir::new().unwrap().path().to_str().unwrap().to_string();
        let temporary_path = TempDir::new().unwrap().path().to_str().unwrap().to_string();
        let frag_data_cache = Arc::new(
            FragDataCache::new(
                backend_storage.clone(),
                frag_data_cache_path,
                temporary_path,
                Duration::from_secs(60),
                ReadableSize::mb(64),
            )
            .unwrap(),
        );
        let frag_meta = FragMeta {
            frag_path: TEST_FRAG_DATA_CACHE_FRAG_PATH.to_string(),
            property: Some(FragProperty {
                size: backend_storage.object_len() as u64,
                split_handles: vec![vec![]],
                ..Default::default()
            }),
            storage_refs: vec![FragStorageRef {
                sub_key: TEST_FRAG_DATA_CACHE_STORAGE_REF_KEY.to_string(),
                offset: 0,
                size: backend_storage.object_len() as u64,
            }],
            ..Default::default()
        };
        let frag_cache_key = super::FragCacheKey::new(frag_meta.frag_path.clone());

        let frag_data_cache_2 = frag_data_cache.clone();
        let frag_meta_2 = frag_meta.clone();
        let (first_result, second_result) = tokio::join!(
            frag_data_cache.acquire_frag(TEST_S3_BUCKET, &frag_meta),
            frag_data_cache_2.acquire_frag(TEST_S3_BUCKET, &frag_meta_2)
        );
        assert!(first_result.is_err());
        assert!(second_result.is_err());
        assert_eq!(backend_storage.get_object_count(), 1);
        assert!(
            !contains_frag_state(&frag_data_cache, &frag_cache_key),
            "failed download state should be removed after the last waiter returns"
        );

        let frag_lease = frag_data_cache
            .acquire_frag(TEST_S3_BUCKET, &frag_meta)
            .await
            .unwrap();
        assert_eq!(backend_storage.get_object_count(), 2);
        assert_eq!(
            frag_lease.local_frag_path(),
            frag_data_cache.local_frag_path(&frag_meta.frag_path)
        );

        let frag_states = frag_data_cache.frag_states.lock();
        match frag_states.get(&frag_cache_key) {
            Some(super::FragCacheValue::Ready(frag)) => assert_eq!(frag.ref_count, 1),
            Some(super::FragCacheValue::Downloading(_)) => {
                panic!("frag should be ready after retry succeeds")
            }
            None => panic!("frag should exist in cache after retry succeeds"),
        }
    }

    #[tokio::test]
    async fn test_from_frag_data_cache_releases_lease_on_missing_property() {
        let shard_meta = test_get_shard_writer_meta();
        let index_meta = TEST_HDFSLOGS_INDEX_META.clone();
        let temp_dir = TempDir::new().unwrap();
        let config = WorkerNodeConfig::new_for_test(temp_dir.path().to_str().unwrap());

        let backend_storage = test_create_backend_storage().await.unwrap();
        let meta_client_pool = Arc::new(MockMetaClientPool::new());
        write_test_data(
            &config,
            std::slice::from_ref(&shard_meta),
            &index_meta,
            backend_storage.clone(),
            meta_client_pool.clone(),
        )
        .await
        .unwrap();

        let frag_data_cache_dir = TempDir::new().unwrap();
        let frag_data_cache = Arc::new(
            FragDataCache::new(
                backend_storage,
                frag_data_cache_dir.path().to_str().unwrap().to_string(),
                config.storage.temporary_path(),
                Duration::from_secs(60),
                ReadableSize::mb(64),
            )
            .unwrap(),
        );

        let mut mock_client = meta_client_pool.get_meta_client().await.unwrap();
        let (frag_metas, _seq) = mock_client
            .get_frags_by_keyrange(
                shard_meta.shard_header.shard_id,
                &shard_meta.index_key,
                KeyRange::default(),
            )
            .await
            .unwrap();
        let frag_meta = frag_metas.into_iter().next().unwrap();

        let _lease = frag_data_cache
            .acquire_frag(TEST_S3_BUCKET, &frag_meta)
            .await
            .unwrap();

        let mut missing_property_frag_meta = frag_meta.clone();
        missing_property_frag_meta.property = None;
        let err = match FragReaderMeta::from_frag_data_cache(
            &missing_property_frag_meta,
            frag_data_cache.clone(),
            TEST_S3_BUCKET,
        )
        .await
        {
            Ok(_) => panic!("missing property should fail validation"),
            Err(err) => err,
        };
        assert!(
            err.to_string().contains("Fragment property is required"),
            "unexpected error: {err:?}"
        );

        let frag_states = frag_data_cache.frag_states.lock();
        match frag_states.get(&super::FragCacheKey::new(frag_meta.frag_path.clone())) {
            Some(super::FragCacheValue::Ready(frag)) => assert_eq!(frag.ref_count, 1),
            Some(super::FragCacheValue::Downloading(_)) => {
                panic!("frag should stay ready after failed metadata validation")
            }
            None => panic!("frag should remain cached after failed metadata validation"),
        }
    }
}