// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors

use std::collections::{HashMap, HashSet};
use std::sync::Arc;

use async_trait::async_trait;
use http::StatusCode;
use lance_io::object_store::StorageOptions;
use lance_namespace_impls::{DynamicContextProvider, OperationInfo};
use moka::future::Cache;
use reqwest::Response;
use reqwest::header::CONTENT_TYPE;

use lance_namespace::models::{
    CreateNamespaceRequest, CreateNamespaceResponse, DescribeNamespaceRequest,
    DescribeNamespaceResponse, DropNamespaceRequest, DropNamespaceResponse, ListNamespacesRequest,
    ListNamespacesResponse, ListTablesRequest, ListTablesResponse,
};

use crate::Error;
use crate::database::{
    CloneTableRequest, CreateTableMode, CreateTableRequest, Database, DatabaseOptions, JobInfo,
    OpenTableRequest, ReadConsistency, TableNamesRequest,
};
use crate::error::Result;
use crate::function::{FunctionRegistrationRequest, FunctionVersion};
use crate::job::Job;
use crate::remote::job::{RemoteJob, job_state_to_client};
use crate::remote::util::stream_as_body;
use crate::table::BaseTable;

use super::client::{
    ClientConfig, HeaderProvider, HttpSend, RequestResultExt, RestfulLanceDbClient, Sender,
};
use super::sql::SqlClient;
use super::table::RemoteTable;
use super::util::parse_server_version;
use super::{ARROW_STREAM_CONTENT_TYPE, extract_job_id};

// Request structure for the remote clone table API
#[derive(serde::Serialize)]
struct RemoteCloneTableRequest {
    source_location: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    source_version: Option<u64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    source_tag: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    is_shallow: Option<bool>,
}

// the versions of the server that we support
// for any new feature that we need to change the SDK behavior, we should bump the server version,
// and add a feature flag as method of `ServerVersion` here.
pub const DEFAULT_SERVER_VERSION: semver::Version = semver::Version::new(0, 1, 0);
#[derive(Debug, Clone)]
pub struct ServerVersion(pub semver::Version);

impl Default for ServerVersion {
    fn default() -> Self {
        Self(DEFAULT_SERVER_VERSION.clone())
    }
}

impl ServerVersion {
    pub fn parse(version: &str) -> Result<Self> {
        let version = Self(
            semver::Version::parse(version).map_err(|e| Error::InvalidInput {
                message: e.to_string(),
            })?,
        );
        Ok(version)
    }

    pub fn support_multivector(&self) -> bool {
        self.0 >= semver::Version::new(0, 2, 0)
    }

    pub fn support_structural_fts(&self) -> bool {
        self.0 >= semver::Version::new(0, 3, 0)
    }

    pub fn support_multipart_write(&self) -> bool {
        self.0 >= semver::Version::new(0, 4, 0)
    }

    pub fn support_blobs(&self) -> bool {
        self.0 >= semver::Version::new(0, 5, 0)
    }

    pub fn support_fts_document_granularity(&self) -> bool {
        self.0 >= semver::Version::new(0, 6, 0)
    }
}

pub const OPT_REMOTE_PREFIX: &str = "remote_database_";
pub const OPT_REMOTE_API_KEY: &str = "remote_database_api_key";
pub const OPT_REMOTE_REGION: &str = "remote_database_region";
pub const OPT_REMOTE_HOST_OVERRIDE: &str = "remote_database_host_override";
pub const OPT_REMOTE_SQL_HOST_OVERRIDE: &str = "remote_database_sql_host_override";
// TODO: add support for configuring client config via key/value options

#[derive(Clone, Debug, Default)]
pub struct RemoteDatabaseOptions {
    /// The LanceDB Cloud API key
    pub api_key: Option<String>,
    /// The LanceDB Cloud region
    pub region: Option<String>,
    /// The LanceDB Enterprise host override
    ///
    /// This is required when connecting to LanceDB Enterprise and should be
    /// provided if using an on-premises LanceDB Enterprise instance.
    pub host_override: Option<String>,
    /// Storage options configure the storage layer (e.g. S3, GCS, Azure, etc.)
    ///
    /// See available options at <https://docs.lancedb.com/storage/>
    ///
    /// These options are only used for LanceDB Enterprise and only a subset of options
    /// are supported.
    pub storage_options: HashMap<String, String>,
}

impl RemoteDatabaseOptions {
    pub fn builder() -> RemoteDatabaseOptionsBuilder {
        RemoteDatabaseOptionsBuilder::new()
    }

    pub(crate) fn parse_from_map(map: &HashMap<String, String>) -> Result<Self> {
        let api_key = map.get(OPT_REMOTE_API_KEY).cloned();
        let region = map.get(OPT_REMOTE_REGION).cloned();
        let host_override = map.get(OPT_REMOTE_HOST_OVERRIDE).cloned();
        let storage_options = map
            .iter()
            .filter(|(key, _)| !key.starts_with(OPT_REMOTE_PREFIX))
            .map(|(key, value)| (key.clone(), value.clone()))
            .collect();
        Ok(Self {
            api_key,
            region,
            host_override,
            storage_options,
        })
    }
}

impl DatabaseOptions for RemoteDatabaseOptions {
    fn serialize_into_map(&self, map: &mut HashMap<String, String>) {
        for (key, value) in &self.storage_options {
            map.insert(key.clone(), value.clone());
        }
        if let Some(api_key) = &self.api_key {
            map.insert(OPT_REMOTE_API_KEY.to_string(), api_key.clone());
        }
        if let Some(region) = &self.region {
            map.insert(OPT_REMOTE_REGION.to_string(), region.clone());
        }
        if let Some(host_override) = &self.host_override {
            map.insert(OPT_REMOTE_HOST_OVERRIDE.to_string(), host_override.clone());
        }
    }
}

#[derive(Clone, Debug, Default)]
pub struct RemoteDatabaseOptionsBuilder {
    options: RemoteDatabaseOptions,
}

impl RemoteDatabaseOptionsBuilder {
    pub fn new() -> Self {
        Self {
            options: RemoteDatabaseOptions::default(),
        }
    }

    /// Set the LanceDB Cloud API key
    ///
    /// # Arguments
    ///
    /// * `api_key` - The LanceDB Cloud API key
    pub fn api_key(mut self, api_key: String) -> Self {
        self.options.api_key = Some(api_key);
        self
    }

    /// Set the LanceDB Cloud region
    ///
    /// # Arguments
    ///
    /// * `region` - The LanceDB Cloud region
    pub fn region(mut self, region: String) -> Self {
        self.options.region = Some(region);
        self
    }

    /// Set the LanceDB Enterprise host override
    ///
    /// # Arguments
    ///
    /// * `host_override` - The LanceDB Enterprise host override
    pub fn host_override(mut self, host_override: String) -> Self {
        self.options.host_override = Some(host_override);
        self
    }
}

#[derive(Debug)]
pub struct RemoteDatabase<S: HttpSend = Sender> {
    client: RestfulLanceDbClient<S>,
    table_cache: Cache<String, Arc<RemoteTable<S>>>,
    uri: String,
    /// Headers to pass to the namespace client for authentication
    namespace_headers: HashMap<String, String>,
    namespace_context_provider: Option<Arc<dyn DynamicContextProvider>>,
    /// TLS configuration for mTLS support
    tls_config: Option<super::client::TlsConfig>,
    sql_client: Option<SqlClient>,
}

#[derive(Clone)]
struct NamespaceHeaderProviderContext {
    header_provider: Arc<dyn HeaderProvider>,
}

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

impl DynamicContextProvider for NamespaceHeaderProviderContext {
    fn provide_context(&self, _info: &OperationInfo) -> HashMap<String, String> {
        let header_provider = Arc::clone(&self.header_provider);
        let handle = match std::thread::Builder::new()
            .name("lancedb-namespace-headers".to_string())
            .spawn(move || {
                tokio::runtime::Builder::new_current_thread()
                    .enable_all()
                    .build()
                    .map_err(|e| Error::Runtime {
                        message: format!(
                            "Failed to create runtime for namespace header provider: {e}"
                        ),
                    })?
                    .block_on(header_provider.get_headers())
            }) {
            Ok(handle) => handle,
            Err(err) => {
                log::warn!("Failed to spawn dynamic namespace header provider thread: {err}");
                return HashMap::new();
            }
        };

        let headers = handle.join();

        match headers {
            Ok(Ok(headers)) => headers
                .into_iter()
                .map(|(key, value)| (format!("headers.{key}"), value))
                .collect(),
            Ok(Err(err)) => {
                log::warn!("Failed to get dynamic namespace headers: {err}");
                HashMap::new()
            }
            Err(_) => {
                log::warn!("Dynamic namespace header provider panicked");
                HashMap::new()
            }
        }
    }
}

pub struct RemoteHostOverrides {
    pub rest: Option<String>,
    pub sql: Option<String>,
}

impl RemoteDatabase {
    pub(crate) fn try_new(
        uri: &str,
        api_key: &str,
        region: &str,
        host_overrides: RemoteHostOverrides,
        client_config: ClientConfig,
        options: RemoteOptions,
        read_consistency_interval: Option<std::time::Duration>,
    ) -> Result<Self> {
        let parsed = super::client::parse_db_url(uri)?;
        let sql_client = SqlClient::new(
            parsed.db_name.clone(),
            parsed.db_prefix.clone(),
            api_key.to_string(),
            host_overrides.rest.clone(),
            host_overrides.sql,
            client_config.clone(),
        );
        let header_map = RestfulLanceDbClient::<Sender>::default_headers(
            api_key,
            region,
            &parsed.db_name,
            host_overrides.rest.is_some(),
            &options,
            parsed.db_prefix.as_deref(),
            &client_config,
        )?;

        let namespace_headers: HashMap<String, String> = header_map
            .iter()
            .filter_map(|(k, v)| {
                v.to_str()
                    .ok()
                    .map(|val| (k.as_str().to_string(), val.to_string()))
            })
            .collect();

        let namespace_context_provider =
            client_config
                .header_provider
                .as_ref()
                .map(|header_provider| {
                    Arc::new(NamespaceHeaderProviderContext {
                        header_provider: Arc::clone(header_provider),
                    }) as Arc<dyn DynamicContextProvider>
                });

        let client = RestfulLanceDbClient::try_new(
            &parsed,
            region,
            host_overrides.rest,
            header_map,
            client_config.clone(),
            read_consistency_interval,
        )?;

        let table_cache = Cache::builder()
            .time_to_live(std::time::Duration::from_secs(300))
            .max_capacity(10_000)
            .build();

        Ok(Self {
            client,
            table_cache,
            uri: uri.to_owned(),
            namespace_headers,
            namespace_context_provider,
            tls_config: client_config.tls_config,
            sql_client: Some(sql_client),
        })
    }
}

impl<S: HttpSend> RemoteDatabase<S> {
    async fn submit_drop_table(
        &self,
        name: &str,
        namespace_path: &[String],
    ) -> Result<(String, Response)> {
        let identifier = build_table_identifier(name, namespace_path, &self.client.id_delimiter);
        let cache_key = build_cache_key(name, namespace_path);
        let req = self.client.post(&format!("/v1/table/{}/drop/", identifier));
        let (request_id, resp) = self.client.send(req).await?;
        let resp = self.client.check_response(&request_id, resp).await?;
        self.table_cache.remove(&cache_key).await;
        Ok((request_id, resp))
    }

