// Copyright 2026 PingCAP, Inc.

use std::{
    fmt::{Debug, Formatter},
    future::Future,
    io::Write,
    num::NonZeroU32,
    path::Path,
    sync::Arc,
    time::Duration,
};

use async_trait::async_trait;
use azure_core::{
    ClientOptions, Context, Error as AzureError, ExponentialRetryOptions, Policy, PolicyResult,
    Request, RetryOptions,
    auth::TokenCredential,
    error::ErrorKind,
    prelude::{Delimiter, MaxResults, Prefix, Range},
};
use azure_identity::DefaultAzureCredential;
use azure_storage::{CloudLocation, ConnectionString, EndpointProtocol, StorageCredentials};
use azure_storage_blobs::prelude::{AccessTier, BlobClient, ClientBuilder, ContainerClient, Tags};
use bytes::{BufMut, Bytes, BytesMut};
use engine_traits::{GetObjectOptions, ListObjectContent, ObjectCacheWithHook, ObjectStorage};
use futures::StreamExt;
use http::StatusCode;
use tikv_util::time::Instant;
use tokio::runtime::Runtime;

use crate::dfs::{
    Dfs, Error, KVENGINE_DFS_LATENCY, KVENGINE_DFS_THROUGHPUT, MetricsFileType, Options,
    ReservableWriter, STORAGE_CLASS_DEFAULT, STORAGE_CLASS_GLACIER_IR,
    STORAGE_CLASS_INTELLIGENT_TIERING, STORAGE_CLASS_STANDARD, STORAGE_CLASS_STANDARD_IA,
    StorageClass,
    config::{AzureConfig, Config, ConnOptions},
    get_file_type_from_key,
};

#[cfg(test)]
const DEVELOPMENT_STORAGE_CONNECTION_STRING: &str = "UseDevelopmentStorage=true";
const TAG_DELETED_KEY: &str = "deleted";
const MAX_DELAY_DURATION: Duration = Duration::from_secs(600);
const AZURE_CLIENT_ID_ENV: &str = "AZURE_CLIENT_ID";
const AZURE_START_FROM_QUERY_KEY: &str = "startFrom";
const AZURE_START_FROM_API_VERSION: &str = "2023-05-03";

#[derive(Clone, Debug)]
struct AzureListStartFrom(String);

#[derive(Clone, Debug, Default)]
struct AzureListStartFromPolicy;

#[async_trait]
impl Policy for AzureListStartFromPolicy {
    async fn send(
        &self,
        ctx: &Context,
        request: &mut Request,
        next: &[Arc<dyn Policy>],
    ) -> PolicyResult {
        if let Some(start_from) = ctx.get::<AzureListStartFrom>() {
            request
                .url_mut()
                .query_pairs_mut()
                .append_pair(AZURE_START_FROM_QUERY_KEY, &start_from.0);
            request.insert_header(azure_core::headers::VERSION, AZURE_START_FROM_API_VERSION);
        }
        next[0].send(ctx, request, &next[1..]).await
    }
}

#[derive(Clone, Debug)]
struct AzureCopySourceAuth(String);

#[derive(Clone, Debug, Default)]
struct AzureCopySourceAuthPolicy;

#[async_trait]
impl Policy for AzureCopySourceAuthPolicy {
    async fn send(
        &self,
        ctx: &Context,
        request: &mut Request,
        next: &[Arc<dyn Policy>],
    ) -> PolicyResult {
        if let Some(auth) = ctx.get::<AzureCopySourceAuth>() {
            request.insert_header(
                azure_core::headers::HeaderName::from_static("x-ms-copy-source-authorization"),
                &auth.0,
            );
        }
        next[0].send(ctx, request, &next[1..]).await
    }
}

#[derive(Clone)]
pub struct AzureFs {
    core: Arc<AzureFsCore>,
}

impl AzureFs {
    pub fn new_from_config(conf: Config) -> Self {
        let core = AzureFsCore::new(conf).unwrap_or_else(|err| {
            panic!("failed to initialize azure dfs: {:?}", err);
        });
        Self {
            core: Arc::new(core),
        }
    }
}

impl std::ops::Deref for AzureFs {
    type Target = AzureFsCore;

    fn deref(&self) -> &Self::Target {
        &self.core
    }
}

pub struct AzureFsCore {
    container_client: Arc<ContainerClient>,
    prefix: String,
    runtime: Option<Runtime>,
    read_only: bool,
    opts: ConnOptions,
    exotic_endpoint: String,
    exotic_token_credential: Option<Arc<dyn TokenCredential>>,
}

impl AzureFsCore {
    fn new(conf: Config) -> crate::dfs::Result<Self> {
        let Config {
            azure,
            mut prefix,
            read_only,
            conn_options,
            ..
        } = conf;
        if azure.container.is_empty() {
            return Err(Error::Other(
                "azure container is required for dfs azure backend".to_string(),
            ));
        }

        // Always create a dedicated runtime for synchronous DFS operations
        // (get_object, list, etc.) which call get_runtime().block_on().
        let runtime = Some(
            tokio::runtime::Builder::new_multi_thread()
                .worker_threads(2)
                .enable_all()
                .thread_name("azure")
                .build()
                .unwrap(),
        );

        // Azure cross-region latency is much higher than same-region S3 (~1s vs
        // <10ms). Increase dispatch_timeout to 30s to avoid spurious timeouts
        // under concurrent flush load.
        let mut conn_options = conn_options;
        let min_azure_dispatch_timeout = Duration::from_secs(30);
        if conn_options.dispatch_timeout.0 < min_azure_dispatch_timeout {
            conn_options.dispatch_timeout.0 = min_azure_dispatch_timeout;
        }

        let retry_options = Self::retry_options(&conn_options);
        let exotic_endpoint = azure.exotic_endpoint.clone();
        let container_client = Self::build_container_client(&azure, &retry_options)?;
        let exotic_token_credential: Option<Arc<dyn TokenCredential>> =
            if !exotic_endpoint.is_empty() {
                Some(Arc::new(DefaultAzureCredential::default()))
            } else {
                None
            };

        if prefix.is_empty() {
            prefix.push_str("default");
        }

        let core = Self {
            container_client,
            prefix,
            runtime,
            read_only,
            opts: conn_options,
            exotic_endpoint,
            exotic_token_credential,
        };
        core.ensure_container()?;
        Ok(core)
    }

    fn retry_options(opts: &ConnOptions) -> RetryOptions {
        RetryOptions::exponential(
            ExponentialRetryOptions::default()
                .initial_delay(opts.retry_sleep_interval.0)
                .max_retries(opts.max_retry_count)
                .max_delay(MAX_DELAY_DURATION),
        )
    }

