// Copyright 2026 PingCAP, Inc.

use std::sync::Arc;

use bytes::Bytes;
use dashmap::DashSet;

use crate::dfs::{Dfs, Error, FileType, Options, Result};

/// An DFS with an "overlay".
///
/// New writes will be put to `overlay` FS.
/// Read will firstly send to the `overlay` fs then `base` fs.
///
/// This can be used to isolate new write files and existing files.
/// Useful when performing light-weight forking write opreations like packing
/// backup.
pub struct OverlaidFs {
    base: Arc<dyn Dfs>,
    overlay: Arc<dyn Dfs>,
    overlay_written: DashSet<(FileType, u64)>,
}

impl OverlaidFs {
    /// Create a new OverlaidFs with the given base and overlay DFS.
    pub fn new(base: Arc<dyn Dfs>, overlay: Arc<dyn Dfs>) -> Self {
        Self {
            base,
            overlay,
            overlay_written: Default::default(),
        }
    }

    /// Get the base DFS.
    pub fn base_fs(&self) -> &Arc<dyn Dfs> {
        &self.base
    }

    /// Get the overlay DFS.
    pub fn overlay_fs(&self) -> &Arc<dyn Dfs> {
        &self.overlay
    }

    /// Check if a file is modified in the overlay layer by the `OverlaidFs`.
    ///
    /// Notice: This only check newly written files and if a file was already
    /// existing when creating this `OverlaidFs`, it won't be counted as a
    /// "modified" file.
    pub fn file_is_modified(&self, file_type: FileType, file_id: u64) -> bool {
        self.overlay_written.contains(&(file_type, file_id))
    }

    /// Get a list of modified files in the overlay layer by the `OverlaidFs`.
    ///
    /// Notice: This only check newly written files and if a file was already
    /// existing when creating this `OverlaidFs`, it won't be counted as a
    /// "modified" file.
    pub fn modified_files(&self) -> Vec<(FileType, u64)> {
        self.overlay_written.iter().map(|v| *v).collect()
    }
}

#[async_trait::async_trait]
impl Dfs for OverlaidFs {
    /// exists checks if the file exists in the DFS.
    async fn exists(&self, file_id: u64, opts: Options) -> Result<bool> {
        // Check in the write layer first, then the base layer.
        let res = self.overlay.exists(file_id, opts).await;
        match res {
            Ok(false) | Err(Error::NoSuchKey(..) | Error::NotExists(..)) => {
                self.base.exists(file_id, opts).await
            }
            any_other => any_other,
        }
    }

    /// read_file reads the whole file to memory.
    /// It can be used by remote compaction server that doesn't have local disk.
    async fn read_file(&self, file_id: u64, opts: Options) -> Result<Bytes> {
        // Try to read from the write layer first, then the base layer.
        let res = self.overlay.read_file(file_id, opts).await;
        match res {
            Err(Error::NoSuchKey(..) | Error::NotExists(..)) => {
                self.base.read_file(file_id, opts).await
            }
            any_other => any_other,
        }
    }

    /// Create creates a new File.
    /// The shard_id and shard_ver can be used determine where to write the
    /// file.
    async fn create(&self, file_id: u64, data: Bytes, opts: Options) -> Result<()> {
        self.overlay.create(file_id, data.clone(), opts).await?;
        self.overlay_written.insert((opts.file_type, file_id));
        Ok(())
    }

    async fn size(&self, file_id: u64, opts: Options) -> Result<u64> {
        let res = self.overlay.size(file_id, opts).await;
        match res {
            Err(Error::NoSuchKey(..) | Error::NotExists(..)) => self.base.size(file_id, opts).await,
            any_other => any_other,
        }
    }

    /// remove removes the file from the DFS.
    /// `file_len` is used to choose proper storage class for cost
    /// efficiency, in bytes.
    /// `None` if file_len is unknown when invoke this method.
    async fn remove(&self, _file_id: u64, _file_len: Option<u64>, _opts: Options) {
        panic!("Call `remove` in OverlayS3Fs isn't allowed.");
    }