    /// Collect the tables of a namespace in name order, for `table_names`.
    ///
    /// `table_names` promises name order and resumes after a table name, but the namespace
    /// route's `page_token` is opaque -- it belongs to the store the listing walks, and a
    /// token this client invented would resume from the wrong place. So the whole namespace is
    /// walked by handing each response's token straight back, and the name semantics are
    /// applied here. Constructing no token is what makes this work against a server on either
    /// side of the change: it only ever repeats what the server said.
    ///
    /// This is the cost `table_names` already paid -- the server used to enumerate and sort the
    /// namespace on every request -- and it is why `list_tables` replaces it.
    async fn table_names_in_namespace(
        &self,
        request: &TableNamesRequest,
    ) -> Result<(Vec<String>, ServerVersion)> {
        let namespace_id =
            build_namespace_identifier(&request.namespace_path, &self.client.id_delimiter);
        let path = format!("/v1/namespace/{}/table/list", namespace_id);

        let mut names = Vec::new();
        // Every page reports the same server, so keep the first page's version.
        let mut version: Option<ServerVersion> = None;
        let mut page_token: Option<String> = None;
        loop {
            let mut req = self.client.get(&path);
            if let Some(ref token) = page_token {
                req = req.query(&[("page_token", token)]);
            }
            let (request_id, rsp) = self.client.send_with_retry(req, None, true).await?;
            let rsp = self.client.check_response(&request_id, rsp).await?;
            if version.is_none() {
                version = Some(parse_server_version(&request_id, &rsp)?);
            }
            let response: ListTablesResponse = rsp.json().await.err_to_http(request_id)?;
            names.extend(response.tables);
            // An empty token is the end of the listing, not a token to send back: a server
            // that reads an empty token as "start from the beginning" would hand back the
            // first page again.
            match response.page_token.filter(|token| !token.is_empty()) {
                // A server that repeated a token would never finish; treat that as the end
                // rather than looping on it.
                Some(token) if Some(&token) != page_token.as_ref() => page_token = Some(token),
                _ => break,
            }
        }

        names.sort();
        if let Some(ref start_after) = request.start_after {
            names.retain(|name| name > start_after);
        }
        if let Some(limit) = request.limit {
            names.truncate(limit as usize);
        }
        Ok((names, version.unwrap_or_default()))
    }
}

#[cfg(all(test, feature = "remote"))]
mod test_utils {
    use super::*;
    use crate::remote::ClientConfig;
    use crate::remote::client::test_utils::MockSender;
    use crate::remote::client::test_utils::{client_with_handler, client_with_handler_and_config};

    impl RemoteDatabase<MockSender> {
        pub fn new_mock<F, T>(handler: F) -> Self
        where
            F: Fn(reqwest::Request) -> http::Response<T> + Send + Sync + 'static,
            T: Into<reqwest::Body>,
        {
            let client = client_with_handler(handler);
            Self {
                client,
                table_cache: Cache::new(0),
                uri: "http://localhost".to_string(),
                namespace_headers: HashMap::new(),
                namespace_context_provider: None,
                tls_config: None,
                sql_client: None,
            }
        }

        pub fn new_mock_with_config<F, T>(handler: F, config: ClientConfig) -> Self
        where
            F: Fn(reqwest::Request) -> http::Response<T> + Send + Sync + 'static,
            T: Into<reqwest::Body>,
        {
            let client = client_with_handler_and_config(handler, config.clone());
            let namespace_context_provider =
                config.header_provider.as_ref().map(|header_provider| {
                    Arc::new(NamespaceHeaderProviderContext {
                        header_provider: Arc::clone(header_provider),
                    }) as Arc<dyn DynamicContextProvider>
                });
            Self {
                client,
                table_cache: Cache::new(0),
                uri: "http://localhost".to_string(),
                namespace_headers: config.extra_headers.clone(),
                namespace_context_provider,
                tls_config: config.tls_config.clone(),
                sql_client: None,
            }
        }
    }
}

impl<S: HttpSend> std::fmt::Display for RemoteDatabase<S> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "RemoteDatabase(host={})", self.client.host())
    }
}

impl From<&CreateTableMode> for &'static str {
    fn from(val: &CreateTableMode) -> Self {
        match val {
            CreateTableMode::Create => "create",
            CreateTableMode::Overwrite => "overwrite",
            CreateTableMode::ExistOk(_) => "exist_ok",
        }
    }
}

fn build_table_identifier(name: &str, namespace: &[String], delimiter: &str) -> String {
    if !namespace.is_empty() {
        let mut parts = namespace.to_vec();
        parts.push(name.to_string());
        parts.join(delimiter)
    } else {
        name.to_string()
    }
}

fn build_namespace_identifier(namespace: &[String], delimiter: &str) -> String {
    if namespace.is_empty() {
        // According to the namespace spec, use delimiter to represent root namespace
        delimiter.to_string()
    } else {
        namespace.join(delimiter)
    }
}

/// Build a secure cache key using length prefixes.
/// This format is completely unambiguous regardless of delimiter or content.
/// Format: [u32_len][namespace1][u32_len][namespace2]...[u32_len][table_name]
/// Returns a hex-encoded string for use as a cache key.
fn build_cache_key(name: &str, namespace: &[String]) -> String {
    let mut key = Vec::new();

    // Add each namespace component with length prefix
    for ns in namespace {
        let bytes = ns.as_bytes();
        key.extend_from_slice(&(bytes.len() as u32).to_le_bytes());
        key.extend_from_slice(bytes);
    }

    // Add table name with length prefix
    let name_bytes = name.as_bytes();
    key.extend_from_slice(&(name_bytes.len() as u32).to_le_bytes());
    key.extend_from_slice(name_bytes);

    // Convert to hex string for use as a cache key
    key.iter().map(|b| format!("{:02x}", b)).collect()
}

#[derive(serde::Deserialize)]
struct RemoteListJobRow {
    job_id: String,
    #[serde(default)]
    table: String,
    #[serde(default)]
    job_type: String,
    #[serde(default)]
    state: String,
    #[serde(default)]
    created_at_millis: i64,
}

#[derive(serde::Deserialize)]
struct RemoteListJobsResponse {
    #[serde(default)]
    jobs: Vec<RemoteListJobRow>,
    #[serde(default)]
    page_token: Option<String>,
}

#[derive(serde::Deserialize)]
struct RemoteListedFunctionVersion {
    definition: FunctionVersion,
}

#[derive(serde::Deserialize)]
struct RemoteListFunctionsResponse {
    #[serde(default)]
    functions: Vec<RemoteListedFunctionVersion>,
    #[serde(default)]
    page_token: Option<String>,
}

#[derive(serde::Deserialize)]
struct RemoteDropFunctionResponse {
    dropped: bool,
}

/// Bound on `list_jobs` page walking; a warning is logged when the listing
/// is truncated at this many pages.
const MAX_LIST_JOBS_PAGES: usize = 100;

#[async_trait]
impl<S: HttpSend> Database for RemoteDatabase<S> {
    fn uri(&self) -> &str {
        &self.uri
    }

    async fn read_consistency(&self) -> Result<ReadConsistency> {
        Err(Error::NotSupported {
            message: "Getting the read consistency of a remote database is not yet supported"
                .to_string(),
        })
    }

    async fn create_function_async(
        &self,
        request: FunctionRegistrationRequest,
    ) -> Result<Job<FunctionVersion>> {
        let req = self.client.post("/v1/functions/create").json(&request);
        let (request_id, response) = self.client.send(req).await?;
        let response = self.client.check_response(&request_id, response).await?;
        let status = response.status();
        let body = response.text().await.err_to_http(request_id.clone())?;
        let job_id = extract_job_id(&body).ok_or_else(|| Error::Http {
            source: "Function registration response did not contain a valid job_id".into(),
            request_id,
            status_code: Some(status),
        })?;
        Ok(Job::new_typed(Box::new(RemoteJob::new(
            self.client.clone(),
            job_id,
        ))))
    }

    async fn get_function(&self, name: &str, version: &str) -> Result<FunctionVersion> {
        let req = self
            .client
            .post("/v1/functions/describe")
            .json(&serde_json::json!({
                "name": name,
                "version": version,
            }));
        let (request_id, response) = self.client.send(req).await?;
        let response = self.client.check_response(&request_id, response).await?;
        response.json().await.err_to_http(request_id)
    }

    async fn list_functions(&self) -> Result<Vec<FunctionVersion>> {
        let mut functions = Vec::new();
        let mut page_token: Option<String> = None;
        let mut seen_page_tokens = HashSet::new();
        loop {
            let mut body = serde_json::json!({ "include_definition": true });
            if let Some(token) = &page_token {
                body["page_token"] = serde_json::Value::String(token.clone());
            }
            let req = self.client.post("/v1/functions/list").json(&body);
            let (request_id, response) = self.client.send(req).await?;
            let response = self.client.check_response(&request_id, response).await?;
            let status = response.status();
            let response: RemoteListFunctionsResponse =
                response.json().await.err_to_http(request_id.clone())?;
            functions.extend(
                response
                    .functions
                    .into_iter()
                    .map(|listed| listed.definition),
            );
            let Some(next_page_token) = response.page_token.filter(|token| !token.is_empty())
            else {
                break;
            };
            if !seen_page_tokens.insert(next_page_token.clone()) {
                return Err(Error::Http {
                    source: "Function listing response repeated a page_token".into(),
                    request_id,
                    status_code: Some(status),
                });
            }
            page_token = Some(next_page_token);
        }
        Ok(functions)
    }

    async fn drop_function(&self, name: &str, version: &str) -> Result<bool> {
        let req = self
            .client
            .post("/v1/functions/drop")
            .json(&serde_json::json!({
                "name": name,
                "version": version,
            }));
        let (request_id, response) = self.client.send(req).await?;
        let response = self.client.check_response(&request_id, response).await?;
        let response: RemoteDropFunctionResponse = response.json().await.err_to_http(request_id)?;
        Ok(response.dropped)
    }

    async fn open_job(&self, job_id: &str) -> Result<Job> {
        let handle = super::job::RemoteJob::new(self.client.clone(), job_id.to_string());
        match crate::job::JobHandle::describe(&handle).await {
            Ok(description) => Ok(Job::opened(Box::new(handle), description)),
            Err(Error::Http {
                status_code: Some(StatusCode::NOT_FOUND),
                ..
            }) => Err(Error::JobNotFound {
                job_id: job_id.to_string(),
            }),
            Err(err) => Err(err),
        }
    }

    async fn list_jobs(&self) -> Result<Vec<JobInfo>> {
        let mut out = Vec::new();
        let mut page_token: Option<String> = None;
        for page in 0..MAX_LIST_JOBS_PAGES {
            let mut body = serde_json::json!({});
            if let Some(token) = &page_token {
                body["page_token"] = serde_json::Value::String(token.clone());
            }
            let req = self.client.post("/v1/jobs/list").json(&body);
            let (request_id, rsp) = self.client.send(req).await?;
            let rsp = self.client.check_response(&request_id, rsp).await?;
            let body: RemoteListJobsResponse = rsp.json().await.err_to_http(request_id)?;
            out.extend(body.jobs.into_iter().map(|row| JobInfo {
                job_id: row.job_id,
                table: row.table,
                job_type: row.job_type,
                state: job_state_to_client(&row.state),
                created_at_millis: row.created_at_millis,
            }));
            page_token = body.page_token;
            if page_token.is_none() {
                break;
            }
            if page + 1 == MAX_LIST_JOBS_PAGES {
                log::warn!(
                    "list_jobs truncated after {} pages ({} jobs)",
                    MAX_LIST_JOBS_PAGES,
                    out.len()
                );
            }
        }
        Ok(out)
    }

    async fn cancel_job(&self, job_id: &str) -> Result<bool> {
        let req = self
            .client
            .post("/v1/jobs/cancel")
            .json(&serde_json::json!({ "job_id": job_id }));
        let (request_id, rsp) = self.client.send(req).await?;
        match self.client.check_response(&request_id, rsp).await {
            Ok(_) => Ok(true),
            Err(Error::Http {
                status_code: Some(StatusCode::NOT_FOUND),
                ..
            }) => Ok(false),
            Err(err) => Err(err),
        }
    }

    async fn execute_query_async(
        &self,
        query: &str,
        default_namespace_path: &[String],
    ) -> Result<crate::sql::Query> {
        let client = self
            .sql_client
            .as_ref()
            .ok_or_else(|| Error::NotSupported {
                message: "SQL is unavailable for this remote database client".to_string(),
            })?;
        client.submit(query, default_namespace_path).await
    }

    async fn describe_query(&self, query_id: uuid::Uuid) -> Result<crate::sql::QueryDescription> {
        let client = self
            .sql_client
            .as_ref()
            .ok_or_else(|| Error::NotSupported {
                message: "SQL is unavailable for this remote database client".to_string(),
            })?;
        client.describe(query_id).await
    }