    fn client_options(retry_options: &RetryOptions) -> ClientOptions {
        let mut client_options = ClientOptions::default().retry(retry_options.clone());
        client_options
            .per_call_policies_mut()
            .push(Arc::new(AzureListStartFromPolicy));
        client_options
            .per_call_policies_mut()
            .push(Arc::new(AzureCopySourceAuthPolicy));
        client_options
    }

    fn build_container_client(
        azure: &AzureConfig,
        retry_options: &RetryOptions,
    ) -> crate::dfs::Result<Arc<ContainerClient>> {
        let container = azure.container.clone();

        // Try connection string authentication first.
        // Skip when endpoint is explicitly set: an explicit endpoint means the
        // caller wants to reach a different account, so the connection_string
        // (which encodes the *original* account credentials) must not take effect.
        if !azure.connection_string.is_empty() && azure.endpoint.is_empty() {
            let connection_string = ConnectionString::new(azure.connection_string.as_str())
                .map_err(|e| Error::Other(format!("azure connection string error: {}", e)))?;
            if connection_string.use_development_storage == Some(true) {
                let options = Self::client_options(retry_options);
                let client = Arc::new(
                    ClientBuilder::emulator()
                        .client_options(options)
                        .container_client(container),
                );
                return Ok(client);
            }
            let account_name = connection_string.account_name.ok_or_else(|| {
                Error::Other("azure connection string missing account name".to_string())
            })?;
            let credentials = connection_string
                .storage_credentials()
                .map_err(|e| Error::Other(format!("azure connection string error: {}", e)))?;
            let cloud_location =
                Self::cloud_location_from_connection_string(&connection_string, account_name);
            let options = Self::client_options(retry_options);
            let container_client = ClientBuilder::with_location(cloud_location, credentials)
                .client_options(options)
                .container_client(container);
            return Ok(Arc::new(container_client));
        }

        // Try account name/key authentication.
        // Skip when AZURE_CLIENT_ID is set so Managed Identity takes priority:
        // cross-account operations need a token credential for the Bearer header
        // in raw_copy_object, and MI provides one while account_key does not.
        if !azure.account_name.is_empty()
            && !azure.account_key.is_empty()
            && std::env::var(AZURE_CLIENT_ID_ENV).is_err()
        {
            let credentials = StorageCredentials::access_key(
                azure.account_name.clone(),
                azure.account_key.clone(),
            );
            let cloud_location = if azure.endpoint.is_empty() {
                CloudLocation::Public {
                    account: azure.account_name.clone(),
                }
            } else {
                CloudLocation::Custom {
                    account: azure.account_name.clone(),
                    uri: azure.endpoint.clone(),
                }
            };
            let options = Self::client_options(retry_options);
            let container_client = ClientBuilder::with_location(cloud_location, credentials)
                .client_options(options)
                .container_client(container);
            return Ok(Arc::new(container_client));
        }

        // Fallback to managed identity authentication if account_name is provided
        if !azure.account_name.is_empty() {
            info!("no credentials provided, attempting Azure Managed Identity authentication");

            let cloud_location = if azure.endpoint.is_empty() {
                CloudLocation::Public {
                    account: azure.account_name.clone(),
                }
            } else {
                info!("using custom endpoint: {}", azure.endpoint);
                CloudLocation::Custom {
                    account: azure.account_name.clone(),
                    uri: azure.endpoint.clone(),
                }
            };

            let azure_client_id = std::env::var(AZURE_CLIENT_ID_ENV).ok();
            if let Some(ref client_id) = azure_client_id {
                info!(
                    "AZURE_CLIENT_ID is set, will use user-assigned managed identity: {}",
                    client_id
                );
            } else {
                info!("AZURE_CLIENT_ID not set, will use system-assigned managed identity");
            }

            let credential: Arc<dyn TokenCredential> = Arc::new(DefaultAzureCredential::default());
            let credentials = StorageCredentials::token_credential(credential);

            let options = Self::client_options(retry_options);
            let container_client = ClientBuilder::with_location(cloud_location, credentials)
                .client_options(options)
                .container_client(container);
            return Ok(Arc::new(container_client));
        }

        Err(Error::Other(
            "azure dfs backend requires one of: connection string, account name/key, or account name (for managed identity)".to_string(),
        ))
    }

    fn cloud_location_from_connection_string(
        connection_string: &ConnectionString<'_>,
        account_name: &str,
    ) -> CloudLocation {
        if let Some(blob_endpoint) = connection_string.blob_endpoint {
            return CloudLocation::Custom {
                account: account_name.to_string(),
                uri: blob_endpoint.to_string(),
            };
        }
        if connection_string.endpoint_suffix.is_some()
            || connection_string.default_endpoints_protocol.is_some()
        {
            let protocol = match connection_string
                .default_endpoints_protocol
                .as_ref()
                .unwrap_or(&EndpointProtocol::Https)
            {
                EndpointProtocol::Http => "http",
                EndpointProtocol::Https => "https",
            };
            let endpoint_suffix = connection_string
                .endpoint_suffix
                .unwrap_or("core.windows.net");
            return CloudLocation::Custom {
                account: account_name.to_string(),
                uri: format!("{protocol}://{account_name}.blob.{endpoint_suffix}"),
            };
        }
        CloudLocation::Public {
            account: account_name.to_string(),
        }
    }

    async fn ensure_container_async(&self) -> crate::dfs::Result<()> {
        let container_name = self.container_client.container_name().to_string();
        let timeout = self.opts.dispatch_timeout.0;

        match tokio::time::timeout(
            timeout,
            self.container_client.get_properties().into_future(),
        )
        .await
        {
            Ok(Ok(_)) => Ok(()),
            Ok(Err(err)) if is_status_code(&err, StatusCode::NOT_FOUND) => Err(Error::Other(
                format!("azure container {} does not exist", container_name),
            )),
            Ok(Err(err)) => Err(Error::Other(format!(
                "azure container {} check error: {}",
                container_name, err
            ))),
            Err(_) => Err(Error::Other(format!(
                "azure container {} check timeout after {:?}",
                container_name, timeout
            ))),
        }
    }

    fn ensure_container(&self) -> crate::dfs::Result<()> {
        let rt = self.runtime.as_ref().unwrap();
        if tokio::runtime::Handle::try_current().is_ok() {
            tokio::task::block_in_place(|| rt.handle().block_on(self.ensure_container_async()))
        } else {
            rt.block_on(self.ensure_container_async())
        }
    }

    fn blob_client(&self, key: &str) -> BlobClient {
        self.container_client.blob_client(key)
    }

    fn source_endpoint_url(&self) -> crate::dfs::Result<azure_core::Url> {
        if !self.exotic_endpoint.is_empty() {
            return azure_core::Url::parse(self.exotic_endpoint.trim_end_matches('/'))
                .map_err(|e| Error::Other(format!("raw_copy_object source endpoint: {}", e)));
        }

        self.container_client
            .service_client()
            .url()
            .map_err(|e| Error::Other(format!("raw_copy_object source endpoint: {}", e)))
    }