    /// Remove the file from DFS permanently.
    async fn permanently_remove(&self, _file_id: u64, _opts: Options) -> Result<()> {
        panic!("Call `permanently_remove` in OverlayS3Fs isn't allowed.");
    }

    /// get_runtime gets the tokio runtime for the DFS.
    fn get_runtime(&self) -> &tokio::runtime::Runtime {
        self.base.get_runtime()
    }

    fn get_prefix(&self) -> String {
        self.overlay.get_prefix()
    }

    fn get_storage_name(&self) -> String {
        self.base.get_storage_name()
    }
}

// Test cases are generated by Copilot, with Claude 4 Sonnet.
#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use bytes::Bytes;

    use super::*;
    use crate::dfs::{FileType, Options, test_util::new_test_s3fs};

    #[test]
    fn test_overlay_fs_read_file() {
        ::test_util::init_log_for_test();

        let base_data = "base_layer_data".to_string().into_bytes();
        let write_data = "write_layer_data".to_string().into_bytes();
        // Create mock instances that return the appropriate data
        let base_s3fs: Arc<dyn Dfs> = Arc::new(new_test_s3fs(&base_data));
        let write_s3fs: Arc<dyn Dfs> = Arc::new(new_test_s3fs(&write_data));
        let overlay_fs = OverlaidFs::new(base_s3fs.clone(), write_s3fs.clone());

        let file_id1 = 456u64;
        let file_id2 = 457u64;
        let opts = Options::default();
        let runtime = overlay_fs.get_runtime();

        // Test reading non-existent file
        let result = runtime.block_on(overlay_fs.read_file(file_id1, opts));
        result.unwrap_err();

        // Create file only in base layer
        runtime
            .block_on(base_s3fs.create(file_id1, Bytes::from(base_data.clone()), opts))
            .unwrap();

        // Should read from base layer (mock returns the base_data)
        let result = runtime.block_on(overlay_fs.read_file(file_id1, opts));
        assert_eq!(result.unwrap(), Bytes::from(base_data));

        // Create file in write layer
        runtime
            .block_on(overlay_fs.create(file_id2, Bytes::from(write_data.clone()), opts))
            .unwrap();

        // Should read from write layer (mock returns the write_data)
        let result = runtime.block_on(overlay_fs.read_file(file_id2, opts));
        assert_eq!(result.unwrap(), Bytes::from(write_data));
    }

    #[test]
    fn test_overlay_fs_create() {
        ::test_util::init_log_for_test();

        let file_data = "overlay_create_test".to_string().into_bytes();
        let base_s3fs: Arc<dyn Dfs> = Arc::new(new_test_s3fs(&file_data));
        let write_s3fs: Arc<dyn Dfs> = Arc::new(new_test_s3fs(&file_data));
        let overlay_fs = OverlaidFs::new(base_s3fs.clone(), write_s3fs.clone());

        let file_id = 789u64;
        let opts = Options::default();
        let runtime = overlay_fs.get_runtime();

        assert!(!overlay_fs.file_is_modified(FileType::Sst, file_id));
        // Create file through overlay
        let create_data = Bytes::from(file_data.clone());
        runtime
            .block_on(overlay_fs.create(file_id, create_data.clone(), opts))
            .unwrap();

        // File should exist in write layer
        let write_exists = runtime.block_on(write_s3fs.exists(file_id, opts));
        assert!(write_exists.unwrap());

        // File should not exist in base layer
        let base_exists = runtime.block_on(base_s3fs.exists(file_id, opts));
        assert!(!base_exists.unwrap());

        // Reading through overlay should work (returns mock data)
        let read_result = runtime
            .block_on(overlay_fs.read_file(file_id, opts))
            .unwrap();
        assert_eq!(read_result, create_data);
        assert!(overlay_fs.file_is_modified(FileType::Sst, file_id));
    }

    #[test]
    fn test_overlay_fs_get_storage_name_delegates_to_base() {
        // get_storage_name() must return the *base* DFS's storage name, not overlay's.
        // For pack_backup, base = source account's container, overlay = tmpfs prefix.
        // The manifest's content_bucket field uses base_fs().get_storage_name().
        let base_data = "base".as_bytes().to_vec();
        let overlay_data = "overlay".as_bytes().to_vec();
        let base: Arc<dyn Dfs> = Arc::new(new_test_s3fs(&base_data));
        let overlay: Arc<dyn Dfs> = Arc::new(new_test_s3fs(&overlay_data));

        // Both test S3Fs instances have bucket="shard-db" by default.
        // Verify OverlaidFs delegates to base.
        let overlaid = OverlaidFs::new(base.clone(), overlay.clone());
        assert_eq!(overlaid.get_storage_name(), base.get_storage_name());
        assert_eq!(overlaid.get_storage_name(), "shard-db");
    }

    #[test]
    fn test_overlay_fs_base_fs_and_overlay_fs_accessors() {
        let base_data = "base".as_bytes().to_vec();
        let overlay_data = "overlay".as_bytes().to_vec();
        let base: Arc<dyn Dfs> = Arc::new(new_test_s3fs(&base_data));
        let overlay: Arc<dyn Dfs> = Arc::new(new_test_s3fs(&overlay_data));
        let overlaid = OverlaidFs::new(base.clone(), overlay.clone());

        // base_fs() returns the base DFS
        assert_eq!(overlaid.base_fs().get_prefix(), base.get_prefix());
        // overlay_fs() returns the overlay DFS
        assert_eq!(overlaid.overlay_fs().get_prefix(), overlay.get_prefix());
        // get_prefix() delegates to overlay (new writes go to overlay prefix)
        assert_eq!(overlaid.get_prefix(), overlay.get_prefix());
    }

    #[test]
    fn test_overlay_fs_write_tracked_in_modified_files() {
        let data = "test".as_bytes().to_vec();
        let base: Arc<dyn Dfs> = Arc::new(new_test_s3fs(&data));
        let overlay: Arc<dyn Dfs> = Arc::new(new_test_s3fs(&data));
        let overlaid = OverlaidFs::new(base, overlay);
        let runtime = overlaid.get_runtime();

        let file_id = 42u64;
        let opts = Options {
            file_type: FileType::Sst,
            ..Default::default()
        };

        // Initially no modified files
        assert!(overlaid.modified_files().is_empty());
        assert!(!overlaid.file_is_modified(FileType::Sst, file_id));

        // Create through overlay
        runtime
            .block_on(overlaid.create(file_id, Bytes::from_static(b"data"), opts))
            .unwrap();

        // Now tracked
        assert!(overlaid.file_is_modified(FileType::Sst, file_id));
        assert_eq!(overlaid.modified_files(), vec![(FileType::Sst, file_id)]);
        // Different file type is not tracked
        assert!(!overlaid.file_is_modified(FileType::Blob, file_id));
    }

    #[test]
    #[should_panic(expected = "Call `remove` in OverlayS3Fs isn't allowed.")]
    fn test_overlay_fs_remove_panics() {
        let file_data = "panic_test_data".to_string().into_bytes();
        let base_s3fs: Arc<dyn Dfs> = Arc::new(new_test_s3fs(&file_data));
        let write_s3fs: Arc<dyn Dfs> = Arc::new(new_test_s3fs(&file_data));
        let overlay_fs = OverlaidFs::new(base_s3fs, write_s3fs);

        let file_id = 999u64;
        let opts = Options::default();
        let runtime = overlay_fs.get_runtime();

        // This should panic
        runtime.block_on(overlay_fs.remove(file_id, None, opts));
    }
}