    async fn table_names(&self, request: TableNamesRequest) -> Result<Vec<String>> {
        let (tables, version) = if request.namespace_path.is_empty() {
            // The flat route resumes after a table name and orders by name, which is exactly
            // what `start_after` means, so the server does the paging.
            let mut req = self.client.get("/v1/table/");
            if let Some(limit) = request.limit {
                req = req.query(&[("limit", limit)]);
            }
            if let Some(ref start_after) = request.start_after {
                req = req.query(&[("page_token", start_after)]);
            }
            let (request_id, rsp) = self.client.send_with_retry(req, None, true).await?;
            let rsp = self.client.check_response(&request_id, rsp).await?;
            let version = parse_server_version(&request_id, &rsp)?;
            let tables = rsp
                .json::<ListTablesResponse>()
                .await
                .err_to_http(request_id)?
                .tables;
            (tables, version)
        } else {
            self.table_names_in_namespace(&request).await?
        };

        for table in &tables {
            let table_identifier =
                build_table_identifier(table, &request.namespace_path, &self.client.id_delimiter);
            let cache_key = build_cache_key(table, &request.namespace_path);
            let remote_table = Arc::new(RemoteTable::new(
                self.client.clone(),
                table.clone(),
                request.namespace_path.clone(),
                table_identifier.clone(),
                version.clone(),
            ));
            self.table_cache.insert(cache_key, remote_table).await;
        }
        Ok(tables)
    }

    async fn list_tables(&self, request: ListTablesRequest) -> Result<ListTablesResponse> {
        let namespace_parts = request.id.as_deref().unwrap_or(&[]);
        let namespace_id = build_namespace_identifier(namespace_parts, &self.client.id_delimiter);
        let mut req = self
            .client
            .get(&format!("/v1/namespace/{}/table/list", namespace_id));

        if let Some(limit) = request.limit {
            req = req.query(&[("limit", limit)]);
        }
        if let Some(ref page_token) = request.page_token {
            req = req.query(&[("page_token", page_token)]);
        }

        let (request_id, rsp) = self.client.send_with_retry(req, None, true).await?;
        let rsp = self.client.check_response(&request_id, rsp).await?;
        let version = parse_server_version(&request_id, &rsp)?;
        let response: ListTablesResponse = rsp.json().await.err_to_http(request_id)?;

        // Cache the tables for future use
        let namespace_vec = namespace_parts.to_vec();
        for table in &response.tables {
            let table_identifier =
                build_table_identifier(table, &namespace_vec, &self.client.id_delimiter);
            let cache_key = build_cache_key(table, &namespace_vec);
            let remote_table = Arc::new(RemoteTable::new(
                self.client.clone(),
                table.clone(),
                namespace_vec.clone(),
                table_identifier.clone(),
                version.clone(),
            ));
            self.table_cache.insert(cache_key, remote_table).await;
        }

        Ok(response)
    }

    async fn create_table(&self, mut request: CreateTableRequest) -> Result<Arc<dyn BaseTable>> {
        let body = stream_as_body(request.data.scan_as_stream())?;

        let identifier = build_table_identifier(
            &request.name,
            &request.namespace_path,
            &self.client.id_delimiter,
        );
        let req = self
            .client
            .post(&format!("/v1/table/{}/create/", identifier))
            .query(&[("mode", Into::<&str>::into(&request.mode))])
            .body(body)
            .header(CONTENT_TYPE, ARROW_STREAM_CONTENT_TYPE);

        let (request_id, rsp) = self.client.send(req).await?;

        if rsp.status() == StatusCode::BAD_REQUEST {
            let body = rsp.text().await.err_to_http(request_id.clone())?;
            if body.contains("already exists") {
                return match request.mode {
                    CreateTableMode::Create => {
                        Err(crate::Error::TableAlreadyExists { name: request.name })
                    }
                    CreateTableMode::ExistOk(callback) => {
                        let req = OpenTableRequest {
                            name: request.name.clone(),
                            namespace_path: request.namespace_path.clone(),
                            index_cache_size: None,
                            lance_read_params: None,
                            location: None,
                            namespace_client: None,
                            managed_versioning: None,
                        };
                        let req = (callback)(req);
                        self.open_table(req).await
                    }

                    // This should not happen, as we explicitly set the mode to overwrite and the server
                    // shouldn't return an error if the table already exists.
                    //
                    // However if the server is an older version that doesn't support the mode parameter,
                    // then we'll get the 400 response.
                    CreateTableMode::Overwrite => Err(crate::Error::Http {
                        source: format!(
                            "unexpected response from server for create mode overwrite: {}",
                            body
                        )
                        .into(),
                        request_id,
                        status_code: Some(StatusCode::BAD_REQUEST),
                    }),
                };
            } else {
                return Err(crate::Error::InvalidInput { message: body });
            }
        }
        let rsp = self.client.check_response(&request_id, rsp).await?;
        let version = parse_server_version(&request_id, &rsp)?;
        let table_identifier = build_table_identifier(
            &request.name,
            &request.namespace_path,
            &self.client.id_delimiter,
        );
        let cache_key = build_cache_key(&request.name, &request.namespace_path);
        let table = Arc::new(RemoteTable::new(
            self.client.clone(),
            request.name.clone(),
            request.namespace_path.clone(),
            table_identifier,
            version,
        ));
        self.table_cache.insert(cache_key, table.clone()).await;

        Ok(table)
    }

    async fn clone_table(&self, request: CloneTableRequest) -> Result<Arc<dyn BaseTable>> {
        let table_identifier = build_table_identifier(
            &request.target_table_name,
            &request.target_namespace_path,
            &self.client.id_delimiter,
        );

        let remote_request = RemoteCloneTableRequest {
            source_location: request.source_uri,
            source_version: request.source_version,
            source_tag: request.source_tag,
            is_shallow: Some(request.is_shallow),
        };

        let req = self
            .client
            .post(&format!("/v1/table/{}/clone", table_identifier.clone()))
            .json(&remote_request);

        let (request_id, rsp) = self.client.send(req).await?;

        let status = rsp.status();
        if status != StatusCode::OK {
            let body = rsp.text().await.err_to_http(request_id.clone())?;
            return Err(crate::Error::Http {
                source: format!("Failed to clone table: {}", body).into(),
                request_id,
                status_code: Some(status),
            });
        }

        let version = parse_server_version(&request_id, &rsp)?;
        let cache_key = build_cache_key(&request.target_table_name, &request.target_namespace_path);
        let table = Arc::new(RemoteTable::new(
            self.client.clone(),
            request.target_table_name.clone(),
            request.target_namespace_path.clone(),
            table_identifier,
            version,
        ));
        self.table_cache.insert(cache_key, table.clone()).await;

        Ok(table)
    }

    async fn open_table(&self, request: OpenTableRequest) -> Result<Arc<dyn BaseTable>> {
        let identifier = build_table_identifier(
            &request.name,
            &request.namespace_path,
            &self.client.id_delimiter,
        );
        let cache_key = build_cache_key(&request.name, &request.namespace_path);

        // We describe the table to confirm it exists before moving on.
        if let Some(table) = self.table_cache.get(&cache_key).await {
            Ok(table.clone())
        } else {
            let req = self
                .client
                .post(&format!("/v1/table/{}/describe/", identifier));
            let (request_id, rsp) = self.client.send_with_retry(req, None, true).await?;
            let rsp =
                RemoteTable::<S>::handle_table_not_found(&request.name, rsp, &request_id).await?;
            let rsp = self.client.check_response(&request_id, rsp).await?;
            let version = parse_server_version(&request_id, &rsp)?;
            let describe_body = rsp.text().await.ok();
            let table_identifier = build_table_identifier(
                &request.name,
                &request.namespace_path,
                &self.client.id_delimiter,
            );
            let table = Arc::new(RemoteTable::new(
                self.client.clone(),
                request.name.clone(),
                request.namespace_path.clone(),
                table_identifier,
                version,
            ));
            // This describe already carries the schema, so hand it to the table
            // instead of making the first schema read fetch it again. A version or
            // branch pin applied after this invalidates the cache.
            if let Some(body) = &describe_body {
                table.seed_schema(body);
            }
            let cache_key = build_cache_key(&request.name, &request.namespace_path);
            self.table_cache.insert(cache_key, table.clone()).await;
            Ok(table)
        }
    }

    async fn rename_table(
        &self,
        current_name: &str,
        new_name: &str,
        cur_namespace_path: &[String],
        new_namespace_path: &[String],
    ) -> Result<()> {
        let current_identifier =
            build_table_identifier(current_name, cur_namespace_path, &self.client.id_delimiter);
        let current_cache_key = build_cache_key(current_name, cur_namespace_path);
        let new_cache_key = build_cache_key(new_name, new_namespace_path);

        let mut body = serde_json::json!({ "new_table_name": new_name });
        if !new_namespace_path.is_empty() {
            body["new_namespace"] = serde_json::Value::Array(
                new_namespace_path
                    .iter()
                    .map(|s| serde_json::Value::String(s.clone()))
                    .collect(),
            );
        }
        let req = self
            .client
            .post(&format!("/v1/table/{}/rename/", current_identifier))
            .json(&body);
        let (request_id, resp) = self.client.send(req).await?;
        self.client.check_response(&request_id, resp).await?;
        let table = self.table_cache.remove(&current_cache_key).await;
        if let Some(table) = table {
            self.table_cache.insert(new_cache_key, table).await;
        }
        Ok(())
    }

    async fn drop_table(&self, name: &str, namespace_path: &[String]) -> Result<()> {
        self.submit_drop_table(name, namespace_path)
            .await
            .map(|_| ())
    }

    async fn drop_table_async(&self, name: &str, namespace_path: &[String]) -> Result<Job> {
        let (request_id, response) = self.submit_drop_table(name, namespace_path).await?;
        let status = response.status();
        let body = response.text().await.err_to_http(request_id.clone())?;
        let job_id = extract_job_id(&body);
        Ok(match job_id {
            Some(job_id) => Job::new(Box::new(RemoteJob::new(self.client.clone(), job_id))),
            None if status == StatusCode::ACCEPTED => {
                return Err(Error::Http {
                    source: "asynchronous drop-table response did not contain a valid job_id"
                        .into(),
                    request_id,
                    status_code: Some(status),
                });
            }
            None => Job::new_done(),
        })
    }

    async fn drop_all_tables(&self, namespace_path: &[String]) -> Result<()> {
        // TODO: Implement namespace-aware drop_all_tables
        let _namespace_path = namespace_path; // Suppress unused warning for now
        Err(crate::Error::NotSupported {
            message: "Dropping all tables is not currently supported in the remote API".to_string(),
        })
    }

    async fn list_namespaces(
        &self,
        request: ListNamespacesRequest,
    ) -> Result<ListNamespacesResponse> {
        let namespace_parts = request.id.as_deref().unwrap_or(&[]);
        let namespace_id = build_namespace_identifier(namespace_parts, &self.client.id_delimiter);
        let mut req = self
            .client
            .get(&format!("/v1/namespace/{}/list", namespace_id));
        if let Some(limit) = request.limit {
            req = req.query(&[("limit", limit)]);
        }
        if let Some(ref page_token) = request.page_token {
            req = req.query(&[("page_token", page_token)]);
        }

        let (request_id, resp) = self.client.send(req).await?;
        let resp = self.client.check_response(&request_id, resp).await?;

        resp.json().await.err_to_http(request_id)
    }

    async fn create_namespace(
        &self,
        request: CreateNamespaceRequest,
    ) -> Result<CreateNamespaceResponse> {
        let namespace_parts = request.id.as_deref().unwrap_or(&[]);
        let namespace_id = build_namespace_identifier(namespace_parts, &self.client.id_delimiter);
        let mut req = self
            .client
            .post(&format!("/v1/namespace/{}/create", namespace_id));

        // Build request body with mode and properties if present
        #[derive(serde::Serialize)]
        struct CreateNamespaceRequestBody {
            #[serde(skip_serializing_if = "Option::is_none")]
            mode: Option<String>,
            #[serde(skip_serializing_if = "Option::is_none")]
            properties: Option<HashMap<String, String>>,
        }

        let body = CreateNamespaceRequestBody {
            mode: request.mode.as_ref().map(|m| format!("{:?}", m)),
            properties: request.properties,
        };

        req = req.json(&body);
        let (request_id, resp) = self.client.send(req).await?;
        let resp = self.client.check_response(&request_id, resp).await?;

        resp.json().await.err_to_http(request_id)
    }