    fn source_url(&self, source_key: &str) -> crate::dfs::Result<azure_core::Url> {
        let endpoint = self.source_endpoint_url()?;
        build_source_url(&endpoint, source_key)
    }

    async fn raw_copy_object_inner(
        &self,
        source_key: &str,
        target_key: &str,
    ) -> crate::dfs::Result<()> {
        if self.read_only {
            return Err(Error::Other("raw_copy_object: read only".into()));
        }
        let target_blob = self.blob_client(target_key);

        let is_cross_account = !self.exotic_endpoint.is_empty();
        let source_url = self.source_url(source_key)?;

        let mut ctx = Context::new();
        if is_cross_account {
            debug!(
                "cross-account copy: {} -> {} via exotic endpoint {}",
                source_key, target_key, self.exotic_endpoint
            );
            if let Some(ref cred) = self.exotic_token_credential {
                let token = cred
                    .get_token(&["https://storage.azure.com/.default"])
                    .await
                    .map_err(|e| Error::Other(format!("raw_copy_object get_token: {}", e)))?;
                ctx.insert(AzureCopySourceAuth(format!(
                    "Bearer {}",
                    token.token.secret()
                )));
            } else {
                warn!(
                    "cross-account copy without exotic token credential, source may reject the request"
                );
            }
        }

        self.timeout_or_io_err(self.opts.dispatch_timeout.0, async {
            target_blob
                .copy_from_url(source_url)
                .is_synchronous(true)
                .context(ctx)
                .await
        })
        .await
        .map_err(|e| Error::Other(format!("raw_copy_object: {}", e)))?;
        Ok(())
    }

    async fn timeout_or_io_err<T, F>(&self, timeout: Duration, fut: F) -> Result<T, AzureError>
    where
        F: Future<Output = Result<T, AzureError>>,
    {
        match tokio::time::timeout(timeout, fut).await {
            Ok(result) => result,
            Err(_) => Err(AzureError::message(
                ErrorKind::Io,
                format!("azure request timeout after {:?}", timeout),
            )),
        }
    }

    async fn object_size_inner(&self, key: &str) -> crate::dfs::Result<u64> {
        let result = self
            .timeout_or_io_err(self.opts.dispatch_timeout.0, async {
                self.blob_client(key).get_properties().await
            })
            .await;
        match result {
            Ok(resp) => Ok(resp.blob.properties.content_length),
            Err(err) if is_status_code(&err, StatusCode::NOT_FOUND) => {
                Err(Error::NoSuchKey(key.to_string()))
            }
            Err(err) => Err(Error::Other(format!("azure head error: {}", err))),
        }
    }

    async fn get_range(
        &self,
        key: &str,
        opts: GetObjectOptions,
    ) -> crate::dfs::Result<Option<Range>> {
        if opts.is_full_range() {
            return Ok(None);
        }
        if let (Some(start_off), Some(end_off)) = (opts.start_off, opts.end_off) {
            return Ok(Some(Range::new(start_off, end_off)));
        }
        let range = match (opts.start_off, opts.end_off) {
            (Some(start_off), None) => (start_off..).into(),
            (None, Some(len)) => {
                let size = self.object_size_inner(key).await?;
                let start = size.saturating_sub(len);
                Range::new(start, size)
            }
            _ => return Ok(None),
        };
        Ok(Some(range))
    }

    pub async fn get_object_to_writer<W, F>(
        &self,
        key: String,
        _file_name: String,
        opts: GetObjectOptions,
        build_writer: F,
    ) -> crate::dfs::Result<(W, u64)>
    where
        W: ReservableWriter + Send + 'static,
        F: Fn() -> std::io::Result<W> + 'static,
    {
        let start_time = Instant::now_coarse();
        let file_type = get_file_type_from_key(&key);
        let range = self.get_range(&key, opts).await?;
        let blob_client = self.blob_client(&key);
        let builder = if let Some(range) = range {
            blob_client.get().range(range)
        } else {
            blob_client.get()
        };
        let mut stream = builder.into_stream();
        let mut writer = build_writer()?;
        let mut read_len = 0;
        let mut err = None;

        loop {
            let response = match self
                .timeout_or_io_err(self.opts.dispatch_timeout.0, async {
                    stream.next().await.transpose()
                })
                .await
            {
                Ok(Some(response)) => response,
                Ok(None) => break,
                Err(e) => {
                    err = Some(e);
                    break;
                }
            };

            let content_len = response.blob.properties.content_length;
            writer.reserve_capacity(content_len);
            let mut body = response.data;
            let response_len = {
                let writer = &mut writer;
                self.timeout_or_io_err(self.opts.read_body_timeout.0, async move {
                    let mut response_len = 0;
                    while let Some(chunk) = body.next().await.transpose()? {
                        response_len += chunk.len() as u64;
                        writer.write_all(&chunk).map_err(|err| {
                            AzureError::message(ErrorKind::Io, format!("write_all: {}", err))
                        })?;
                    }
                    Ok(response_len)
                })
                .await
            };
            match response_len {
                Ok(response_len) => {
                    read_len += response_len;
                }
                Err(e) => {
                    err = Some(e);
                    break;
                }
            }
        }

        if let Some(err) = err {
            if is_status_code(&err, StatusCode::NOT_FOUND) {
                return Err(Error::NoSuchKey(key));
            }
            return Err(Error::Other(format!("azure get error: {}", err)));
        }

        KVENGINE_DFS_THROUGHPUT
            .azure
            .get
            .get(MetricsFileType::from(file_type))
            .inc_by(read_len);
        KVENGINE_DFS_LATENCY
            .get
            .get(MetricsFileType::from(file_type))
            .observe(start_time.saturating_elapsed().as_secs_f64());
        Ok((writer, read_len))
    }

    async fn put_object_with_tier(
        &self,
        key: String,
        data: Bytes,
        access_tier: Option<AccessTier>,
    ) -> crate::dfs::Result<()> {
        if self.read_only {
            return Err(Error::ReadOnly);
        }
        let start_time = Instant::now_coarse();
        let file_type = get_file_type_from_key(&key);
        let data_len = data.len();
        let data_mb = data_len as u32 / (1024 * 1024);
        let timeout = self.opts.dispatch_timeout.0 + self.opts.write_timeout_per_mb.0 * data_mb;
        let blob_client = self.blob_client(&key);
        let mut builder = blob_client.put_block_blob(data);
        if let Some(access_tier) = access_tier {
            builder = builder.access_tier(access_tier);
        }
        let result = self
            .timeout_or_io_err(timeout, async { builder.await })
            .await;
        match result {
            Ok(_) => {
                KVENGINE_DFS_THROUGHPUT
                    .azure
                    .put
                    .get(MetricsFileType::from(file_type))
                    .inc_by(data_len as u64);
                KVENGINE_DFS_LATENCY
                    .put
                    .get(MetricsFileType::from(file_type))
                    .observe(start_time.saturating_elapsed().as_secs_f64());
                Ok(())
            }
            Err(err) => Err(Error::Other(format!("azure put error: {}", err))),
        }
    }