    async fn drop_namespace(&self, request: DropNamespaceRequest) -> Result<DropNamespaceResponse> {
        let namespace_parts = request.id.as_deref().unwrap_or(&[]);
        let namespace_id = build_namespace_identifier(namespace_parts, &self.client.id_delimiter);
        let mut req = self
            .client
            .post(&format!("/v1/namespace/{}/drop", namespace_id));

        // Build request body with mode and behavior if present
        #[derive(serde::Serialize)]
        struct DropNamespaceRequestBody {
            #[serde(skip_serializing_if = "Option::is_none")]
            mode: Option<String>,
            #[serde(skip_serializing_if = "Option::is_none")]
            behavior: Option<String>,
        }

        let body = DropNamespaceRequestBody {
            mode: request.mode.as_ref().map(|m| format!("{:?}", m)),
            behavior: request.behavior.as_ref().map(|b| format!("{:?}", b)),
        };

        req = req.json(&body);
        let (request_id, resp) = self.client.send(req).await?;
        let resp = self.client.check_response(&request_id, resp).await?;

        resp.json().await.err_to_http(request_id)
    }

    async fn describe_namespace(
        &self,
        request: DescribeNamespaceRequest,
    ) -> Result<DescribeNamespaceResponse> {
        let namespace_parts = request.id.as_deref().unwrap_or(&[]);
        let namespace_id = build_namespace_identifier(namespace_parts, &self.client.id_delimiter);
        let req = self
            .client
            .post(&format!("/v1/namespace/{}/describe", namespace_id))
            .json(&DescribeNamespaceRequest::default());

        let (request_id, resp) = self.client.send(req).await?;
        let resp = self.client.check_response(&request_id, resp).await?;

        resp.json().await.err_to_http(request_id)
    }

    fn as_any(&self) -> &dyn std::any::Any {
        self
    }

    async fn namespace_client(&self) -> Result<Arc<dyn lance_namespace::LanceNamespace>> {
        // Create a RestNamespace pointing to the same remote host with the same authentication headers
        let mut builder = lance_namespace_impls::RestNamespaceBuilder::new(self.client.host())
            .delimiter(&self.client.id_delimiter)
            .headers(self.namespace_headers.clone());

        if let Some(context_provider) = &self.namespace_context_provider {
            builder = builder.context_provider(Arc::clone(context_provider));
        }

        // Apply mTLS configuration if present
        if let Some(tls_config) = &self.tls_config {
            if let Some(cert_file) = &tls_config.cert_file {
                builder = builder.cert_file(cert_file);
            }
            if let Some(key_file) = &tls_config.key_file {
                builder = builder.key_file(key_file);
            }
            if let Some(ssl_ca_cert) = &tls_config.ssl_ca_cert {
                builder = builder.ssl_ca_cert(ssl_ca_cert);
            }
            builder = builder.assert_hostname(tls_config.assert_hostname);
        }

        let namespace = builder.build();
        Ok(Arc::new(namespace) as Arc<dyn lance_namespace::LanceNamespace>)
    }

    async fn namespace_client_config(&self) -> Result<(String, HashMap<String, String>)> {
        if self.namespace_context_provider.is_some() {
            return Err(Error::NotSupported {
                message:
                    "Cannot export a namespace client config when dynamic headers are configured; use LanceDB connection namespace methods instead"
                        .to_string(),
            });
        }

        let mut properties = HashMap::new();
        properties.insert("uri".to_string(), self.client.host().to_string());
        properties.insert("delimiter".to_string(), self.client.id_delimiter.clone());
        for (key, value) in &self.namespace_headers {
            properties.insert(format!("header.{}", key), value.clone());
        }
        // Add TLS configuration if present
        if let Some(tls_config) = &self.tls_config {
            if let Some(cert_file) = &tls_config.cert_file {
                properties.insert("tls.cert_file".to_string(), cert_file.clone());
            }
            if let Some(key_file) = &tls_config.key_file {
                properties.insert("tls.key_file".to_string(), key_file.clone());
            }
            if let Some(ssl_ca_cert) = &tls_config.ssl_ca_cert {
                properties.insert("tls.ssl_ca_cert".to_string(), ssl_ca_cert.clone());
            }
            properties.insert(
                "tls.assert_hostname".to_string(),
                tls_config.assert_hostname.to_string(),
            );
        }
        Ok(("rest".to_string(), properties))
    }
}

/// RemoteOptions contains a subset of StorageOptions that are compatible with Remote LanceDB connections
#[derive(Clone, Debug, Default)]
pub struct RemoteOptions(pub HashMap<String, String>);

impl RemoteOptions {
    pub fn new(options: HashMap<String, String>) -> Self {
        Self(options)
    }
}

impl From<StorageOptions> for RemoteOptions {
    fn from(options: StorageOptions) -> Self {
        let supported_opts = vec!["account_name", "azure_storage_account_name"];
        let mut filtered = HashMap::new();
        for opt in supported_opts {
            if let Some(v) = options.0.get(opt) {
                filtered.insert(opt.to_string(), v.clone());
            }
        }
        Self::new(filtered)
    }
}

#[cfg(test)]
mod tests {
    use super::{NamespaceHeaderProviderContext, build_cache_key};
    use std::collections::HashMap;
    use std::sync::atomic::{AtomicUsize, Ordering};
    use std::sync::{Arc, OnceLock};

    use arrow_array::{Int32Array, RecordBatch};
    use arrow_schema::{DataType, Field, Schema};
    use lance_namespace_impls::{DynamicContextProvider, OperationInfo};

    use crate::connection::ConnectBuilder;
    use crate::{
        Connection, Error,
        database::CreateTableMode,
        job::JobEventsRequest,
        remote::{ARROW_STREAM_CONTENT_TYPE, ClientConfig, HeaderProvider, JSON_CONTENT_TYPE},
    };

    #[test]
    fn test_cache_key_security() {
        // Test that cache keys are unique regardless of delimiter manipulation

        // Case 1: Different delimiters should not affect cache key
        let key1 = build_cache_key("table1", &["ns1".to_string(), "ns2".to_string()]);
        let key2 = build_cache_key("table1", &["ns1$ns2".to_string()]);
        assert_ne!(
            key1, key2,
            "Cache keys should differ for different namespace structures"
        );

        // Case 2: Table name containing delimiter should not cause collision
        let key3 = build_cache_key("ns2$table1", &["ns1".to_string()]);
        assert_ne!(
            key1, key3,
            "Cache key should be different when table name contains delimiter"
        );

        // Case 3: Empty namespace vs namespace with empty string
        let key4 = build_cache_key("table1", &[]);
        let key5 = build_cache_key("table1", &["".to_string()]);
        assert_ne!(
            key4, key5,
            "Empty namespace should differ from namespace with empty string"
        );

        // Case 4: Verify same inputs produce same key (consistency)
        let key6 = build_cache_key("table1", &["ns1".to_string(), "ns2".to_string()]);
        assert_eq!(key1, key6, "Same inputs should produce same cache key");
    }

    #[tokio::test]
    async fn test_retries() {
        // We'll record the request_id here, to check it matches the one in the error.
        let seen_request_id = Arc::new(OnceLock::new());
        let seen_request_id_ref = seen_request_id.clone();
        let conn = Connection::new_with_handler(move |request| {
            // Request id should be the same on each retry.
            let request_id = request.headers()["x-request-id"]
                .to_str()
                .unwrap()
                .to_string();
            let seen_id = seen_request_id_ref.get_or_init(|| request_id.clone());
            assert_eq!(&request_id, seen_id);

            http::Response::builder()
                .status(500)
                .body("internal server error")
                .unwrap()
        });
        let result = conn.table_names().execute().await;
        if let Err(Error::Retry {
            request_id,
            request_failures,
            max_request_failures,
            source,
            ..
        }) = result
        {
            let expected_id = seen_request_id.get().unwrap();
            assert_eq!(&request_id, expected_id);
            assert_eq!(request_failures, max_request_failures);
            assert!(
                source.to_string().contains("internal server error"),
                "source: {:?}",
                source
            );
        } else {
            panic!("unexpected result: {:?}", result);
        };
    }