    async fn set_deleted_tag(&self, key: &str, deleted: bool) -> crate::dfs::Result<()> {
        if self.read_only {
            return Err(Error::ReadOnly);
        }
        let blob_client = self.blob_client(key);
        let mut tags = match self
            .timeout_or_io_err(self.opts.dispatch_timeout.0, async {
                blob_client.get_tags().await
            })
            .await
        {
            Ok(resp) => resp.tags,
            Err(err) if is_status_code(&err, StatusCode::NOT_FOUND) => Tags::new(),
            Err(err) => return Err(Error::Other(format!("azure get tags error: {}", err))),
        };
        if deleted {
            let mut updated = false;
            for tag in tags.tag_set.tags.iter_mut() {
                if tag.key == TAG_DELETED_KEY {
                    tag.value = "true".to_string();
                    updated = true;
                    break;
                }
            }
            if !updated {
                tags.insert(TAG_DELETED_KEY, "true");
            }
        } else {
            tags.tag_set.tags.retain(|tag| tag.key != TAG_DELETED_KEY);
        }
        let result = self
            .timeout_or_io_err(self.opts.dispatch_timeout.0, async {
                blob_client.set_tags(tags).await
            })
            .await;
        match result {
            Ok(_) => Ok(()),
            Err(err) => Err(Error::Other(format!("azure set tags error: {}", err))),
        }
    }

    async fn get_tag_deleted(&self, key: &str) -> crate::dfs::Result<bool> {
        let blob_client = self.blob_client(key);
        let result = self
            .timeout_or_io_err(self.opts.dispatch_timeout.0, async {
                blob_client.get_tags().await
            })
            .await;
        match result {
            Ok(resp) => Ok(resp
                .tags
                .tag_set
                .tags
                .iter()
                .find(|tag| tag.key == TAG_DELETED_KEY)
                .map(|tag| tag.value == "true")
                .unwrap_or(false)),
            Err(err) if is_status_code(&err, StatusCode::NOT_FOUND) => Ok(false),
            Err(err) => Err(Error::Other(format!("azure get tags error: {}", err))),
        }
    }

    fn storage_class_str(storage_class: StorageClass) -> &'static str {
        match storage_class {
            StorageClass::Standard => STORAGE_CLASS_STANDARD,
            StorageClass::InfrequentAccess => STORAGE_CLASS_STANDARD_IA,
            StorageClass::GlacierInstantRetrieval => STORAGE_CLASS_GLACIER_IR,
            StorageClass::IntelligentTiering => STORAGE_CLASS_INTELLIGENT_TIERING,
            StorageClass::Default => STORAGE_CLASS_DEFAULT,
        }
    }

    fn storage_class_from_tier(tier: Option<AccessTier>) -> StorageClass {
        match tier {
            Some(AccessTier::Hot) => StorageClass::Standard,
            Some(AccessTier::Cool) => StorageClass::InfrequentAccess,
            Some(AccessTier::Archive) => StorageClass::GlacierInstantRetrieval,
            None => StorageClass::Default,
        }
    }

    fn access_tier_from_storage_class(storage_class: StorageClass) -> Option<AccessTier> {
        match storage_class {
            StorageClass::InfrequentAccess => Some(AccessTier::Cool),
            StorageClass::GlacierInstantRetrieval => Some(AccessTier::Archive),
            StorageClass::Standard | StorageClass::IntelligentTiering | StorageClass::Default => {
                Some(AccessTier::Hot)
            }
        }
    }
}

impl Drop for AzureFsCore {
    fn drop(&mut self) {
        if let Some(runtime) = self.runtime.take() {
            std::thread::spawn(|| drop(runtime));
        }
    }
}

impl Debug for AzureFs {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("AzureFs")
            .field("prefix", &self.prefix)
            .finish()
    }
}

impl ObjectStorage for AzureFs {
    fn put_objects(&self, objects: Vec<(String, Bytes)>) -> Result<(), String> {
        let put_object = |key: String, data: Bytes| {
            let full_key = format!("{}/{}", self.prefix, key);
            let fs = self.clone();
            async move {
                fs.put_object(full_key, data, key.clone())
                    .await
                    .map_err(|err| format!("put {} failed {:?}", &key, err))
            }
        };

        let runtime = self.get_runtime();
        let len = objects.len();

        if len == 1 {
            let (key, data) = objects.into_iter().next().unwrap();
            return runtime.block_on(put_object(key, data));
        }

        let mut handles = Vec::with_capacity(len);
        for (key, data) in objects {
            handles.push(runtime.spawn(put_object(key, data)));
        }

        let errs: Vec<_> = runtime
            .block_on(futures::future::join_all(handles))
            .into_iter()
            .filter_map(|r| r.unwrap().err())
            .collect();
        if !errs.is_empty() {
            return Err(format!("{:?}", errs));
        }
        Ok(())
    }

    fn get_objects(
        &self,
        keys: Vec<(String, GetObjectOptions)>,
        cache: Option<&ObjectCacheWithHook>,
    ) -> Result<Vec<(String, Bytes)>, String> {
        let get_object =
            |key: String, opts: GetObjectOptions, cache: Option<&ObjectCacheWithHook>| {
                let full_key = format!("{}/{}", self.prefix, key);
                let fs = self.clone();
                let cache = cache.cloned();
                async move {
                    match fs
                        .get_object_with_cache(full_key, key.clone(), opts, cache.as_ref())
                        .await
                    {
                        Ok(data) => Ok((key, data)),
                        Err(err) => Err(format!("get {} failed {:?}", &key, err)),
                    }
                }
            };

        let runtime = self.get_runtime();
        if keys.len() == 1 {
            let mut keys = keys;
            let (key, opts) = keys.pop().unwrap();
            let res = runtime.block_on(get_object(key, opts, cache))?;
            return Ok(vec![res]);
        }

        let mut join_set = tokio::task::JoinSet::new();
        for (key, opts) in keys {
            join_set.spawn_on(get_object(key, opts, cache), runtime.handle());
        }

        runtime.block_on(async move {
            let mut objects: Vec<(String, Bytes)> = vec![];
            while let Some(res) = join_set.join_next().await {
                let res = res.map_err(|e| e.to_string())??;
                objects.push(res);
            }
            Ok(objects)
        })
    }

    fn list_objects(
        &self,
        start_after: &str,
        prefix: Option<&str>,
        max_keys: Option<u32>,
    ) -> Result<(Vec<ListObjectContent>, Option<String>), String> {
        let runtime = self.get_runtime();
        runtime
            .block_on(self.list(start_after, prefix, max_keys))
            .map(|v| (v.0, v.2))
            .map_err(|err| format!("list failed {:?}", err))
    }
}

#[async_trait]
impl Dfs for AzureFs {
    async fn exists(&self, file_id: u64, opts: Options) -> crate::dfs::Result<bool> {
        let file_key = self.file_key(file_id, opts.file_type);
        self.exist(file_key, format!("{}.{}", file_id, opts.file_type.suffix()))
            .await
    }

    async fn size(&self, file_id: u64, opts: Options) -> crate::dfs::Result<u64> {
        let file_key = self.file_key(file_id, opts.file_type);
        self.object_size(file_key, format!("{}.{}", file_id, opts.file_type.suffix()))
            .await
    }

    async fn read_file(&self, file_id: u64, opts: Options) -> crate::dfs::Result<Bytes> {
        let filename = if let Some(end_off) = opts.end_off {
            format!("{}-{}-{}.seg", file_id, opts.start_off, end_off)
        } else {
            format!("{}.{}", file_id, opts.file_type.suffix())
        };
        self.get_object(
            self.file_key(file_id, opts.file_type),
            filename,
            GetObjectOptions {
                start_off: Some(opts.start_off),
                end_off: opts.end_off,
            },
        )
        .await
    }

    async fn create(&self, file_id: u64, data: Bytes, opts: Options) -> crate::dfs::Result<()> {
        self.put_object_with_storage_class(
            self.file_key(file_id, opts.file_type),
            data,
            format!("{}.{}", file_id, opts.file_type.suffix()),
            StorageClass::Default,
        )
        .await
    }

    async fn remove(&self, file_id: u64, _file_len: Option<u64>, opts: Options) {
        let file_key = self.file_key(file_id, opts.file_type);
        if let Err(err) = self.set_deleted_tag(&file_key, true).await {
            error!("failed to mark azure object removed: {:?}", err);
        }
    }

    async fn permanently_remove(&self, file_id: u64, opts: Options) -> crate::dfs::Result<()> {
        self.delete_object(self.file_key(file_id, opts.file_type), file_id.to_string())
            .await
    }

    async fn list(
        &self,
        start_after: &str,
        prefix: Option<&str>,
        max_keys: Option<u32>,
    ) -> crate::dfs::Result<(Vec<ListObjectContent>, bool, Option<String>)> {
        // The caller is responsible for routing to the correct account via
        // apply_as_local_override before building this DFS instance.
        let (client, base_prefix) = (&self.container_client, self.prefix.as_str());
        let prefix = format!("{}/{}", base_prefix, prefix.unwrap_or_default());
        let start_from = build_list_start_from(&prefix, start_after);
        let mut builder = client.list_blobs().prefix(Prefix::new(prefix.clone()));
        if let Some(request_max_keys) = build_list_request_max_keys(max_keys, start_from.is_some())
        {
            if let Some(max_results) = NonZeroU32::new(request_max_keys) {
                builder = builder.max_results(MaxResults::new(max_results));
            }
        }
        let mut context = Context::new();
        if let Some(start_from) = &start_from {
            context.insert(AzureListStartFrom(start_from.clone()));
        }
        let mut stream = builder.context(context).into_stream();
        let response = match self
            .timeout_or_io_err(self.opts.dispatch_timeout.0, async {
                stream.next().await.transpose()
            })
            .await
        {
            Ok(Some(resp)) => resp,
            Ok(None) => return Ok((vec![], false, None)),
            Err(err) if is_status_code(&err, StatusCode::NOT_FOUND) => {
                return Ok((vec![], false, None));
            }
            Err(err) => return Err(Error::Other(format!("azure list error: {}", err))),
        };

        let contents = response
            .blobs
            .blobs()
            .map(|blob| ListObjectContent {
                key: blob.name.clone(),
                last_modified: azure_core::date::to_rfc3339(&blob.properties.last_modified),
                storage_class: AzureFsCore::storage_class_str(
                    AzureFsCore::storage_class_from_tier(blob.properties.access_tier),
                )
                .to_string(),
                size: blob.properties.content_length,
            })
            .collect::<Vec<_>>();

        let (contents, has_more, next_start_after) = finalize_list_result(
            &prefix,
            start_from.as_deref(),
            max_keys,
            response.next_marker.is_some(),
            contents,
        );
        Ok((contents, has_more, next_start_after))
    }

    async fn get_object(
        &self,
        key: String,
        file_name: String,
        opts: GetObjectOptions,
    ) -> crate::dfs::Result<Bytes> {
        let (writer, _) = self
            .get_object_to_writer(key, file_name, opts, &|| Ok(BytesMut::new().writer()))
            .await?;
        Ok(writer.into_inner().freeze())
    }

    async fn get_object_to_path(
        &self,
        key: String,
        file_name: String,
        opts: GetObjectOptions,
        path: &Path,
    ) -> crate::dfs::Result<u64> {
        let path = path.to_path_buf();
        let build_writer = move || -> std::io::Result<_> {
            let file = std::fs::OpenOptions::new()
                .create(true)
                .write(true)
                .read(true)
                .truncate(true)
                .open(&path)?;
            Ok(std::io::BufWriter::new(file))
        };
        let (mut writer, len) = self
            .get_object_to_writer(key, file_name, opts, build_writer)
            .await?;
        writer.flush()?;
        Ok(len)
    }

    async fn put_object(
        &self,
        key: String,
        data: Bytes,
        _file_name: String,
    ) -> crate::dfs::Result<()> {
        self.put_object_with_tier(key, data, None).await
    }

    async fn put_object_with_storage_class(
        &self,
        key: String,
        data: Bytes,
        _file_name: String,
        storage_class: StorageClass,
    ) -> crate::dfs::Result<()> {
        let access_tier = AzureFsCore::access_tier_from_storage_class(storage_class);
        self.put_object_with_tier(key, data, access_tier).await
    }

    async fn exist(&self, key: String, _file_name: String) -> crate::dfs::Result<bool> {
        match self.object_size_inner(&key).await {
            Ok(_) => Ok(true),
            Err(Error::NoSuchKey(_)) => Ok(false),
            Err(err) => Err(err),
        }
    }

    async fn retain_file(&self, file_key: &str) -> crate::dfs::Result<()> {
        self.set_deleted_tag(file_key, false).await
    }

    async fn delete_object(&self, key: String, _file_name: String) -> crate::dfs::Result<()> {
        if self.read_only {
            return Err(Error::ReadOnly);
        }
        let result = self
            .timeout_or_io_err(self.opts.dispatch_timeout.0, async {
                self.blob_client(&key).delete().await
            })
            .await;
        match result {
            Ok(_) => Ok(()),
            Err(err) if is_status_code(&err, StatusCode::NOT_FOUND) => Ok(()),
            Err(err) => Err(Error::Other(format!("azure delete error: {}", err))),
        }
    }