    #[tokio::test]
    async fn test_table_names() {
        let conn = Connection::new_with_handler(|request| {
            assert_eq!(request.method(), &reqwest::Method::GET);
            assert_eq!(request.url().path(), "/v1/table/");
            assert_eq!(request.url().query(), None);

            http::Response::builder()
                .status(200)
                .body(r#"{"tables": ["table1", "table2"]}"#)
                .unwrap()
        });
        let names = conn.table_names().execute().await.unwrap();
        assert_eq!(names, vec!["table1", "table2"]);
    }

    #[tokio::test]
    async fn test_table_names_in_a_namespace_never_invents_a_page_token() {
        // The namespace route's token belongs to the store, so `table_names` cannot build one
        // from `start_after`. It walks the namespace on the server's own tokens and applies the
        // name semantics itself, which is what keeps it working either side of the change.
        let page = Arc::new(AtomicUsize::new(0));
        let conn = Connection::new_with_handler(move |request| {
            assert_eq!(request.url().path(), "/v1/namespace/ns/table/list");
            let query = request.url().query().unwrap_or("");
            assert!(
                !query.contains("page_token=users"),
                "a table name must never be sent as a page token: {query}"
            );
            match page.fetch_add(1, Ordering::SeqCst) {
                0 => {
                    assert!(
                        !query.contains("page_token"),
                        "the walk starts with no token"
                    );
                    http::Response::builder()
                        .status(200)
                        .body(r#"{"tables": ["users", "orders"], "page_token": "opaque-1"}"#)
                        .unwrap()
                }
                _ => {
                    assert!(query.contains("page_token=opaque-1"));
                    http::Response::builder()
                        .status(200)
                        .body(r#"{"tables": ["widgets"]}"#)
                        .unwrap()
                }
            }
        });

        let names = conn
            .table_names()
            .namespace(vec!["ns".to_string()])
            .start_after("users")
            .execute()
            .await
            .unwrap();
        // Name order, resumed after "users": "orders" sorts before it and is dropped.
        assert_eq!(names, vec!["widgets"]);
    }

    #[tokio::test]
    async fn test_table_names_in_a_namespace_stops_on_a_repeated_token() {
        // A server that handed back the token it was given would never finish the walk.
        let conn = Connection::new_with_handler(|_request| {
            http::Response::builder()
                .status(200)
                .body(r#"{"tables": ["a"], "page_token": "same"}"#)
                .unwrap()
        });

        let names = conn
            .table_names()
            .namespace(vec!["ns".to_string()])
            .execute()
            .await
            .unwrap();
        // The guard bounds the walk instead of letting it run forever. The repeat is the
        // server breaking the token contract and is not papered over here.
        assert_eq!(names, vec!["a", "a"]);
    }

    #[tokio::test]
    async fn test_table_names_in_a_namespace_stops_on_an_empty_token() {
        // An empty token ends the listing. Sending it back would ask a server that reads it
        // as "start from the beginning" for the first page a second time, and every name on
        // that page would be collected twice.
        let requests = Arc::new(AtomicUsize::new(0));
        let seen = requests.clone();
        let conn = Connection::new_with_handler(move |request| {
            seen.fetch_add(1, Ordering::SeqCst);
            assert!(
                !request.url().query().unwrap_or("").contains("page_token"),
                "an empty token must never be sent back"
            );
            http::Response::builder()
                .status(200)
                .body(r#"{"tables": ["a"], "page_token": ""}"#)
                .unwrap()
        });

        let names = conn
            .table_names()
            .namespace(vec!["ns".to_string()])
            .execute()
            .await
            .unwrap();
        assert_eq!(names, vec!["a"]);
        assert_eq!(requests.load(Ordering::SeqCst), 1);
    }

    #[tokio::test]
    async fn test_table_names_pagination() {
        let conn = Connection::new_with_handler(|request| {
            assert_eq!(request.method(), &reqwest::Method::GET);
            assert_eq!(request.url().path(), "/v1/table/");
            assert!(request.url().query().unwrap().contains("limit=2"));
            assert!(request.url().query().unwrap().contains("page_token=table2"));

            http::Response::builder()
                .status(200)
                .body(r#"{"tables": ["table3", "table4"], "page_token": "token"}"#)
                .unwrap()
        });
        let names = conn
            .table_names()
            .start_after("table2")
            .limit(2)
            .execute()
            .await
            .unwrap();
        assert_eq!(names, vec!["table3", "table4"]);
    }

    #[tokio::test]
    async fn test_open_table() {
        let conn = Connection::new_with_handler(|request| {
            assert_eq!(request.method(), &reqwest::Method::POST);
            assert_eq!(request.url().path(), "/v1/table/table1/describe/");
            assert_eq!(request.url().query(), None);

            http::Response::builder()
                .status(200)
                .body(r#"{"table": "table1"}"#)
                .unwrap()
        });
        let table = conn.open_table("table1").execute().await.unwrap();
        assert_eq!(table.name(), "table1");

        // Storage options should be ignored.
        let table = conn
            .open_table("table1")
            .storage_option("key", "value")
            .execute()
            .await
            .unwrap();
        assert_eq!(table.name(), "table1");
    }

    #[tokio::test]
    async fn test_open_table_seeds_the_schema_from_its_describe() {
        let describe_calls = Arc::new(AtomicUsize::new(0));
        let counted = describe_calls.clone();
        let conn = Connection::new_with_handler(move |request| {
            assert_eq!(request.url().path(), "/v1/table/table1/describe/");
            counted.fetch_add(1, Ordering::SeqCst);
            http::Response::builder()
                .status(200)
                .body(
                    r#"{"version": 1, "schema": {"fields": [
                        {"name": "id", "type": {"type": "int64"}, "nullable": false}
                    ]}}"#
                        .to_string(),
                )
                .unwrap()
        });

        let table = conn.open_table("table1").execute().await.unwrap();
        let schema = table.schema().await.unwrap();

        assert_eq!(schema.field(0).name(), "id");
        assert_eq!(describe_calls.load(Ordering::SeqCst), 1);
    }

    #[tokio::test]
    async fn test_open_table_survives_a_describe_body_it_cannot_parse() {
        let conn = Connection::new_with_handler(|request| {
            assert_eq!(request.url().path(), "/v1/table/table1/describe/");
            http::Response::builder()
                .status(200)
                .body(r#"{"table": "table1"}"#.to_string())
                .unwrap()
        });

        let table = conn.open_table("table1").execute().await.unwrap();

        assert_eq!(table.name(), "table1");
    }

    #[tokio::test]
    async fn test_open_table_branch_and_version() {
        let conn = Connection::new_with_handler(|request| {
            let body = if request.url().path() == "/v1/table/t/branches/list/" {
                // checkout_branch validates the branch exists via list_branches.
                r#"{"branches":{"exp":{"parentVersion":1,"createAt":1,"manifestSize":1}}}"#
            } else {
                // describe (table open + version/branch validation)
                r#"{"table": "t", "version": 2, "schema": {"fields": [
                    {"name": "a", "type": { "type": "int32" }, "nullable": false}
                ]}}"#
            };
            http::Response::builder().status(200).body(body).unwrap()
        });

        // version-only (and "main" + version) time-travel the main chain
        let v2 = conn.open_table("t").version(2).execute().await.unwrap();
        assert_eq!(v2.current_branch(), None);
        let main_v2 = conn
            .open_table("t")
            .branch("main")
            .version(2)
            .execute()
            .await
            .unwrap();
        assert_eq!(main_v2.current_branch(), None);

        // a non-main branch opens a handle scoped to that branch
        let exp = conn.open_table("t").branch("exp").execute().await.unwrap();
        assert_eq!(exp.current_branch(), Some("exp".to_string()));
        let exp_v2 = conn
            .open_table("t")
            .branch("exp")
            .version(2)
            .execute()
            .await
            .unwrap();
        assert_eq!(exp_v2.current_branch(), Some("exp".to_string()));
    }

    #[tokio::test]
    async fn test_open_table_not_found() {
        let conn = Connection::new_with_handler(|_| {
            http::Response::builder()
                .status(404)
                .body("table not found")
                .unwrap()
        });
        let result = conn.open_table("table1").execute().await;
        assert!(result.is_err());
        assert!(matches!(result, Err(crate::Error::TableNotFound { .. })));
    }

    #[tokio::test]
    async fn test_create_table() {
        let conn = Connection::new_with_handler(|request| {
            assert_eq!(request.method(), &reqwest::Method::POST);
            assert_eq!(request.url().path(), "/v1/table/table1/create/");
            assert_eq!(
                request
                    .headers()
                    .get(reqwest::header::CONTENT_TYPE)
                    .unwrap(),
                ARROW_STREAM_CONTENT_TYPE.as_bytes()
            );

            http::Response::builder().status(200).body("").unwrap()
        });
        let data = RecordBatch::try_new(
            Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])),
            vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
        )
        .unwrap();
        let table = conn.create_table("table1", data).execute().await.unwrap();
        assert_eq!(table.name(), "table1");
    }

    #[tokio::test]
    async fn test_create_table_already_exists() {
        let conn = Connection::new_with_handler(|_| {
            http::Response::builder()
                .status(400)
                .body("table table1 already exists")
                .unwrap()
        });
        let data = RecordBatch::try_new(
            Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])),
            vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
        )
        .unwrap();
        let result = conn.create_table("table1", data).execute().await;
        assert!(result.is_err());
        assert!(
            matches!(result, Err(crate::Error::TableAlreadyExists { name }) if name == "table1")
        );
    }

    #[tokio::test]
    async fn test_create_table_modes() {
        let test_cases = [
            (None, "mode=create"),
            (Some(CreateTableMode::Create), "mode=create"),
            (Some(CreateTableMode::Overwrite), "mode=overwrite"),
            (
                Some(CreateTableMode::ExistOk(Box::new(|b| b))),
                "mode=exist_ok",
            ),
        ];

        for (mode, expected_query_string) in test_cases {
            let conn = Connection::new_with_handler(move |request| {
                assert_eq!(request.method(), &reqwest::Method::POST);
                assert_eq!(request.url().path(), "/v1/table/table1/create/");
                assert_eq!(request.url().query(), Some(expected_query_string));

                http::Response::builder().status(200).body("").unwrap()
            });

            let data = RecordBatch::try_new(
                Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])),
                vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
            )
            .unwrap();
            let mut builder = conn.create_table("table1", data.clone());
            if let Some(mode) = mode {
                builder = builder.mode(mode);
            }
            builder.execute().await.unwrap();
        }

        // check that the open table callback is called with exist_ok
        let conn = Connection::new_with_handler(|request| match request.url().path() {
            "/v1/table/table1/create/" => http::Response::builder()
                .status(400)
                .body("Table table1 already exists")
                .unwrap(),
            "/v1/table/table1/describe/" => http::Response::builder().status(200).body("").unwrap(),
            _ => {
                panic!("unexpected path: {:?}", request.url().path());
            }
        });
        let data = RecordBatch::try_new(
            Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])),
            vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
        )
        .unwrap();

        let called: Arc<OnceLock<bool>> = Arc::new(OnceLock::new());
        let called_in_cb = called.clone();
        conn.create_table("table1", data)
            .mode(CreateTableMode::ExistOk(Box::new(move |b| {
                called_in_cb.clone().set(true).unwrap();
                b
            })))
            .execute()
            .await
            .unwrap();

        let called = *called.get().unwrap_or(&false);
        assert!(called);
    }

    #[tokio::test]
    async fn test_create_table_empty() {
        let conn = Connection::new_with_handler(|request| {
            assert_eq!(request.method(), &reqwest::Method::POST);
            assert_eq!(request.url().path(), "/v1/table/table1/create/");
            assert_eq!(
                request
                    .headers()
                    .get(reqwest::header::CONTENT_TYPE)
                    .unwrap(),
                ARROW_STREAM_CONTENT_TYPE.as_bytes()
            );

            http::Response::builder().status(200).body("").unwrap()
        });
        let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]));
        conn.create_empty_table("table1", schema)
            .execute()
            .await
            .unwrap();
    }

    #[tokio::test]
    async fn test_drop_table() {
        let conn = Connection::new_with_handler(|request| {
            assert_eq!(request.method(), &reqwest::Method::POST);
            assert_eq!(request.url().path(), "/v1/table/table1/drop/");
            assert_eq!(request.url().query(), None);
            assert!(request.body().is_none());

            http::Response::builder().status(200).body("").unwrap()
        });
        conn.drop_table("table1", &[]).await.unwrap();
        // NOTE: the API will return 200 even if the table does not exist. So we shouldn't expect 404.
    }

    #[tokio::test]
    async fn test_drop_table_does_not_read_response_body() {
        let conn = Connection::new_with_handler(|_| {
            http::Response::builder()
                .status(200)
                .body(vec![0xff])
                .unwrap()
        });

        conn.drop_table("table1", &[]).await.unwrap();
    }

    #[tokio::test]
    async fn test_drop_table_async_returns_job() {
        let conn = Connection::new_with_handler(|request| {
            assert_eq!(request.method(), &reqwest::Method::POST);
            assert_eq!(request.url().path(), "/v1/table/table1/drop/");
            http::Response::builder()
                .status(202)
                .body(r#"{"job_id":"drop-job-123"}"#)
                .unwrap()
        });

        let job = conn.drop_table_async("table1", &[]).await.unwrap();
        assert_eq!(job.id(), Some("drop-job-123"));
    }

    #[tokio::test]
    async fn test_drop_table_async_old_server_returns_done_job() {
        let conn = Connection::new_with_handler(|_| {
            http::Response::builder().status(200).body("").unwrap()
        });

        let job = conn.drop_table_async("table1", &[]).await.unwrap();
        assert_eq!(job.id(), None);
        assert_eq!(job.status().await.unwrap(), "finished");
    }

    #[tokio::test]
    async fn test_drop_table_async_rejects_accepted_response_without_job_id() {
        let conn = Connection::new_with_handler(|_| {
            http::Response::builder().status(202).body("{}").unwrap()
        });

        let error = conn.drop_table_async("table1", &[]).await.err().unwrap();
        assert!(error.to_string().contains("valid job_id"));
    }

    #[tokio::test]
    async fn test_drop_table_async_rejects_empty_job_id() {
        let conn = Connection::new_with_handler(|_| {
            http::Response::builder()
                .status(202)
                .body(r#"{"job_id":""}"#)
                .unwrap()
        });

        let error = conn.drop_table_async("table1", &[]).await.err().unwrap();
        assert!(error.to_string().contains("valid job_id"));
    }

    #[tokio::test]
    async fn test_rename_table() {
        let conn = Connection::new_with_handler(|request| {
            assert_eq!(request.method(), &reqwest::Method::POST);
            assert_eq!(request.url().path(), "/v1/table/table1/rename/");
            assert_eq!(
                request.headers().get("Content-Type").unwrap(),
                JSON_CONTENT_TYPE
            );

            let body = request.body().unwrap().as_bytes().unwrap();
            let body: serde_json::Value = serde_json::from_slice(body).unwrap();
            assert_eq!(body["new_table_name"], "table2");

            http::Response::builder().status(200).body("").unwrap()
        });
        conn.rename_table("table1", "table2", &[], &[])
            .await
            .unwrap();
    }

    #[tokio::test]
    async fn test_connect_remote_options() {
        let db_uri = "db://my-container/my-prefix";
        let _ = ConnectBuilder::new(db_uri)
            .region("us-east-1")
            .api_key("my-api-key")
            .storage_options(vec![("azure_storage_account_name", "my-storage-account")])
            .execute()
            .await
            .unwrap();
    }

    #[tokio::test]
    async fn test_table_names_with_root_namespace() {
        // When namespace is empty (root namespace), should use /v1/table/ for backwards compatibility
        let conn = Connection::new_with_handler(|request| {
            assert_eq!(request.method(), &reqwest::Method::GET);
            assert_eq!(request.url().path(), "/v1/table/");
            assert_eq!(request.url().query(), None);

            http::Response::builder()
                .status(200)
                .body(r#"{"tables": ["table1", "table2"]}"#)
                .unwrap()
        });
        let names = conn
            .table_names()
            .namespace(vec![])
            .execute()
            .await
            .unwrap();
        assert_eq!(names, vec!["table1", "table2"]);
    }

    #[tokio::test]
    async fn test_table_names_with_namespace() {
        // When namespace is non-empty, should use /v1/namespace/{id}/table/list
        let conn = Connection::new_with_handler(|request| {
            assert_eq!(request.method(), &reqwest::Method::GET);
            assert_eq!(request.url().path(), "/v1/namespace/test/table/list");
            assert_eq!(request.url().query(), None);

            http::Response::builder()
                .status(200)
                .body(r#"{"tables": ["table1", "table2"]}"#)
                .unwrap()
        });
        let names = conn
            .table_names()
            .namespace(vec!["test".to_string()])
            .execute()
            .await
            .unwrap();
        assert_eq!(names, vec!["table1", "table2"]);
    }

    #[tokio::test]
    async fn test_table_names_with_nested_namespace() {
        // When namespace is vec!["ns1", "ns2"], should use /v1/namespace/ns1$ns2/table/list
        let conn = Connection::new_with_handler(|request| {
            assert_eq!(request.method(), &reqwest::Method::GET);
            assert_eq!(request.url().path(), "/v1/namespace/ns1$ns2/table/list");
            assert_eq!(request.url().query(), None);

            http::Response::builder()
                .status(200)
                .body(r#"{"tables": ["ns1$ns2$table1", "ns1$ns2$table2"]}"#)
                .unwrap()
        });
        let names = conn
            .table_names()
            .namespace(vec!["ns1".to_string(), "ns2".to_string()])
            .execute()
            .await
            .unwrap();
        assert_eq!(names, vec!["ns1$ns2$table1", "ns1$ns2$table2"]);
    }

    #[tokio::test]
    async fn test_open_table_with_namespace() {
        let conn = Connection::new_with_handler(|request| {
            assert_eq!(request.method(), &reqwest::Method::POST);
            assert_eq!(request.url().path(), "/v1/table/ns1$ns2$table1/describe/");
            assert_eq!(request.url().query(), None);

            http::Response::builder()
                .status(200)
                .body(r#"{"table": "table1"}"#)
                .unwrap()
        });
        let table = conn
            .open_table("table1")
            .namespace(vec!["ns1".to_string(), "ns2".to_string()])
            .execute()
            .await
            .unwrap();
        assert_eq!(table.name(), "table1");
    }

    #[tokio::test]
    async fn test_create_table_with_namespace() {
        let conn = Connection::new_with_handler(|request| {
            assert_eq!(request.method(), &reqwest::Method::POST);
            assert_eq!(request.url().path(), "/v1/table/ns1$table1/create/");
            assert_eq!(
                request
                    .headers()
                    .get(reqwest::header::CONTENT_TYPE)
                    .unwrap(),
                ARROW_STREAM_CONTENT_TYPE.as_bytes()
            );

            http::Response::builder().status(200).body("").unwrap()
        });
        let data = RecordBatch::try_new(
            Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])),
            vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
        )
        .unwrap();
        let table = conn
            .create_table("table1", data)
            .namespace(vec!["ns1".to_string()])
            .execute()
            .await
            .unwrap();
        assert_eq!(table.name(), "table1");
    }

    #[tokio::test]
    async fn test_drop_table_with_namespace() {
        let conn = Connection::new_with_handler(|request| {
            assert_eq!(request.method(), &reqwest::Method::POST);
            assert_eq!(request.url().path(), "/v1/table/ns1$ns2$table1/drop/");
            assert_eq!(request.url().query(), None);
            assert!(request.body().is_none());

            http::Response::builder().status(200).body("").unwrap()
        });
        conn.drop_table("table1", &["ns1".to_string(), "ns2".to_string()])
            .await
            .unwrap();
    }

    #[tokio::test]
    async fn test_rename_table_with_namespace() {
        let conn = Connection::new_with_handler(|request| {
            assert_eq!(request.method(), &reqwest::Method::POST);
            assert_eq!(request.url().path(), "/v1/table/ns1$table1/rename/");
            assert_eq!(
                request.headers().get("Content-Type").unwrap(),
                JSON_CONTENT_TYPE
            );

            let body = request.body().unwrap().as_bytes().unwrap();
            let body: serde_json::Value = serde_json::from_slice(body).unwrap();
            assert_eq!(body["new_table_name"], "table2");
            assert_eq!(body["new_namespace"], serde_json::json!(["ns2"]));

            http::Response::builder().status(200).body("").unwrap()
        });
        conn.rename_table(
            "table1",
            "table2",
            &["ns1".to_string()],
            &["ns2".to_string()],
        )
        .await
        .unwrap();
    }

    #[tokio::test]
    async fn test_create_empty_table_with_namespace() {
        let conn = Connection::new_with_handler(|request| {
            assert_eq!(request.method(), &reqwest::Method::POST);
            assert_eq!(request.url().path(), "/v1/table/prod$data$metrics/create/");
            assert_eq!(
                request
                    .headers()
                    .get(reqwest::header::CONTENT_TYPE)
                    .unwrap(),
                ARROW_STREAM_CONTENT_TYPE.as_bytes()
            );

            http::Response::builder().status(200).body("").unwrap()
        });
        let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]));
        conn.create_empty_table("metrics", schema)
            .namespace(vec!["prod".to_string(), "data".to_string()])
            .execute()
            .await
            .unwrap();
    }

    #[tokio::test]
    async fn test_header_provider_in_request() {
        // Test HeaderProvider implementation that adds custom headers
        #[derive(Debug, Clone)]
        struct TestHeaderProvider {
            headers: HashMap<String, String>,
        }

        #[async_trait::async_trait]
        impl HeaderProvider for TestHeaderProvider {
            async fn get_headers(&self) -> crate::Result<HashMap<String, String>> {
                Ok(self.headers.clone())
            }
        }

        // Create a test header provider with custom headers
        let mut headers = HashMap::new();
        headers.insert("X-Custom-Auth".to_string(), "test-token".to_string());
        headers.insert("X-Request-Id".to_string(), "test-123".to_string());
        let provider = Arc::new(TestHeaderProvider { headers }) as Arc<dyn HeaderProvider>;

        // Create client config with the header provider
        let client_config = ClientConfig {
            header_provider: Some(provider),
            ..Default::default()
        };

        // Create connection with handler that verifies the headers are present
        let conn = Connection::new_with_handler_and_config(
            move |request| {
                // Verify that our custom headers are present
                assert_eq!(
                    request.headers().get("X-Custom-Auth").unwrap(),
                    "test-token"
                );
                assert_eq!(request.headers().get("X-Request-Id").unwrap(), "test-123");

                // Also check standard headers are still there
                assert_eq!(request.method(), &reqwest::Method::GET);
                assert_eq!(request.url().path(), "/v1/table/");

                http::Response::builder()
                    .status(200)
                    .body(r#"{"tables": ["table1", "table2"]}"#)
                    .unwrap()
            },
            client_config,
        );

        // Make a request that should include the custom headers
        let names = conn.table_names().execute().await.unwrap();
        assert_eq!(names, vec!["table1", "table2"]);
    }

    #[tokio::test]
    async fn test_header_provider_error_handling() {
        // Test HeaderProvider that returns an error
        #[derive(Debug)]
        struct ErrorHeaderProvider;

        #[async_trait::async_trait]
        impl HeaderProvider for ErrorHeaderProvider {
            async fn get_headers(&self) -> crate::Result<HashMap<String, String>> {
                Err(crate::Error::Runtime {
                    message: "Failed to fetch auth token".to_string(),
                })
            }
        }

        let provider = Arc::new(ErrorHeaderProvider) as Arc<dyn HeaderProvider>;
        let client_config = ClientConfig {
            header_provider: Some(provider),
            ..Default::default()
        };

        // Create connection - handler won't be called because header provider fails
        let conn = Connection::new_with_handler_and_config(
            move |_request| -> http::Response<&'static str> {
                panic!("Handler should not be called when header provider fails");
            },
            client_config,
        );

        // Request should fail due to header provider error
        let result = conn.table_names().execute().await;
        assert!(result.is_err());

        match result.unwrap_err() {
            crate::Error::Runtime { message } => {
                assert_eq!(message, "Failed to fetch auth token");
            }
            _ => panic!("Expected Runtime error from header provider"),
        }
    }

    #[tokio::test]
    async fn test_clone_table() {
        let conn = Connection::new_with_handler(|request| {
            assert_eq!(request.method(), &reqwest::Method::POST);
            assert_eq!(request.url().path(), "/v1/table/cloned_table/clone");
            assert_eq!(
                request.headers().get("Content-Type").unwrap(),
                JSON_CONTENT_TYPE
            );

            let body = request.body().unwrap().as_bytes().unwrap();
            let body: serde_json::Value = serde_json::from_slice(body).unwrap();
            assert_eq!(body["source_location"], "s3://bucket/source_table");
            assert_eq!(body["is_shallow"], true);

            http::Response::builder().status(200).body("").unwrap()
        });

        let table = conn
            .clone_table("cloned_table", "s3://bucket/source_table")
            .execute()
            .await
            .unwrap();
        assert_eq!(table.name(), "cloned_table");
    }

    #[tokio::test]
    async fn test_clone_table_with_version() {
        let conn = Connection::new_with_handler(|request| {
            assert_eq!(request.method(), &reqwest::Method::POST);
            assert_eq!(request.url().path(), "/v1/table/cloned_table/clone");

            let body = request.body().unwrap().as_bytes().unwrap();
            let body: serde_json::Value = serde_json::from_slice(body).unwrap();
            assert_eq!(body["source_location"], "s3://bucket/source_table");
            assert_eq!(body["source_version"], 42);
            assert_eq!(body["is_shallow"], true);

            http::Response::builder().status(200).body("").unwrap()
        });

        let table = conn
            .clone_table("cloned_table", "s3://bucket/source_table")
            .source_version(42)
            .execute()
            .await
            .unwrap();
        assert_eq!(table.name(), "cloned_table");
    }

    #[tokio::test]
    async fn test_clone_table_with_tag() {
        let conn = Connection::new_with_handler(|request| {
            assert_eq!(request.method(), &reqwest::Method::POST);
            assert_eq!(request.url().path(), "/v1/table/cloned_table/clone");

            let body = request.body().unwrap().as_bytes().unwrap();
            let body: serde_json::Value = serde_json::from_slice(body).unwrap();
            assert_eq!(body["source_location"], "s3://bucket/source_table");
            assert_eq!(body["source_tag"], "v1.0");
            assert_eq!(body["is_shallow"], true);

            http::Response::builder().status(200).body("").unwrap()
        });

        let table = conn
            .clone_table("cloned_table", "s3://bucket/source_table")
            .source_tag("v1.0")
            .execute()
            .await
            .unwrap();
        assert_eq!(table.name(), "cloned_table");
    }

    #[tokio::test]
    async fn test_clone_table_deep_clone() {
        let conn = Connection::new_with_handler(|request| {
            assert_eq!(request.method(), &reqwest::Method::POST);
            assert_eq!(request.url().path(), "/v1/table/cloned_table/clone");

            let body = request.body().unwrap().as_bytes().unwrap();
            let body: serde_json::Value = serde_json::from_slice(body).unwrap();
            assert_eq!(body["source_location"], "s3://bucket/source_table");
            assert_eq!(body["is_shallow"], false);

            http::Response::builder().status(200).body("").unwrap()
        });

        let table = conn
            .clone_table("cloned_table", "s3://bucket/source_table")
            .is_shallow(false)
            .execute()
            .await
            .unwrap();
        assert_eq!(table.name(), "cloned_table");
    }

    #[tokio::test]
    async fn test_clone_table_with_namespace() {
        let conn = Connection::new_with_handler(|request| {
            assert_eq!(request.method(), &reqwest::Method::POST);
            assert_eq!(request.url().path(), "/v1/table/ns1$ns2$cloned_table/clone");

            let body = request.body().unwrap().as_bytes().unwrap();
            let body: serde_json::Value = serde_json::from_slice(body).unwrap();
            assert_eq!(body["source_location"], "s3://bucket/source_table");
            assert_eq!(body["is_shallow"], true);

            http::Response::builder().status(200).body("").unwrap()
        });

        let table = conn
            .clone_table("cloned_table", "s3://bucket/source_table")
            .target_namespace(vec!["ns1".to_string(), "ns2".to_string()])
            .execute()
            .await
            .unwrap();
        assert_eq!(table.name(), "cloned_table");
    }

    #[tokio::test]
    async fn test_clone_table_error() {
        let conn = Connection::new_with_handler(|_| {
            http::Response::builder()
                .status(500)
                .body("Internal server error")
                .unwrap()
        });

        let result = conn
            .clone_table("cloned_table", "s3://bucket/source_table")
            .execute()
            .await;

        assert!(result.is_err());
        if let Err(crate::Error::Http { source, .. }) = result {
            assert!(source.to_string().contains("Failed to clone table"));
        } else {
            panic!("Expected HTTP error");
        }
    }

    #[tokio::test]
    async fn test_namespace_client() {
        let conn = Connection::new_with_handler(|_| {
            http::Response::builder()
                .status(200)
                .body(r#"{"tables": []}"#)
                .unwrap()
        });

        // Get the namespace client from the connection's internal database
        let namespace_client = conn.namespace_client().await;
        assert!(namespace_client.is_ok());
    }

    #[tokio::test]
    async fn test_namespace_client_with_tls_config() {
        use crate::remote::client::TlsConfig;

        let tls_config = TlsConfig {
            cert_file: Some("/path/to/cert.pem".to_string()),
            key_file: Some("/path/to/key.pem".to_string()),
            ssl_ca_cert: Some("/path/to/ca.pem".to_string()),
            assert_hostname: true,
        };

        let client_config = ClientConfig {
            tls_config: Some(tls_config),
            ..Default::default()
        };

        let conn = Connection::new_with_handler_and_config(
            |_| {
                http::Response::builder()
                    .status(200)
                    .body(r#"{"tables": []}"#)
                    .unwrap()
            },
            client_config,
        );

        // Get the namespace client - it should be created with the TLS config
        let namespace_client = conn.namespace_client().await;
        assert!(namespace_client.is_ok());
    }

    #[tokio::test]
    async fn test_namespace_client_with_headers() {
        let mut extra_headers = HashMap::new();
        extra_headers.insert("X-Custom-Header".to_string(), "custom-value".to_string());

        let client_config = ClientConfig {
            extra_headers,
            ..Default::default()
        };

        let conn = Connection::new_with_handler_and_config(
            |_| {
                http::Response::builder()
                    .status(200)
                    .body(r#"{"tables": []}"#)
                    .unwrap()
            },
            client_config,
        );

        // Get the namespace client - it should be created with the extra headers
        let namespace_client = conn.namespace_client().await;
        assert!(namespace_client.is_ok());
    }

    #[test]
    fn test_namespace_header_provider_context_maps_headers() {
        #[derive(Debug)]
        struct TestHeaderProvider;

        #[async_trait::async_trait]
        impl HeaderProvider for TestHeaderProvider {
            async fn get_headers(&self) -> crate::Result<HashMap<String, String>> {
                Ok(HashMap::from([(
                    "authorization".to_string(),
                    "Bearer token".to_string(),
                )]))
            }
        }

        let context_provider = NamespaceHeaderProviderContext {
            header_provider: Arc::new(TestHeaderProvider) as Arc<dyn HeaderProvider>,
        };

        let context =
            context_provider.provide_context(&OperationInfo::new("list_tables", "namespace"));

        assert_eq!(
            context.get("headers.authorization"),
            Some(&"Bearer token".to_string())
        );
    }

    #[tokio::test]
    async fn test_namespace_client_supports_dynamic_headers() {
        #[derive(Debug)]
        struct TestHeaderProvider;

        #[async_trait::async_trait]
        impl HeaderProvider for TestHeaderProvider {
            async fn get_headers(&self) -> crate::Result<HashMap<String, String>> {
                Ok(HashMap::from([(
                    "authorization".to_string(),
                    "Bearer token".to_string(),
                )]))
            }
        }

        let client_config = ClientConfig {
            header_provider: Some(Arc::new(TestHeaderProvider) as Arc<dyn HeaderProvider>),
            ..Default::default()
        };

        let conn = Connection::new_with_handler_and_config(
            |_| {
                http::Response::builder()
                    .status(200)
                    .body(r#"{"tables": []}"#)
                    .unwrap()
            },
            client_config,
        );

        let namespace_client = conn.namespace_client().await;
        assert!(namespace_client.is_ok());

        match conn.namespace_client_config().await {
            Err(Error::NotSupported { message })
                if message.contains("dynamic headers are configured") => {}
            Err(err) => panic!("expected NotSupported, got {err:?}"),
            Ok(_) => panic!("expected namespace_client_config to reject dynamic headers"),
        }
    }

    /// Integration tests using RestAdapter to run RemoteDatabase against a real namespace server
    mod rest_adapter_integration {
        use super::*;
        use lance_namespace::models::ListTablesRequest;
        use lance_namespace_impls::{DirectoryNamespaceBuilder, RestAdapter, RestAdapterConfig};
        use std::sync::Arc;
        use tempfile::TempDir;

        /// Test fixture that manages a REST server backed by DirectoryNamespace
        struct RestServerFixture {
            _temp_dir: TempDir,
            server_handle: lance_namespace_impls::RestAdapterHandle,
            server_url: String,
        }

        impl RestServerFixture {
            async fn new() -> Self {
                let temp_dir = TempDir::new().unwrap();
                let temp_path = temp_dir.path().to_str().unwrap().to_string();

                // Create DirectoryNamespace backend
                let backend = DirectoryNamespaceBuilder::new(&temp_path)
                    .build()
                    .await
                    .unwrap();
                let backend = Arc::new(backend);

                // Start REST server with port 0 (OS assigns available port)
                let config = RestAdapterConfig {
                    port: 0,
                    ..Default::default()
                };

                let server = RestAdapter::new(backend, config);
                let server_handle = server.start().await.unwrap();

                // Get the actual port assigned by OS
                let actual_port = server_handle.port();
                let server_url = format!("http://127.0.0.1:{}", actual_port);

                Self {
                    _temp_dir: temp_dir,
                    server_handle,
                    server_url,
                }
            }
        }

        impl Drop for RestServerFixture {
            fn drop(&mut self) {
                self.server_handle.shutdown();
            }
        }

        #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
        async fn test_remote_database_with_rest_adapter() {
            use lance_namespace::models::CreateNamespaceRequest;

            let fixture = RestServerFixture::new().await;

            // Connect to the REST server using lancedb Connection
            // Use db://dummy as URI and set actual server URL via host_override
            let conn = ConnectBuilder::new("db://dummy")
                .api_key("test-api-key")
                .region("us-east-1")
                .host_override(&fixture.server_url)
                .execute()
                .await
                .unwrap();

            // Create a child namespace first
            let namespace = vec!["test_ns".to_string()];
            conn.create_namespace(CreateNamespaceRequest {
                id: Some(namespace.clone()),
                ..Default::default()
            })
            .await
            .expect("Failed to create namespace");

            // Create a table in the child namespace
            let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]));
            let data = RecordBatch::try_new(
                schema.clone(),
                vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
            )
            .unwrap();
            let table = conn
                .create_table("test_table", data)
                .namespace(namespace.clone())
                .execute()
                .await;
            assert!(table.is_ok(), "Failed to create table: {:?}", table.err());

            // List tables in the child namespace
            let list_response = conn
                .list_tables(ListTablesRequest {
                    id: Some(namespace.clone()),
                    ..Default::default()
                })
                .await
                .expect("Failed to list tables");
            assert_eq!(list_response.tables, vec!["test_table"]);

            // Get namespace client and verify it can also list tables
            let namespace_client = conn.namespace_client().await.unwrap();
            let list_response = namespace_client
                .list_tables(ListTablesRequest {
                    id: Some(namespace.clone()),
                    ..Default::default()
                })
                .await
                .unwrap();
            assert_eq!(list_response.tables, vec!["test_table"]);

            // Open the table from the child namespace
            let opened_table = conn
                .open_table("test_table")
                .namespace(namespace.clone())
                .execute()
                .await;
            assert!(
                opened_table.is_ok(),
                "Failed to open table: {:?}",
                opened_table.err()
            );
            assert_eq!(opened_table.unwrap().name(), "test_table");
        }

        #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
        async fn test_remote_database_with_multiple_tables() {
            use lance_namespace::models::CreateNamespaceRequest;

            let fixture = RestServerFixture::new().await;

            // Connect to the REST server
            // Use db://dummy as URI and set actual server URL via host_override
            let conn = ConnectBuilder::new("db://dummy")
                .api_key("test-api-key")
                .region("us-east-1")
                .host_override(&fixture.server_url)
                .execute()
                .await
                .unwrap();

            // Create a child namespace first
            let namespace = vec!["multi_table_ns".to_string()];
            conn.create_namespace(CreateNamespaceRequest {
                id: Some(namespace.clone()),
                ..Default::default()
            })
            .await
            .expect("Failed to create namespace");

            // Create multiple tables in the child namespace
            let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));

            for i in 1..=3 {
                let data =
                    RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![i]))])
                        .unwrap();
                conn.create_table(format!("table{}", i), data)
                    .namespace(namespace.clone())
                    .execute()
                    .await
                    .unwrap_or_else(|e| panic!("Failed to create table{}: {:?}", i, e));
            }

            // List tables in the child namespace
            let list_response = conn
                .list_tables(ListTablesRequest {
                    id: Some(namespace.clone()),
                    ..Default::default()
                })
                .await
                .unwrap();
            assert_eq!(list_response.tables.len(), 3);
            assert!(list_response.tables.contains(&"table1".to_string()));
            assert!(list_response.tables.contains(&"table2".to_string()));
            assert!(list_response.tables.contains(&"table3".to_string()));
        }
    }

    #[tokio::test]
    async fn test_list_jobs_paginates() {
        let page = Arc::new(AtomicUsize::new(0));
        let conn = Connection::new_with_handler(move |request| {
            assert_eq!(request.method(), &reqwest::Method::POST);
            assert_eq!(request.url().path(), "/v1/jobs/list");
            let body: serde_json::Value =
                serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
            match page.fetch_add(1, Ordering::SeqCst) {
                0 => {
                    assert!(body.get("page_token").is_none());
                    http::Response::builder()
                        .status(200)
                        .body(
                            r#"{"jobs": [{"job_id": "job-1", "table": "t1", "job_type": "create_index", "state": "in_progress", "created_at_millis": 1000}], "page_token": "next"}"#,
                        )
                        .unwrap()
                }
                _ => {
                    assert_eq!(body["page_token"], "next");
                    http::Response::builder()
                        .status(200)
                        .body(
                            r#"{"jobs": [{"job_id": "job-2", "table": "t2", "job_type": "create_index", "state": "succeeded", "created_at_millis": 2000}, {"job_id": "job-3", "table": "t3", "job_type": "create_index", "state": "timed_out", "created_at_millis": 3000}]}"#,
                        )
                        .unwrap()
                }
            }
        });
        let jobs = conn.list_jobs().await.unwrap();
        assert_eq!(jobs.len(), 3);
        assert_eq!(jobs[0].job_id, "job-1");
        assert_eq!(jobs[0].table, "t1");
        assert_eq!(jobs[0].state, "running");
        assert_eq!(jobs[1].job_id, "job-2");
        assert_eq!(jobs[1].state, "finished");
        assert_eq!(jobs[1].created_at_millis, 2000);
        assert_eq!(jobs[2].job_id, "job-3");
        assert_eq!(jobs[2].state, "failed");
    }

    #[tokio::test]
    async fn test_open_job() {
        let conn = Connection::new_with_handler(|request| {
            assert_eq!(request.method(), &reqwest::Method::POST);
            assert_eq!(request.url().path(), "/v1/jobs/describe");
            let body: serde_json::Value =
                serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
            assert_eq!(body["job_id"], "job-1");
            http::Response::builder()
                .status(200)
                .body(
                    r#"{"job_id": "job-1", "job_type": "create_index", "job_state": "FAILED", "creation_ms": 1000, "spec": {"column": "vec"}, "failure": {"phase": "execute", "message": "worker died", "retryable": true}}"#,
                )
                .unwrap()
        });
        // Opening populates the handle, so the accessors answer without a
        // second round trip.
        let job = conn.open_job("job-1").await.unwrap();
        assert_eq!(job.id(), Some("job-1"));
        assert_eq!(job.job_type().as_deref(), Some("create_index"));
        assert_eq!(job.state().as_deref(), Some("failed"));
        assert_eq!(job.creation_ms(), Some(1000));
        assert_eq!(job.spec().unwrap()["column"], "vec");
        assert!(job.result().is_none());
        let failure = job.failure().unwrap();
        assert_eq!(failure.phase.as_deref(), Some("execute"));
        assert_eq!(failure.message.as_deref(), Some("worker died"));
        assert_eq!(failure.retryable, Some(true));
    }

    #[tokio::test]
    async fn test_open_job_reports_the_terminal_result() {
        let conn = Connection::new_with_handler(|_| {
            http::Response::builder()
                .status(200)
                .body(
                    r#"{"job_id": "job-1", "job_type": "refresh_column", "job_state": "DONE", "creation_ms": 1000, "result": {"rows_assigned": 1000000, "rows_failed": 0}}"#,
                )
                .unwrap()
        });
        let job = conn.open_job("job-1").await.unwrap();
        assert_eq!(job.state().as_deref(), Some("finished"));
        let result = job.result().unwrap();
        assert_eq!(result["rows_assigned"], 1_000_000);
        assert_eq!(result["rows_failed"], 0);
    }

    #[tokio::test]
    async fn test_open_job_missing_fails() {
        let conn = Connection::new_with_handler(|_| {
            http::Response::builder()
                .status(404)
                .body("no such job")
                .unwrap()
        });
        let err = conn.open_job("nope").await.unwrap_err();
        assert!(
            matches!(&err, Error::JobNotFound { job_id } if job_id == "nope"),
            "{err:?}"
        );
    }

    #[tokio::test]
    async fn test_job_events_scope_to_that_job() {
        let schema = Arc::new(Schema::new(vec![Field::new(
            "state",
            DataType::Utf8,
            false,
        )]));
        let batch = RecordBatch::try_new(
            schema.clone(),
            vec![Arc::new(arrow_array::StringArray::from(vec![
                "claim_complete",
            ]))],
        )
        .unwrap();
        let mut events = Vec::new();
        {
            let mut writer =
                arrow_ipc::writer::StreamWriter::try_new(&mut events, &schema).unwrap();
            writer.write(&batch).unwrap();
            writer.finish().unwrap();
        }
        let conn = Connection::new_with_handler(move |request| {
            let body: serde_json::Value =
                serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
            if request.url().path() == "/v1/jobs/describe" {
                return http::Response::builder()
                    .status(200)
                    .body(
                        r#"{"job_id": "job-1", "job_type": "refresh_column", "job_state": "IN_PROGRESS", "creation_ms": 1}"#
                            .as_bytes()
                            .to_vec(),
                    )
                    .unwrap();
            }
            assert_eq!(request.url().path(), "/v1/jobs/query_events");
            // The handle supplies job_id; the caller only narrows the query.
            assert_eq!(body["job_id"], "job-1");
            assert_eq!(body["limit"], 500);
            assert_eq!(body["filter"], "state = 'claim_complete'");
            http::Response::builder()
                .status(200)
                .body(events.clone())
                .unwrap()
        });
        let job = conn.open_job("job-1").await.unwrap();
        let batches = job
            .events(
                JobEventsRequest::default()
                    .limit(500)
                    .filter("state = 'claim_complete'"),
            )
            .await
            .unwrap();
        assert_eq!(batches.len(), 1);
        assert_eq!(batches[0].num_rows(), 1);
    }

    #[tokio::test]
    async fn test_job_events_keep_the_schema_when_nothing_matches() {
        let schema = Arc::new(Schema::new(vec![Field::new(
            "state",
            DataType::Utf8,
            false,
        )]));
        let mut events = Vec::new();
        {
            let mut writer =
                arrow_ipc::writer::StreamWriter::try_new(&mut events, &schema).unwrap();
            writer.finish().unwrap();
        }
        let conn = Connection::new_with_handler(move |request| {
            if request.url().path() == "/v1/jobs/describe" {
                return http::Response::builder()
                    .status(200)
                    .body(
                        r#"{"job_id": "job-1", "job_type": "refresh_column", "job_state": "IN_PROGRESS", "creation_ms": 1}"#
                            .as_bytes()
                            .to_vec(),
                    )
                    .unwrap();
            }
            let body: serde_json::Value =
                serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
            // Only the job id when the caller narrows nothing.
            assert_eq!(body, serde_json::json!({ "job_id": "job-1" }));
            http::Response::builder()
                .status(200)
                .body(events.clone())
                .unwrap()
        });
        let job = conn.open_job("job-1").await.unwrap();
        let batches = job.events(JobEventsRequest::default()).await.unwrap();
        assert_eq!(batches.len(), 1);
        assert_eq!(batches[0].num_rows(), 0);
        assert_eq!(batches[0].schema(), schema);
    }

    #[tokio::test]
    async fn test_create_function_async_sends_canonical_request_and_decodes_typed_job() {
        const REQUEST: &str = include_str!(
            "../../tests/fixtures/first_class_functions/v1/remote_function_registration_request.json"
        );
        const FUNCTION_JOB: &str =
            include_str!("../../tests/fixtures/first_class_functions/v1/remote_function_job.json");
        let expected: serde_json::Value = serde_json::from_str(REQUEST).unwrap();
        let conn = Connection::new_with_handler(move |request| match request.url().path() {
            "/v1/functions/create" => {
                assert_eq!(request.method(), &reqwest::Method::POST);
                let body: serde_json::Value =
                    serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
                assert_eq!(body, expected);
                http::Response::builder()
                    .status(202)
                    .body(r#"{"job_id":"job-function-1"}"#)
                    .unwrap()
            }
            "/v1/jobs/describe" => http::Response::builder()
                .status(200)
                .body(FUNCTION_JOB)
                .unwrap(),
            path => panic!("unexpected path: {path}"),
        });
        let request = crate::function::FunctionRegistrationRequest::from_json(REQUEST).unwrap();
        let job = conn.create_function_async(request).await.unwrap();
        assert_eq!(job.id(), Some("job-function-1"));
        let version = job.wait().await.unwrap();
        assert_eq!(version.name(), "embed");
        assert_eq!(version.version(), "fv_01K3EXACT");
    }

    #[tokio::test]
    async fn test_get_function_requires_and_sends_exact_version() {
        const VERSION: &str = include_str!(
            "../../tests/fixtures/first_class_functions/v1/remote_function_version.canonical.json"
        );
        let conn = Connection::new_with_handler(|request| {
            assert_eq!(request.method(), &reqwest::Method::POST);
            assert_eq!(request.url().path(), "/v1/functions/describe");
            let body: serde_json::Value =
                serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
            assert_eq!(
                body,
                serde_json::json!({"name": "embed", "version": "fv_01K3EXACT"})
            );
            http::Response::builder().status(200).body(VERSION).unwrap()
        });
        let version = conn.get_function("embed", "fv_01K3EXACT").await.unwrap();
        assert_eq!(version.name(), "embed");
        assert_eq!(version.version(), "fv_01K3EXACT");
    }

    #[tokio::test]
    async fn test_list_functions_requests_definitions_and_paginates() {
        const VERSION: &str = include_str!(
            "../../tests/fixtures/first_class_functions/v1/remote_function_version.canonical.json"
        );
        let version: serde_json::Value = serde_json::from_str(VERSION).unwrap();
        let page = Arc::new(AtomicUsize::new(0));
        let conn = Connection::new_with_handler(move |request| {
            assert_eq!(request.method(), &reqwest::Method::POST);
            assert_eq!(request.url().path(), "/v1/functions/list");
            let body: serde_json::Value =
                serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
            assert_eq!(body["include_definition"], true);
            match page.fetch_add(1, Ordering::SeqCst) {
                0 => {
                    assert!(body.get("page_token").is_none());
                    http::Response::builder()
                        .status(200)
                        .body(r#"{"functions": [], "page_token": "next"}"#.to_string())
                        .unwrap()
                }
                _ => {
                    assert_eq!(body["page_token"], "next");
                    http::Response::builder()
                        .status(200)
                        .body(
                            serde_json::json!({
                                "functions": [{
                                    "name": "embed",
                                    "version": "fv_01K3EXACT",
                                    "definition": version.clone(),
                                }],
                            })
                            .to_string(),
                        )
                        .unwrap()
                }
            }
        });
        let functions = conn.list_functions().await.unwrap();
        assert_eq!(functions.len(), 1);
        assert_eq!(functions[0].name(), "embed");
        assert_eq!(functions[0].version(), "fv_01K3EXACT");
    }

    #[tokio::test]
    async fn test_list_functions_stops_on_an_empty_page_token() {
        let requests = Arc::new(AtomicUsize::new(0));
        let seen = requests.clone();
        let conn = Connection::new_with_handler(move |request| {
            seen.fetch_add(1, Ordering::SeqCst);
            let body: serde_json::Value =
                serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
            assert!(body.get("page_token").is_none());
            http::Response::builder()
                .status(200)
                .body(r#"{"functions": [], "page_token": ""}"#)
                .unwrap()
        });

        let functions = conn.list_functions().await.unwrap();
        assert!(functions.is_empty());
        assert_eq!(requests.load(Ordering::SeqCst), 1);
    }

    #[tokio::test]
    async fn test_list_functions_rejects_a_page_token_cycle() {
        let page = Arc::new(AtomicUsize::new(0));
        let requests = page.clone();
        let conn = Connection::new_with_handler(move |request| {
            let body: serde_json::Value =
                serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
            let next_page_token = match page.fetch_add(1, Ordering::SeqCst) {
                0 => {
                    assert!(body.get("page_token").is_none());
                    "one"
                }
                1 => {
                    assert_eq!(body["page_token"], "one");
                    "two"
                }
                2 => {
                    assert_eq!(body["page_token"], "two");
                    "one"
                }
                page => panic!("unexpected page: {page}"),
            };
            http::Response::builder()
                .status(200)
                .body(
                    serde_json::json!({
                        "functions": [],
                        "page_token": next_page_token,
                    })
                    .to_string(),
                )
                .unwrap()
        });

        let error = conn.list_functions().await.unwrap_err();
        assert!(
            matches!(
                &error,
                Error::Http {
                    status_code: Some(http::StatusCode::OK),
                    ..
                }
            ),
            "got {error:?}"
        );
        assert_eq!(requests.load(Ordering::SeqCst), 3);
    }

    #[tokio::test]
    async fn test_drop_function_sends_exact_version_and_decodes_replay() {
        let conn = Connection::new_with_handler(|request| {
            assert_eq!(request.method(), &reqwest::Method::POST);
            assert_eq!(request.url().path(), "/v1/functions/drop");
            let body: serde_json::Value =
                serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
            assert_eq!(
                body,
                serde_json::json!({"name": "embed", "version": "fv_01K3EXACT"})
            );
            http::Response::builder()
                .status(200)
                .body(r#"{"dropped":false}"#)
                .unwrap()
        });
        assert!(!conn.drop_function("embed", "fv_01K3EXACT").await.unwrap());
    }

    #[tokio::test]
    async fn test_conn_job_waits_to_done() {
        let polls = Arc::new(AtomicUsize::new(0));
        let polls_ref = polls.clone();
        let conn = Connection::new_with_handler(move |request| {
            assert_eq!(request.url().path(), "/v1/jobs/describe");
            // Two in-progress answers: one for the load, one for the first
            // status poll.
            let state = if polls_ref.fetch_add(1, Ordering::SeqCst) < 2 {
                "IN_PROGRESS"
            } else {
                "DONE"
            };
            http::Response::builder()
                .status(200)
                .body(format!(
                    r#"{{"job_id": "job-1", "job_type": "create_function", "job_state": "{}", "creation_ms": 1, "result": {{"name": "embed", "version": "fv_1"}}}}"#,
                    state
                ))
                .unwrap()
        });
        let job = conn.open_job("job-1").await.unwrap();
        assert_eq!(job.id(), Some("job-1"));
        // Opening already answered the state; no extra call needed for it.
        assert_eq!(job.state().as_deref(), Some("running"));
        assert_eq!(job.status().await.unwrap(), "running");
        job.wait().await.unwrap();
        assert_eq!(job.status().await.unwrap(), "finished");
        assert!(polls.load(Ordering::SeqCst) >= 4);
    }
}