    async fn object_size(&self, key: String, _file_name: String) -> crate::dfs::Result<u64> {
        self.object_size_inner(&key).await
    }

    async fn list_folders(
        &self,
        prefix: &str,
        delimiter: Option<&str>,
    ) -> crate::dfs::Result<Vec<String>> {
        let delimiter = delimiter.unwrap_or("/");
        let prefix = format!("{}/{}", self.prefix, prefix);
        let mut folders = Vec::new();
        let builder = self
            .container_client
            .list_blobs()
            .prefix(Prefix::new(prefix.clone()))
            .delimiter(Delimiter::new(delimiter.to_string()));
        let mut stream = builder.into_stream();
        let mut err = None;
        loop {
            let response = match self
                .timeout_or_io_err(self.opts.dispatch_timeout.0, async {
                    stream.next().await.transpose()
                })
                .await
            {
                Ok(Some(resp)) => resp,
                Ok(None) => break,
                Err(e) => {
                    err = Some(e);
                    break;
                }
            };
            folders.extend(response.blobs.prefixes().map(|p| p.name.clone()));
        }
        if let Some(err) = err {
            if is_status_code(&err, StatusCode::NOT_FOUND) {
                return Ok(vec![]);
            }
            return Err(Error::Other(format!("azure list error: {}", err)));
        }
        Ok(folders)
    }

    async fn is_removed(&self, file_key: &str) -> crate::dfs::Result<bool> {
        self.get_tag_deleted(file_key).await
    }

    fn list_objects(
        &self,
        start_after: &str,
        prefix: Option<&str>,
        max_keys: Option<u32>,
    ) -> std::result::Result<(Vec<ListObjectContent>, Option<String>), String> {
        ObjectStorage::list_objects(self, start_after, prefix, max_keys)
    }

    fn put_objects(&self, objects: Vec<(String, Bytes)>) -> std::result::Result<(), String> {
        ObjectStorage::put_objects(self, objects)
    }

    fn get_runtime(&self) -> &Runtime {
        self.runtime.as_ref().unwrap()
    }

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

    fn get_storage_name(&self) -> String {
        self.container_client.container_name().to_string()
    }

    async fn raw_copy_object(
        &self,
        source_key: &str,
        target_key: &str,
        _target_tagging: Option<&crate::dfs::Tagging>,
        _target_storage_class: Option<&str>,
    ) -> crate::dfs::Result<()> {
        self.raw_copy_object_inner(source_key, target_key).await
    }
}

fn is_status_code(err: &(dyn std::error::Error + 'static), status: StatusCode) -> bool {
    let Some(azure_err) = err.downcast_ref::<AzureError>() else {
        return false;
    };
    if let Some(http_err) = azure_err.as_http_error() {
        return http_err.status() == status.as_u16();
    }
    match azure_err.kind() {
        ErrorKind::HttpResponse {
            status: err_status, ..
        } => *err_status == status.as_u16(),
        _ => false,
    }
}

fn build_source_url(
    endpoint: &azure_core::Url,
    source_key: &str,
) -> crate::dfs::Result<azure_core::Url> {
    let mut base = endpoint.clone();
    if !base.path().ends_with('/') {
        base.set_path(&format!("{}/", base.path()));
    }
    base.join(source_key)
        .map_err(|e| Error::Other(format!("raw_copy_object source url: {}", e)))
}

fn build_list_start_from(prefix: &str, start_after: &str) -> Option<String> {
    (!start_after.is_empty()).then(|| format!("{}{}", prefix, start_after))
}

fn build_list_request_max_keys(max_keys: Option<u32>, has_start_after: bool) -> Option<u32> {
    if !has_start_after {
        return max_keys;
    }
    max_keys.map(|max_keys| {
        if max_keys == 0 {
            0
        } else {
            max_keys.saturating_add(1)
        }
    })
}

fn build_next_start_after(
    prefix: &str,
    has_more: bool,
    contents: &[ListObjectContent],
) -> Option<String> {
    if !has_more {
        return None;
    }
    contents.last().map(|content| {
        content
            .key
            .strip_prefix(prefix)
            .unwrap_or(&content.key)
            .to_string()
    })
}

fn finalize_list_result(
    prefix: &str,
    start_from: Option<&str>,
    max_keys: Option<u32>,
    response_has_more: bool,
    mut contents: Vec<ListObjectContent>,
) -> (Vec<ListObjectContent>, bool, Option<String>) {
    if let Some(start_from) = start_from {
        if contents
            .first()
            .is_some_and(|content| content.key == start_from)
        {
            contents.remove(0);
        }
    }

    let mut has_more = response_has_more;
    if let Some(max_keys) = max_keys.map(|max_keys| max_keys as usize) {
        if max_keys > 0 && contents.len() > max_keys {
            contents.truncate(max_keys);
            has_more = true;
        }
    }

    let next_start_after = build_next_start_after(prefix, has_more, &contents);
    (contents, has_more, next_start_after)
}

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

    use bytes::Bytes;
    use engine_traits::{GetObjectOptions, ListObjectContent, ObjectCache, ObjectCacheWithHook};
    use rand::Rng;
    use tokio::runtime::Runtime;

    use super::{
        DEVELOPMENT_STORAGE_CONNECTION_STRING, build_list_request_max_keys, build_list_start_from,
        build_next_start_after, build_source_url, finalize_list_result,
    };
    use crate::dfs::{
        AzureFsCore, FileType, Options, StorageClass,
        config::{AzureConfig, Config},
        new_dfs_from_config,
    };

    const REAL_AZURE_ENV: &str = "TEST_REAL_AZURE";

    fn azurite_config() -> Config {
        let mut rng = rand::thread_rng();
        let suffix = rng.r#gen::<u64>();
        Config {
            backend: "azure".to_string(),
            prefix: format!("azurite-{}", suffix),
            azure: AzureConfig {
                connection_string: DEVELOPMENT_STORAGE_CONNECTION_STRING.to_string(),
                container: format!("azurite-{}", suffix),
                ..Default::default()
            },
            ..Default::default()
        }
    }

    fn create_azurite_container(cfg: &Config) {
        let retry_options = AzureFsCore::retry_options(&cfg.conn_options);
        let container_client = AzureFsCore::build_container_client(&cfg.azure, &retry_options)
            .expect("build azurite container client");
        let runtime = Runtime::new().unwrap();
        runtime.block_on(async {
            container_client.create().await.unwrap();
        });
    }

    #[test]
    fn test_build_list_start_from() {
        assert_eq!(build_list_start_from("root/prefix/", ""), None);
        assert_eq!(
            build_list_start_from("root/prefix/", "obj-001"),
            Some("root/prefix/obj-001".to_string())
        );
    }

    #[test]
    fn test_build_next_start_after() {
        let contents = vec![
            ListObjectContent {
                key: "root/prefix/obj-001".to_string(),
                ..Default::default()
            },
            ListObjectContent {
                key: "root/prefix/obj-002".to_string(),
                ..Default::default()
            },
        ];
        assert_eq!(
            build_next_start_after("root/prefix/", true, &contents),
            Some("obj-002".to_string())
        );
        assert_eq!(
            build_next_start_after("root/prefix/", false, &contents),
            None
        );
        assert_eq!(build_next_start_after("root/prefix/", true, &[]), None);
    }

    #[test]
    fn test_build_list_request_max_keys() {
        assert_eq!(build_list_request_max_keys(None, false), None);
        assert_eq!(build_list_request_max_keys(Some(3), false), Some(3));
        assert_eq!(build_list_request_max_keys(Some(3), true), Some(4));
        assert_eq!(build_list_request_max_keys(Some(0), true), Some(0));
    }

    #[test]
    fn test_build_source_url() {
        let url = |s| azure_core::Url::parse(s).unwrap();
        assert_eq!(
            build_source_url(
                &url("https://acct.blob.core.windows.net"),
                "src/path/to/blob"
            )
            .unwrap()
            .as_str(),
            "https://acct.blob.core.windows.net/src/path/to/blob"
        );
        assert_eq!(
            build_source_url(
                &url("http://127.0.0.1:10000/devstoreaccount1"),
                "src/path/to/blob"
            )
            .unwrap()
            .as_str(),
            "http://127.0.0.1:10000/devstoreaccount1/src/path/to/blob"
        );
    }

    #[test]
    fn test_finalize_list_result_skip_start_after() {
        let prefix = "root/prefix/";
        let contents = vec![
            ListObjectContent {
                key: "root/prefix/obj-001".to_string(),
                ..Default::default()
            },
            ListObjectContent {
                key: "root/prefix/obj-002".to_string(),
                ..Default::default()
            },
        ];
        let (contents, has_more, next_start_after) =
            finalize_list_result(prefix, Some("root/prefix/obj-001"), Some(1), true, contents);
        assert_eq!(contents.len(), 1);
        assert_eq!(contents[0].key, "root/prefix/obj-002");
        assert!(has_more);
        assert_eq!(next_start_after, Some("obj-002".to_string()));
    }

    #[test]
    fn test_finalize_list_result_trim_to_max_keys() {
        let prefix = "root/prefix/";
        let contents = vec![
            ListObjectContent {
                key: "root/prefix/obj-001".to_string(),
                ..Default::default()
            },
            ListObjectContent {
                key: "root/prefix/obj-002".to_string(),
                ..Default::default()
            },
            ListObjectContent {
                key: "root/prefix/obj-003".to_string(),
                ..Default::default()
            },
        ];
        let (contents, has_more, next_start_after) =
            finalize_list_result(prefix, None, Some(2), false, contents);
        assert_eq!(contents.len(), 2);
        assert_eq!(contents[1].key, "root/prefix/obj-002");
        assert!(has_more);
        assert_eq!(next_start_after, Some("obj-002".to_string()));
    }

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

        if std::env::var(REAL_AZURE_ENV).ok().as_deref() != Some("1") {
            eprintln!(
                "skip real azure list test; set {}=1 to enable",
                REAL_AZURE_ENV
            );
            return;
        }
        std::env::set_var("TIKV_DISABLE_SYSTEM_PROXY", "1");
        let mut dfs_config = Config::default();
        dfs_config.override_from_env();

        let dfs = new_dfs_from_config(dfs_config);
        let runtime = dfs.get_runtime();

        let mut rng = rand::thread_rng();
        let list_prefix = format!("list-api-{}/", rng.gen::<u64>());
        let object_prefix = format!("{}/{}", dfs.get_prefix(), list_prefix);
        let object_names = ["obj-001", "obj-002", "obj-003"];

        for name in object_names {
            let relative_key = format!("{}{}", list_prefix, name);
            let full_key = format!("{}/{}", dfs.get_prefix(), relative_key);
            runtime
                .block_on(dfs.put_object(
                    full_key,
                    Bytes::from(name.as_bytes().to_vec()),
                    name.to_string(),
                ))
                .unwrap();
        }

        let (all_contents, has_more, next_start_after) = runtime
            .block_on(dfs.list("", Some(&list_prefix), None))
            .unwrap();
        let mut listed_keys = all_contents
            .iter()
            .map(|content| content.key.clone())
            .collect::<Vec<_>>();
        listed_keys.sort();
        assert_eq!(
            listed_keys,
            vec![
                format!("{}obj-001", object_prefix),
                format!("{}obj-002", object_prefix),
                format!("{}obj-003", object_prefix),
            ]
        );
        assert!(!has_more);
        assert_eq!(next_start_after, None);

        // Existing start_after: should exclude the start key and report the next key.
        let (from_existing, has_more, next_start_after) = runtime
            .block_on(dfs.list("obj-001", Some(&list_prefix), Some(1)))
            .unwrap();
        assert_eq!(from_existing.len(), 1);
        assert!(from_existing[0].key.ends_with("obj-002"));
        assert!(has_more);
        assert_eq!(next_start_after, Some("obj-002".to_string()));

        // Missing start_after: validates the branch where no first-element filtering is
        // needed.
        let (from_missing, has_more, next_start_after) = runtime
            .block_on(dfs.list("obj-001-zzz", Some(&list_prefix), Some(1)))
            .unwrap();
        assert_eq!(from_missing.len(), 1);
        assert!(from_missing[0].key.ends_with("obj-002"));
        assert!(has_more);
        assert_eq!(next_start_after, Some("obj-002".to_string()));

        // max_keys=0: list does not truncate after start_after filtering.
        let (max_zero, has_more, next_start_after) = runtime
            .block_on(dfs.list("obj-001", Some(&list_prefix), Some(0)))
            .unwrap();
        assert_eq!(max_zero.len(), 2);
        assert!(
            max_zero
                .iter()
                .all(|content| !content.key.ends_with("obj-001"))
        );
        assert!(!has_more);
        assert_eq!(next_start_after, None);

        let (after_last, has_more, next_start_after) = runtime
            .block_on(dfs.list("obj-003", Some(&list_prefix), None))
            .unwrap();
        assert!(after_last.is_empty());
        assert!(!has_more);
        assert_eq!(next_start_after, None);

        for name in object_names {
            let full_key = format!("{}/{list_prefix}{name}", dfs.get_prefix());
            runtime
                .block_on(dfs.delete_object(full_key, name.to_string()))
                .unwrap();
        }
    }

    #[test]
    fn test_azurite() {
        test_util::init_log_for_test();
        if std::env::var("TIKV_AZURITE_TEST").ok().as_deref() != Some("1") {
            eprintln!("skip azurite test; set TIKV_AZURITE_TEST=1 to enable");
            return;
        }
        std::env::set_var("TIKV_DISABLE_SYSTEM_PROXY", "1");

        let cfg = azurite_config();
        create_azurite_container(&cfg);
        let dfs = new_dfs_from_config(cfg);
        let runtime = dfs.get_runtime();

        let file_data = Bytes::from_static(b"azurite-file");
        runtime
            .block_on(dfs.create(1, file_data.clone(), Options::default()))
            .unwrap();
        let read_data = runtime
            .block_on(dfs.read_file(1, Options::default()))
            .unwrap();
        assert_eq!(read_data, file_data);

        let delete_data = Bytes::from_static(b"azurite-delete");
        runtime
            .block_on(dfs.create(2, delete_data, Options::default()))
            .unwrap();
        runtime
            .block_on(dfs.permanently_remove(2, Options::default()))
            .unwrap();

        let object_key = format!("{}/objects/test", dfs.get_prefix());
        let object_data = Bytes::from_static(b"azurite-object");
        runtime
            .block_on(dfs.put_object_with_storage_class(
                object_key.clone(),
                object_data.clone(),
                "test".to_string(),
                StorageClass::Standard,
            ))
            .unwrap();

        let plain_key = format!("{}/objects/plain", dfs.get_prefix());
        let plain_data = Bytes::from_static(b"azurite-plain");
        runtime
            .block_on(dfs.put_object(plain_key.clone(), plain_data.clone(), "plain".to_string()))
            .unwrap();
        let plain_fetched = runtime
            .block_on(dfs.get_object(
                plain_key.clone(),
                "plain".to_string(),
                GetObjectOptions::default(),
            ))
            .unwrap();
        assert_eq!(plain_fetched, plain_data);

        let exists = runtime
            .block_on(dfs.exist(object_key.clone(), "test".to_string()))
            .unwrap();
        assert!(exists);
        let size = runtime
            .block_on(dfs.object_size(object_key.clone(), "test".to_string()))
            .unwrap();
        assert_eq!(size as usize, object_data.len());

        let fetched = runtime
            .block_on(dfs.get_object(
                object_key.clone(),
                "test".to_string(),
                GetObjectOptions::default(),
            ))
            .unwrap();
        assert_eq!(fetched, object_data);

        let range_key = format!("{}/objects/range", dfs.get_prefix());
        let range_data = Bytes::from_static(b"0123456789abcdef");
        runtime
            .block_on(dfs.put_object(range_key.clone(), range_data.clone(), "range".to_string()))
            .unwrap();
        let range_fetched = runtime
            .block_on(dfs.get_object(
                range_key,
                "range".to_string(),
                GetObjectOptions {
                    start_off: Some(2),
                    end_off: Some(6),
                },
            ))
            .unwrap();
        assert_eq!(range_fetched, Bytes::from_static(b"2345"));

        let hook_calls = Arc::new(AtomicUsize::new(0));
        let hook_calls_clone = Arc::clone(&hook_calls);
        let cache = ObjectCache::new(1024 * 1024, 256);
        let cache = ObjectCacheWithHook::from(cache).with_hook(Arc::new(move |data| {
            hook_calls_clone.fetch_add(1, Ordering::SeqCst);
            Ok(data)
        }));
        let cached = runtime
            .block_on(dfs.get_object_with_cache(
                object_key.clone(),
                "test".to_string(),
                GetObjectOptions::default(),
                Some(&cache),
            ))
            .unwrap();
        assert_eq!(cached, object_data);
        let cached_again = runtime
            .block_on(dfs.get_object_with_cache(
                object_key.clone(),
                "test".to_string(),
                GetObjectOptions::default(),
                Some(&cache),
            ))
            .unwrap();
        assert_eq!(cached_again, object_data);
        assert_eq!(hook_calls.load(Ordering::SeqCst), 1);

        let temp_dir = tempfile::tempdir().unwrap();
        let path = temp_dir.path().join("downloaded");
        let saved = runtime
            .block_on(dfs.get_object_to_path(
                object_key.clone(),
                "test".to_string(),
                GetObjectOptions::default(),
                &path,
            ))
            .unwrap();
        assert_eq!(saved as usize, object_data.len());
        assert_eq!(std::fs::read(&path).unwrap(), object_data);

        dfs.put_objects(vec![
            (
                "folders/alpha/item".to_string(),
                Bytes::from_static(b"alpha"),
            ),
            ("folders/beta/item".to_string(), Bytes::from_static(b"beta")),
        ])
        .unwrap();

        let folders = runtime
            .block_on(dfs.list_folders("folders/", None))
            .unwrap();
        assert!(
            folders
                .iter()
                .any(|folder| folder.contains("folders/alpha"))
        );
        assert!(folders.iter().any(|folder| folder.contains("folders/beta")));

        let (objects, ..) = runtime
            .block_on(dfs.list("", Some("objects/"), None))
            .unwrap();
        assert!(objects.iter().any(|o| o.key.ends_with("objects/test")));

        let (listed, _) = dfs.list_objects("", Some("objects/"), None).unwrap();
        assert!(listed.iter().any(|o| o.key.ends_with("objects/test")));

        let file_key = dfs.file_key(1, FileType::Sst);
        runtime.block_on(async {
            dfs.remove(1, None, Options::default()).await;
        });
        let removed = runtime.block_on(dfs.is_removed(&file_key)).unwrap();
        assert!(removed);
        runtime.block_on(dfs.retain_file(&file_key)).unwrap();
        let removed = runtime.block_on(dfs.is_removed(&file_key)).unwrap();
        assert!(!removed);

        runtime
            .block_on(dfs.delete_object(object_key, "test".to_string()))
            .unwrap();
    }

    #[test]
    fn test_azurite_init_inside_runtime() {
        test_util::init_log_for_test();
        if std::env::var("TIKV_AZURITE_TEST").ok().as_deref() != Some("1") {
            eprintln!("skip azurite runtime-init test; set TIKV_AZURITE_TEST=1 to enable");
            return;
        }
        std::env::set_var("TIKV_DISABLE_SYSTEM_PROXY", "1");

        let cfg = azurite_config();
        create_azurite_container(&cfg);

        let outer_runtime = Runtime::new().unwrap();
        let _guard = outer_runtime.enter();
        let dfs = new_dfs_from_config(cfg);
        let runtime = dfs.get_runtime();
        let (objects, has_more, next_start_after) = runtime
            .block_on(dfs.list("", Some("objects/"), Some(1)))
            .unwrap();
        assert!(objects.is_empty());
        assert!(!has_more);
        assert_eq!(next_start_after, None);
    }
}