//! DataFusion executor implementation for the Hamelin Executor trait.
//!
//! This module provides `DataFusionExecutor`, which executes Hamelin queries
//! directly in DataFusion without generating SQL strings.

use std::collections::{HashMap, HashSet};
use std::fmt::Debug;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;

use tokio::sync::Mutex as TokioMutex;

use crate::arrow::{arrow_batches_to_result_set, arrow_schema_to_columns};
use crate::config::{extract_s3_buckets, register_s3_bucket, resolve_files, resolve_mem};
use crate::iceberg::{
    build_iceberg_rest_catalog, resolve_iceberg_metadata, DefaultIcebergTableResolver,
    IcebergTableResolver,
};

use async_recursion::async_recursion;
use async_trait::async_trait;
use aws_config::BehaviorVersion;
use datafusion::arrow::array::{Array, UInt64Array};
use datafusion::catalog::view::ViewTable;
use datafusion::catalog::{
    CatalogProvider as DFCatalogProvider, MemoryCatalogProvider, MemorySchemaProvider,
    TableProvider,
};
use datafusion::common::TableReference;
use datafusion::datasource::file_format::parquet::ParquetSink;
use datafusion::datasource::physical_plan::FileScanConfig;
use datafusion::datasource::sink::DataSinkExec;
use datafusion::datasource::source::DataSourceExec;
use datafusion::execution::config::SessionConfig;
use datafusion::execution::session_state::SessionStateBuilder;
use datafusion::physical_plan::{displayable, ExecutionPlan};
use datafusion::prelude::SessionContext;

use hamelin_executor::config::{
    CatalogConfig, CatalogEntry, DatasetConfig, DatasetStore, ResolvedDataset,
};
use hamelin_executor::executor::{CompiledQuery, Executor, ExecutorError};
use hamelin_executor::results::ResultSet;
use hamelin_executor::ExecutorStats;
use hamelin_lib::catalog::{CatalogProvider, Column};
use hamelin_lib::err::{ContextualTranslationErrors, WithTranslationErrors};
use hamelin_lib::interner::Interner;
use hamelin_lib::parse_with_options;
use hamelin_lib::provider::EnvironmentProvider;
use hamelin_lib::tree::ast::expression::Expression;
use hamelin_lib::tree::ast::identifier::{CompoundIdentifier, Identifier, SimpleIdentifier};
use hamelin_lib::tree::ast::ParseWithErrors;
use hamelin_lib::tree::builder::ident as ident_builder;
use hamelin_lib::tree::options::{ParseOptions, TypeCheckOptions};
use hamelin_lib::type_check_with_options;
use hamelin_translation::IRStatement;

use crate::exec::{register_extensions, TrimStringsExec};
use crate::statement::{translate_statement, TranslatedStatement};

/// A table resolved from configuration, carrying both its schema columns and
/// its DataFusion `TableProvider` so callers can choose what they need.
struct ResolvedTable {
    identifier: Identifier,
    columns: Vec<Column>,
    provider: Arc<dyn TableProvider>,
}

/// User-facing executor configuration deserialized from the
/// `[executors.datafusion]` TOML section. The CLI converts the raw
/// string map here into a typed `SessionConfig` before handing it to
/// [`DataFusionExecutorConfig::with_session_config`]. Rust callers
/// (tests, embeddings) build a `SessionConfig` directly and never
/// touch this struct.
#[derive(serde::Deserialize, serde::Serialize, Default, Clone, Debug, PartialEq)]
#[serde(default)]
pub struct DataFusionExecutorConf {
    /// String-keyed DataFusion session options. Keys use DataFusion's
    /// dotted format (e.g. `"datafusion.execution.batch_size"`).
    ///
    /// ```toml
    /// [executors.datafusion.session]
    /// "datafusion.execution.batch_size" = "8192"
    /// "datafusion.optimizer.repartition_file_min_size" = "0"
    /// ```
    pub session: HashMap<String, String>,
}

impl DataFusionExecutorConf {
    /// Convert the string-keyed TOML overrides into a typed
    /// `SessionConfig`. Returns an error on the first bad key or
    /// value — DataFusion validates each key against its registered
    /// config schema and fails with a clear message.
    pub fn into_session_config(self) -> Result<SessionConfig, ExecutorError> {
        let mut sc = SessionConfig::new();
        for (key, value) in self.session {
            sc.options_mut().set(&key, &value).map_err(|e| {
                ExecutorError::ConfigurationError(
                    anyhow::anyhow!("invalid datafusion session option {key:?} = {value:?}: {e}")
                        .into(),
                )
            })?;
        }
        Ok(sc)
    }
}

/// Configuration for building a `DataFusionExecutor` from CLI config.
#[derive(Default, Clone)]
pub struct DataFusionExecutorConfig {
    pub datasets: HashMap<String, HashMap<String, HashMap<String, DatasetConfig>>>,
    pub catalogs: HashMap<String, CatalogEntry>,
    /// Optional pre-built DataFusion `SessionConfig`. Rust callers
    /// build one directly; the CLI bridge derives it from
    /// [`DataFusionExecutorConf::into_session_config`].
    pub session_config: Option<SessionConfig>,
}

impl DataFusionExecutorConfig {
    /// Register a dataset by AST identifier.
    ///
    /// Supports 2-part (`schema.table`) and 3-part (`catalog.schema.table`)
    /// identifiers. For 2-part identifiers, the catalog defaults to `"default"`.
    ///
    /// When the catalog is inferred as `"default"` (2-part identifiers),
    /// a default catalog entry is auto-inserted so that `DatasetStore` can
    /// resolve 2-part query references against it.
    pub fn with_dataset(mut self, identifier: Identifier, config: DatasetConfig) -> Self {
        let segments = identifier.segments();
        let (catalog, schema, table, needs_default_catalog) = match segments {
            [s, t] => ("default", s.as_str(), t.as_str(), true),
            [c, s, t] => (c.as_str(), s.as_str(), t.as_str(), false),
            _ => return self,
        };

        self.datasets
            .entry(catalog.to_string())
            .or_default()
            .entry(schema.to_string())
            .or_default()
            .insert(table.to_string(), config);

        if needs_default_catalog && !self.catalogs.contains_key(catalog) {
            self.catalogs.insert(
                catalog.to_string(),
                CatalogEntry {
                    default: true,
                    config: CatalogConfig::Reflection,
                },
            );
        }
        self
    }

    pub fn with_catalog(mut self, name: &str, entry: CatalogEntry) -> Self {
        self.catalogs.insert(name.to_string(), entry);
        self
    }

    /// Override the DataFusion `SessionConfig` used when building
    /// the executor's `SessionContext`. Rust callers (tests,
    /// embeddings) build a `SessionConfig` and pass it here; the CLI
    /// derives one from
    /// [`DataFusionExecutorConf::into_session_config`].
    pub fn with_session_config(mut self, session_config: SessionConfig) -> Self {
        self.session_config = Some(session_config);
        self
    }
}

/// Cumulative atomic counters backing `ExecutorStats`.
///
/// Shared across all queries on a `DataFusionExecutor` instance via `Arc`. The executor
/// folds per-query metrics into these after each execution; `stats()` returns a snapshot.
#[derive(Default, Debug)]
struct ExecutorStatsInner {
    rows_read: AtomicU64,
    bytes_read: AtomicU64,
    rows_written: AtomicU64,
    bytes_written: AtomicU64,
    trim_clamped: AtomicU64,
}

impl ExecutorStatsInner {
    fn snapshot(&self) -> ExecutorStats {
        ExecutorStats {
            rows_read: self.rows_read.load(Ordering::Relaxed),
            bytes_read: self.bytes_read.load(Ordering::Relaxed),
            rows_written: self.rows_written.load(Ordering::Relaxed),
            bytes_written: self.bytes_written.load(Ordering::Relaxed),
            trim_clamped: self.trim_clamped.load(Ordering::Relaxed),
        }
    }
}

/// DataFusion executor that uses caller-provided table schemas.
///
/// Tables are resolved on-demand from configured datasets and catalogs
/// (e.g. Iceberg REST) during `compile_query`.
pub struct DataFusionExecutor {
    ctx: SessionContext,
    /// Shared resolution layer for matching query refs to config.
    dataset_store: DatasetStore,
    /// Lock to serialize the check-resolve-register sequence for local tables.
    /// Prevents concurrent callers from redundantly resolving and
    /// double-registering the same table.
    local_table_lock: TokioMutex<()>,
    iceberg_table_resolver: Arc<dyn IcebergTableResolver>,
    /// Cumulative metric counters, folded from physical-plan `MetricsSet` after each query.
    stats_inner: Arc<ExecutorStatsInner>,
}

impl DataFusionExecutor {
    /// Build a `SessionContext` wiring up Hamelin's custom analyzer/optimizer
    /// rules on top of DataFusion's defaults.
    ///
    /// All Hamelin-specific wiring (extension planners, optimizer-rule
    /// overrides, the `JsonRewrite` analyzer) is applied via
    /// [`register_extensions`].
    fn build_session_context(session_config: Option<SessionConfig>) -> SessionContext {
        let mut builder = SessionStateBuilder::new().with_default_features();
        if let Some(sc) = session_config {
            builder = builder.with_config(sc);
        }
        let builder = register_extensions(builder);
        SessionContext::new_with_state(builder.build())
    }

    /// Create a new DataFusion executor with an empty catalog.
    pub fn new() -> Self {
        Self {
            ctx: Self::build_session_context(None),
            dataset_store: DatasetStore::new(HashMap::new(), HashMap::new()),
            local_table_lock: TokioMutex::new(()),
            iceberg_table_resolver: Arc::new(DefaultIcebergTableResolver::default()),
            stats_inner: Arc::new(ExecutorStatsInner::default()),
        }
    }

    /// Snapshot the cumulative executor stats.
    ///
    /// Counters are cumulative across the executor's lifetime — a snapshot reflects
    /// the running totals, not per-query values. See `ExecutorStats` field docs.
    pub fn stats(&self) -> ExecutorStats {
        self.stats_inner.snapshot()
    }

    /// Build a DataFusion executor from CLI configuration.
    ///
    /// Stores config only — datasets are resolved lazily on first use in
    /// `compile_query`. Iceberg catalogs are also built on-demand.
    pub fn from_config(config: DataFusionExecutorConfig) -> Self {
        Self {
            ctx: Self::build_session_context(config.session_config),
            dataset_store: DatasetStore::new(config.datasets, config.catalogs),
            local_table_lock: TokioMutex::new(()),
            iceberg_table_resolver: Arc::new(DefaultIcebergTableResolver::default()),
            stats_inner: Arc::new(ExecutorStatsInner::default()),
        }
    }

    pub fn with_iceberg_table_resolver(
        self,
        iceberg_table_resolver: Arc<dyn IcebergTableResolver>,
    ) -> Self {
        Self {
            ctx: self.ctx,
            dataset_store: self.dataset_store,
            local_table_lock: self.local_table_lock,
            iceberg_table_resolver,
            stats_inner: self.stats_inner,
        }
    }

    /// Get the underlying SessionContext for advanced configuration.
    pub fn session_context(&self) -> &SessionContext {
        &self.ctx
    }

    /// Register a table provider with the executor.
    ///
    /// Registers the provider in DataFusion's SessionContext for query execution.
    /// Schema information is derived from the provider's Arrow schema on demand.
    async fn register_provider(
        &self,
        table: Identifier,
        provider: Arc<dyn TableProvider>,
    ) -> Result<(), ExecutorError> {
        let segments = table.segments();
        let table_ref = match segments {
            [name] => TableReference::bare(name.as_str()),
            [schema, name] => TableReference::partial(schema.as_str(), name.as_str()),
            [catalog, schema, name] => {
                TableReference::full(catalog.as_str(), schema.as_str(), name.as_str())
            }
            _ => {
                return Err(ExecutorError::ConfigurationError(
                    anyhow::anyhow!(
                        "Invalid table identifier: expected 1-3 parts, got {}",
                        segments.len()
                    )
                    .into(),
                ));
            }
        };

        self.ensure_catalog_schema(&table_ref)?;

        self.ctx
            .register_table(table_ref, provider)
            .map_err(|e| ExecutorError::ConfigurationError(e.into()))?;

        Ok(())
    }

    fn ensure_catalog_schema(&self, table_ref: &TableReference) -> Result<(), ExecutorError> {
        let state = self.ctx.state();
        let options = state.config().options();
        let default_catalog = options.catalog.default_catalog.clone();
        let default_schema = options.catalog.default_schema.clone();
        let (catalog_name, schema_name) = match table_ref {
            TableReference::Bare { .. } => (default_catalog, default_schema),
            TableReference::Partial { schema, .. } => (default_catalog, schema.to_string()),
            TableReference::Full {
                catalog, schema, ..
            } => (catalog.to_string(), schema.to_string()),
        };

        let catalog_list = state.catalog_list().clone();
        let catalog = catalog_list.catalog(&catalog_name).unwrap_or_else(|| {
            let catalog: Arc<dyn DFCatalogProvider> = Arc::new(MemoryCatalogProvider::new());
            catalog_list.register_catalog(catalog_name.clone(), Arc::clone(&catalog));
            catalog
        });

        if catalog.schema(&schema_name).is_none() {
            catalog
                .register_schema(&schema_name, Arc::new(MemorySchemaProvider::new()))
                .map_err(|e| ExecutorError::ConfigurationError(e.into()))?;
        }

        Ok(())
    }

    /// Resolve a dataset configuration into columns and a TableProvider.
    ///
    /// For file-based and iceberg-metadata datasets, this also registers S3 object
    /// stores (with optional per-dataset region) so callers don't need to handle
    /// bucket setup separately.
    ///
    /// View datasets are compiled to DataFusion LogicalPlans and wrapped as
    /// ViewTables. The `views_in_progress` set tracks views currently being
    /// compiled to detect circular dependencies.
    #[async_recursion]
    async fn resolve_dataset(
        &self,
        identifier: &Identifier,
        canonical: &Identifier,
        config: &DatasetConfig,
        views_in_progress: &mut HashSet<Identifier>,
    ) -> Result<(Vec<Column>, Arc<dyn TableProvider>), ExecutorError> {
        match config {
            DatasetConfig::Files(files) => {
                let s3_buckets = extract_s3_buckets(&files.paths);
                if !s3_buckets.is_empty() {
                    let sdk_config = aws_config::load_defaults(BehaviorVersion::latest()).await;
                    for (scheme, bucket) in s3_buckets {
                        register_s3_bucket(
                            &self.ctx,
                            &sdk_config,
                            &scheme,
                            &bucket,
                            files.region.as_deref(),
                            files.role_arn.as_deref(),
                            &files.static_credentials,
                            files.endpoint_url.as_deref(),
                            files.force_path_style,
                            files.allow_http,
                        )
                        .await?;
                    }
                }
                resolve_files(files, &self.ctx).await
            }
            DatasetConfig::Mem(mem) => resolve_mem(&mem.columns),
            DatasetConfig::Iceberg(iceberg) => {
                resolve_iceberg_metadata(&iceberg.metadata_location, iceberg.region.as_deref())
                    .await
            }
            DatasetConfig::View(view) => {
                if !views_in_progress.insert(canonical.clone()) {
                    return Err(ExecutorError::ConfigurationError(
                        anyhow::anyhow!(
                            "Circular view dependency detected: '{}' references itself \
                             (directly or transitively)",
                            identifier
                        )
                        .into(),
                    ));
                }

                let view_name = format!("view '{}'", identifier);
                let translated = self
                    .compile_query_helper(&view.query, None, Some(&view_name), views_in_progress)
                    .await?;

                views_in_progress.remove(canonical);

                match translated {
                    TranslatedStatement::Query {
                        plan,
                        output_schema,
                    } => {
                        let provider: Arc<dyn TableProvider> = Arc::new(ViewTable::new(plan, None));
                        Ok((output_schema.clone(), provider))
                    }
                    TranslatedStatement::Dml { .. } => Err(ExecutorError::ConfigurationError(
                        anyhow::anyhow!(
                            "View '{}' query must be a SELECT, not a DML statement",
                            identifier
                        )
                        .into(),
                    )),
                }
            }
        }
    }

    /// Resolve a single local dataset on demand, registering it in the SessionContext
    /// if not already present.
    async fn resolve_local_table(
        &self,
        identifier: &Identifier,
        canonical: &Identifier,
        dataset_config: &DatasetConfig,
        views_in_progress: &mut HashSet<Identifier>,
    ) -> Result<Arc<dyn TableProvider>, ExecutorError> {
        let segments = identifier.segments();
        let table_ref = match segments {
            [name] => TableReference::bare(name.as_str()),
            [schema, name] => TableReference::partial(schema.as_str(), name.as_str()),
            [catalog, schema, name] => {
                TableReference::full(catalog.as_str(), schema.as_str(), name.as_str())
            }
            _ => {
                return Err(ExecutorError::ConfigurationError(
                    anyhow::anyhow!(
                        "Invalid table identifier: expected 1-3 parts, got {}",
                        segments.len()
                    )
                    .into(),
                ));
            }
        };

        {
            let _guard = self.local_table_lock.lock().await;

            if let Ok(provider) = self.ctx.table_provider(table_ref.clone()).await {
                return Ok(provider);
            }
        }

        // Do the expensive/recursive dataset resolution without holding the mutex.
        // View-backed datasets recurse through compile_query_helper -> resolve_tables
        // -> resolve_local_table, so holding the lock here deadlocks chained views.
        let (_columns, provider) = self
            .resolve_dataset(identifier, canonical, dataset_config, views_in_progress)
            .await?;

        let _guard = self.local_table_lock.lock().await;

        if let Ok(provider) = self.ctx.table_provider(table_ref).await {
            return Ok(provider);
        }

        self.register_provider(identifier.clone(), provider.clone())
            .await?;

        Ok(provider)
    }

    /// Classify and resolve a set of AST identifiers against the dataset store,
    /// handling both local file-based datasets and Iceberg REST catalog tables.
    ///
    /// Returns a `ResolvedTable` per identifier that could be resolved. Identifiers
    /// that don't match any configured dataset are silently skipped (they'll fail
    /// later at type-check with a proper error).
    async fn resolve_tables(
        &self,
        identifiers: &[Identifier],
        views_in_progress: &mut HashSet<Identifier>,
    ) -> Result<Vec<ResolvedTable>, ExecutorError> {
        let mut local_refs: Vec<(Identifier, Identifier, DatasetConfig)> = vec![];
        let mut iceberg_refs: HashMap<String, Vec<(Identifier, CatalogConfig, String, String)>> =
            HashMap::new();

        let mut already_registered: Vec<Identifier> = vec![];

        for identifier in identifiers {
            let segments = identifier.segments();
            let seg_strs: Vec<&str> = segments.iter().map(|s| s.as_str()).collect();
            let resolved = match self.dataset_store.resolve(&seg_strs) {
                Ok(r) => r,
                Err(_) => {
                    already_registered.push(identifier.clone());
                    continue;
                }
            };

            match resolved {
                ResolvedDataset::Local {
                    catalog: cat,
                    schema: sch,
                    table: tbl,
                    config,
                    ..
                } => {
                    let canonical: Identifier = CompoundIdentifier::new(
                        SimpleIdentifier::new(&cat),
                        SimpleIdentifier::new(&sch),
                        vec![SimpleIdentifier::new(&tbl)],
                    )
                    .into();

                    local_refs.push((identifier.clone(), canonical, config));
                }
                ResolvedDataset::Catalog {
                    catalog_name,
                    catalog_config,
                    namespace,
                    table,
                } => {
                    iceberg_refs.entry(catalog_name).or_default().push((
                        identifier.clone(),
                        catalog_config,
                        namespace,
                        table,
                    ));
                }
            }
        }

        let mut results = Vec::new();

        // Resolve local tables
        for (identifier, canonical, dataset_config) in &local_refs {
            let provider = self
                .resolve_local_table(identifier, canonical, dataset_config, views_in_progress)
                .await?;
            let columns = arrow_schema_to_columns(&provider.schema());
            results.push(ResolvedTable {
                identifier: identifier.clone(),
                columns,
                provider,
            });
        }

        // Resolve Iceberg tables — one client per catalog, tables resolved in parallel
        for (catalog_name, tables) in &iceberg_refs {
            let (uri, warehouse, extra_properties) = match &tables[0].1 {
                CatalogConfig::IcebergRest {
                    uri,
                    warehouse,
                    extra_properties,
                } => (uri.clone(), warehouse.clone(), extra_properties.clone()),
                CatalogConfig::Reflection => {
                    return Err(ExecutorError::ConfigurationError(
                        anyhow::anyhow!(
                            "Catalog '{}' is configured as 'reflection', which is not supported by the DataFusion executor. \
                             Reflection catalogs are designed to work only with database-style executors, like Trino.",
                            catalog_name
                        )
                        .into(),
                    ));
                }
            };

            let iceberg_cat = build_iceberg_rest_catalog(
                &catalog_name,
                &uri,
                warehouse.as_deref(),
                extra_properties,
            )
            .await?;

            let mut join_set = tokio::task::JoinSet::new();
            for (identifier, _, namespace, table) in tables {
                let cat = iceberg_cat.clone();
                let ns = namespace.clone();
                let tbl = table.clone();
                let id = identifier.clone();
                let resolver = self.iceberg_table_resolver.clone();
                join_set.spawn(async move {
                    let provider = resolver.resolve_table(cat, &ns, &tbl).await?;

                    let schema = provider.schema();
                    let columns = arrow_schema_to_columns(&schema);
                    Ok::<_, ExecutorError>((id, columns, provider))
                });
            }
            while let Some(result) = join_set.join_next().await {
                let (identifier, columns, provider) = result
                    .map_err(|e| ExecutorError::ConfigurationError(anyhow::anyhow!(e).into()))??;
                self.register_iceberg_provider(identifier.clone(), provider.clone())
                    .await?;
                results.push(ResolvedTable {
                    identifier,
                    columns,
                    provider,
                });
            }
        }

        // Tables already registered on the SessionContext aren't in the
        // dataset_store config — resolve them here.
        for identifier in &already_registered {
            let segments = identifier.segments();
            let table_ref = match segments {
                [name] => TableReference::bare(name.as_str()),
                [schema, name] => TableReference::partial(schema.as_str(), name.as_str()),
                [catalog, schema, name] => {
                    TableReference::full(catalog.as_str(), schema.as_str(), name.as_str())
                }
                _ => continue,
            };
            if let Ok(provider) = self.ctx.table_provider(table_ref).await {
                let columns = arrow_schema_to_columns(&provider.schema());
                results.push(ResolvedTable {
                    identifier: identifier.clone(),
                    columns,
                    provider,
                });
            }
        }

        Ok(results)
    }

    /// Parse, resolve, type-check, and lower a Hamelin query to IR.
    ///
    /// Returns the IR statement and the original source string. Tables are
    /// resolved and registered on `self.ctx` so the subsequent translation
    /// step can look them up via `SessionContext::table_provider`.
    ///
    /// When `source_name` is provided (e.g. a view name), it appears in
    /// error messages instead of `<unknown>`.
    async fn compile_to_ir(
        &self,
        hamelin: &str,
        time_range: Option<&str>,
        source_name: Option<&str>,
        views_in_progress: &mut HashSet<Identifier>,
    ) -> Result<(IRStatement, String), ExecutorError> {
        let source = hamelin.to_string();
        let sn = source_name.map(|s| s.to_string());

        // Shared interner for parse and typecheck phases
        let interner = Arc::new(Interner::new());

        // Parse the query
        let WithTranslationErrors {
            output: ast,
            errors: parse_errors,
        } = parse_with_options(
            &source,
            ParseOptions::builder().interner(interner.clone()).build(),
        );
        if !parse_errors.is_empty() {
            return Err(ContextualTranslationErrors::new_with_source_name(
                source,
                sn,
                parse_errors,
            )
            .into());
        }

        // Resolve all table references into table sources and a per-query catalog
        let direct_identifiers = ast.datasets().map_err(|e| {
            ExecutorError::from(ContextualTranslationErrors::new_with_source_name(
                source.clone(),
                sn.clone(),
                e.as_ref().clone().single(),
            ))
        })?;

        let resolved = self
            .resolve_tables(&direct_identifiers, views_in_progress)
            .await?;
        let catalog = CatalogProvider::default();
        for rt in &resolved {
            register_resolved(&rt.provider, &rt.identifier, &catalog)?;
        }

        let catalog = Arc::new(catalog);

        // Type-check with the per-query catalog (optionally inject time range)
        let parsed_time_range = time_range
            .map(|tr| {
                Expression::parse_result(tr)
                    .map(|e| Arc::new(e))
                    .map_err(|errors| ExecutorError::from(errors.contextualize(tr)))
            })
            .transpose()?;

        let WithTranslationErrors {
            output: typed,
            errors: type_errors,
        } = type_check_with_options(
            ast,
            TypeCheckOptions::builder()
                .provider(catalog.clone() as Arc<dyn EnvironmentProvider>)
                .interner(interner)
                .maybe_time_range(parsed_time_range)
                .build(),
        );

        if !type_errors.is_empty() {
            return Err(
                ContextualTranslationErrors::new_with_source_name(source, sn, type_errors).into(),
            );
        }

        // Lower to IR - must pass the catalog provider for re-typechecking during normalization
        // Enable transform lowering since DataFusion doesn't support lambdas
        let ir = hamelin_translation::normalize_with()
            .with_provider(catalog)
            .with_lower_filter()
            .with_lower_transform()
            .lower(typed.into())
            .map_err(|e| ExecutorError::from(e.with_source_name(sn.clone(), true)))?;

        Ok((ir, source))
    }

    pub async fn compile_query(
        &self,
        hamelin: &str,
        time_range: Option<&str>,
    ) -> Result<TranslatedStatement, ExecutorError> {
        self.compile_query_helper(hamelin, time_range, None, &mut HashSet::new())
            .await
    }

    async fn compile_query_helper(
        &self,
        hamelin: &str,
        time_range: Option<&str>,
        source_name: Option<&str>,
        views_in_progress: &mut HashSet<Identifier>,
    ) -> Result<TranslatedStatement, ExecutorError> {
        let (ir, source) = self
            .compile_to_ir(hamelin, time_range, source_name, views_in_progress)
            .await?;

        // Translate to DataFusion LogicalPlan
        let translated = translate_statement(&ir, &self.ctx).await.map_err(|e| {
            ContextualTranslationErrors::new_with_source_name(
                source,
                source_name.map(|s| s.to_string()),
                e.as_ref().clone().single(),
            )
        })?;

        Ok(translated)
    }

    /// Register an Iceberg table provider on the SessionContext (for execution).
    ///
    /// Always deregisters any existing provider first so re-resolution picks up
    /// schema/snapshot changes between queries.
    async fn register_iceberg_provider(
        &self,
        table: Identifier,
        provider: Arc<dyn TableProvider>,
    ) -> Result<(), ExecutorError> {
        let segments = table.segments();
        let table_ref = match segments {
            [name] => TableReference::bare(name.as_str()),
            [schema, name] => TableReference::partial(schema.as_str(), name.as_str()),
            [catalog, schema, name] => {
                TableReference::full(catalog.as_str(), schema.as_str(), name.as_str())
            }
            _ => {
                return Err(ExecutorError::ConfigurationError(
                    anyhow::anyhow!(
                        "Invalid table identifier: expected 1-3 parts, got {}",
                        segments.len()
                    )
                    .into(),
                ));
            }
        };

        if self.ctx.table_provider(table_ref.clone()).await.is_ok() {
            self.ctx
                .deregister_table(table_ref.clone())
                .map_err(|e| ExecutorError::ConfigurationError(e.into()))?;
        }

        self.ensure_catalog_schema(&table_ref)?;
        self.ctx
            .register_table(table_ref, provider)
            .map_err(|e| ExecutorError::ConfigurationError(e.into()))?;

        Ok(())
    }

    /// Compile a Hamelin query and return its physical plan.
    ///
    /// This is the same plan-construction path the trait's `explain`
    /// uses; `explain` is now a thin stringification on top of this
    /// method. Tests use this to inspect the post-`JsonRewrite` plan
    /// structurally without going through the string-formatted output.
    pub async fn create_physical_plan(
        &self,
        hamelin: &str,
        time_range: Option<&str>,
    ) -> Result<Arc<dyn ExecutionPlan>, ExecutorError> {
        let translated = self.compile_query(hamelin, time_range).await?;

        let logical_plan = match translated {
            TranslatedStatement::Query { plan, .. } => plan,
            dml @ TranslatedStatement::Dml { .. } => dml.into_dml_plan(&self.ctx).await?,
        };

        self.ctx
            .state()
            .create_physical_plan(&logical_plan)
            .await
            .map_err(|e| ExecutorError::QueryError(e.into()))
    }
}

impl Default for DataFusionExecutor {
    fn default() -> Self {
        Self::new()
    }
}

/// Walk the physical plan tree and fold its `MetricsSet` values into `stats`.
///
/// Per-node policy (named to match the metric names DataFusion's stock operators emit):
/// - `DataSourceExec` → contributes `output_rows` → `rows_read`, and `bytes_read` from
///   either a `bytes_scanned` metric (parquet emits one, accounting for projection
///   pushdown) or — when no metric is present — the byte ranges declared in the
///   `FileScanConfig` (csv/json/lines, which read full files end-to-end).
/// - `DataSinkExec(ParquetSink)` → contributes `bytes_written` (read post-hoc from
///   `ParquetSink::written`'s row-group metadata; stock parquet sink doesn't expose
///   metrics). `rows_written` is captured separately from `execute_dml`'s output batch.
/// - `TrimStringsExec` → contributes `trim_clamped`
///
/// Other nodes' metrics are intentionally ignored; summing `output_rows` across every
/// operator would double-count rows passing through the plan.
fn fold_metrics_into_stats(plan: &dyn ExecutionPlan, stats: &Arc<ExecutorStatsInner>) {
    let any = plan.as_any();
    if let Some(sink_exec) = any.downcast_ref::<DataSinkExec>() {
        if let Some(parquet_sink) = sink_exec.sink().as_any().downcast_ref::<ParquetSink>() {
            // `ParquetSink::written()` returns the per-file `ParquetMetaData` after
            // `write_all` completes. `RowGroupMetaData::compressed_size()` sums column
            // chunk byte sizes; this excludes the parquet footer and magic (typically
            // <1% of file), but matches parquet's `bytes_scanned` semantics on the
            // read side: bytes that flowed through compressed I/O against storage.
            let total: i64 = parquet_sink
                .written()
                .values()
                .flat_map(|meta| meta.row_groups().iter())
                .map(|rg| rg.compressed_size())
                .sum();
            stats
                .bytes_written
                .fetch_add(total as u64, Ordering::Relaxed);
        }
    } else if let Some(source_exec) = any.downcast_ref::<DataSourceExec>() {
        if let Some(metrics) = plan.metrics() {
            // `output_rows` is the typed `MetricValue::OutputRows` variant, accessed via
            // the `output_rows()` convenience method — `sum_by_name("output_rows")` does
            // NOT match it (it filters out the typed variants).
            if let Some(n) = metrics.output_rows() {
                stats.rows_read.fetch_add(n as u64, Ordering::Relaxed);
            }
            // Prefer a metric-emitted `bytes_scanned` (parquet, accounting for
            // projection/predicate pushdown). When no metric is present, fall back to
            // the byte ranges in `FileScanConfig` — for our streaming text formats
            // (csv, json, lines) the opener pulls the full declared range from
            // storage, so this is exact.
            if let Some(v) = metrics.sum_by_name("bytes_scanned") {
                stats
                    .bytes_read
                    .fetch_add(v.as_usize() as u64, Ordering::Relaxed);
            } else if let Some(fsc) = source_exec
                .data_source()
                .as_any()
                .downcast_ref::<FileScanConfig>()
            {
                let total: u64 = fsc
                    .file_groups
                    .iter()
                    .flat_map(|g| g.iter())
                    .map(|pf| match &pf.range {
                        Some(r) => (r.end - r.start) as u64,
                        None => pf.object_meta.size,
                    })
                    .sum();
                stats.bytes_read.fetch_add(total, Ordering::Relaxed);
            }
        }
    } else if any.is::<TrimStringsExec>() {
        if let Some(metrics) = plan.metrics() {
            // Our own Count metric registered via MetricBuilder::counter — generic Count
            // keyed by name.
            if let Some(v) = metrics.sum_by_name("trim_clamped") {
                stats
                    .trim_clamped
                    .fetch_add(v.as_usize() as u64, Ordering::Relaxed);
            }
        }
    }

    for child in plan.children() {
        fold_metrics_into_stats(child.as_ref(), stats);
    }
}

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

#[async_trait]
impl Executor for DataFusionExecutor {
    async fn translate(
        &self,
        hamelin: &str,
        time_range: Option<&str>,
    ) -> Result<CompiledQuery, ExecutorError> {
        let translated = self.compile_query(hamelin, time_range).await?;

        match translated {
            TranslatedStatement::Query {
                plan,
                output_schema,
            } => Ok(CompiledQuery {
                display: format!("{}", plan.display_indent()),
                columns: output_schema.clone(),
            }),
            TranslatedStatement::Dml {
                source_plan,
                target_table,
                distinct_by,
            } => {
                let mut display = format!("INSERT INTO {}", target_table);
                if !distinct_by.is_empty() {
                    let cols: Vec<String> = distinct_by.iter().map(|id| id.to_string()).collect();
                    display.push_str(&format!(" DISTINCT {}", cols.join(", ")));
                }
                display.push_str(&format!("\n{}", source_plan.display_indent()));
                Ok(CompiledQuery {
                    display,
                    columns: vec![],
                })
            }
        }
    }

    async fn ir(
        &self,
        hamelin: &str,
        time_range: Option<&str>,
    ) -> Result<CompiledQuery, ExecutorError> {
        let (ir, _) = self
            .compile_to_ir(hamelin, time_range, None, &mut HashSet::new())
            .await?;
        let columns = ir.pipeline.output_schema.as_struct().into();
        Ok(CompiledQuery {
            display: ir.to_string(),
            columns,
        })
    }

    async fn explain(
        &self,
        hamelin: &str,
        time_range: Option<&str>,
    ) -> Result<String, ExecutorError> {
        let physical_plan = self.create_physical_plan(hamelin, time_range).await?;
        let display = displayable(physical_plan.as_ref()).indent(true).to_string();
        Ok(display)
    }

    async fn execute_query(
        &self,
        hamelin: &str,
        time_range: Option<&str>,
    ) -> Result<ResultSet, ExecutorError> {
        let translated = self.compile_query(hamelin, time_range).await?;

        let (plan, output_schema) = translated.query().ok_or_else(|| {
            ExecutorError::QueryError(anyhow::anyhow!("expected query statement, got DML").into())
        })?;

        // Build the physical plan ourselves (instead of going through `DataFrame`) so we
        // can walk it post-execution and read its `MetricsSet`. `DataFrame::collect` drops
        // the `ExecutionPlan` before returning, taking the metrics with it.
        let physical_plan = self
            .ctx
            .state()
            .create_physical_plan(&plan)
            .await
            .map_err(|e| ExecutorError::QueryError(e.into()))?;

        let batches =
            datafusion::physical_plan::collect(Arc::clone(&physical_plan), self.ctx.task_ctx())
                .await
                .map_err(|e| ExecutorError::QueryError(e.into()))?;

        fold_metrics_into_stats(physical_plan.as_ref(), &self.stats_inner);

        arrow_batches_to_result_set(&batches, output_schema)
    }

    async fn execute_dml(
        &self,
        hamelin: &str,
        time_range: Option<&str>,
    ) -> Result<usize, ExecutorError> {
        let translated = self.compile_query(hamelin, time_range).await?;
        let insert_plan = translated.into_dml_plan(&self.ctx).await?;

        let physical_plan = self
            .ctx
            .state()
            .create_physical_plan(&insert_plan)
            .await
            .map_err(|e| ExecutorError::QueryError(e.into()))?;

        let batches =
            datafusion::physical_plan::collect(Arc::clone(&physical_plan), self.ctx.task_ctx())
                .await
                .map_err(|e| ExecutorError::QueryError(e.into()))?;

        fold_metrics_into_stats(physical_plan.as_ref(), &self.stats_inner);

        // The `DataSinkExec` output is a one-row UInt64 RecordBatch holding the row count
        // returned by `DataSink::write_all`. That number is also our `rows_written` stat —
        // stock parquet sinks don't emit a `rows_written` counter for the metrics fold,
        // so this is where we capture it.
        let count = batches
            .first()
            .filter(|b| b.num_columns() > 0 && b.num_rows() > 0)
            .and_then(|b| b.column(0).as_any().downcast_ref::<UInt64Array>())
            .map(|a| a.value(0))
            .unwrap_or(0);
        self.stats_inner
            .rows_written
            .fetch_add(count, Ordering::Relaxed);
        Ok(count as usize)
    }

    async fn create_dataset(
        &self,
        _dataset: Identifier,
        _columns: Vec<Column>,
    ) -> Result<(), ExecutorError> {
        Err(ExecutorError::ConfigurationError(
            anyhow::anyhow!(
                "DataFusion local tables must be configured via DatasetConfig::Mem, not create_dataset"
            )
            .into(),
        ))
    }

    async fn resolve_datasets(
        &self,
        datasets: &[Identifier],
    ) -> Result<Vec<Vec<Column>>, ExecutorError> {
        let resolved = self.resolve_tables(datasets, &mut HashSet::new()).await?;

        // Build a lookup map from the resolved tables
        let resolved_map: HashMap<&Identifier, &Vec<Column>> = resolved
            .iter()
            .map(|rt| (&rt.identifier, &rt.columns))
            .collect();

        // Return columns in input order, falling back to SessionContext lookup
        let mut results = Vec::with_capacity(datasets.len());
        for id in datasets {
            if let Some(cols) = resolved_map.get(id) {
                results.push((*cols).clone());
                continue;
            }
            // Try the SessionContext directly (e.g. tables registered internally)
            let segments = id.segments();
            let table_ref = match segments {
                [name] => TableReference::bare(name.as_str()),
                [schema, name] => TableReference::partial(schema.as_str(), name.as_str()),
                [catalog, schema, name] => {
                    TableReference::full(catalog.as_str(), schema.as_str(), name.as_str())
                }
                _ => {
                    return Err(ExecutorError::QueryError(
                        anyhow::anyhow!("Dataset '{}' not found", id).into(),
                    ));
                }
            };
            let provider = self.ctx.table_provider(table_ref).await.map_err(|_| {
                ExecutorError::QueryError(anyhow::anyhow!("Dataset '{}' not found", id).into())
            })?;
            results.push(arrow_schema_to_columns(&provider.schema()));
        }
        Ok(results)
    }

    async fn reflect_catalog(&self) -> Result<hamelin_lib::catalog::Catalog, ExecutorError> {
        let mut tables = HashMap::new();

        match self.dataset_store.default_catalog_config() {
            Some((
                catalog_name,
                CatalogConfig::IcebergRest {
                    uri,
                    warehouse,
                    extra_properties,
                },
            )) => {
                // Connect to the Iceberg REST catalog and enumerate all tables.
                let iceberg_cat = build_iceberg_rest_catalog(
                    catalog_name,
                    uri,
                    warehouse.as_deref(),
                    extra_properties.clone(),
                )
                .await?;

                let namespaces = iceberg_cat
                    .list_namespaces(None)
                    .await
                    .map_err(|e| ExecutorError::ServerError(e.into()))?;

                for ns in &namespaces {
                    let table_idents = iceberg_cat
                        .list_tables(ns)
                        .await
                        .map_err(|e| ExecutorError::ServerError(e.into()))?;

                    for table_ident in &table_idents {
                        let iceberg_table = iceberg_cat
                            .load_table(table_ident)
                            .await
                            .map_err(|e| ExecutorError::ServerError(e.into()))?;

                        let iceberg_schema = iceberg_table.metadata().current_schema();
                        let arrow_schema = iceberg::arrow::schema_to_arrow_schema(iceberg_schema)
                            .map_err(|e| ExecutorError::ServerError(e.into()))?;
                        let columns = arrow_schema_to_columns(&arrow_schema);
                        let key = match ns.as_ref().split_first() {
                            Some((first, rest)) => {
                                let ns_ident = rest
                                    .iter()
                                    .fold(ident_builder(first), |acc, part| acc.dot(part));
                                ns_ident.dot(table_ident.name()).build()
                            }
                            _ => ident_builder(table_ident.name()).build(),
                        };
                        tables.insert(key, columns);
                    }
                }
            }
            Some((_, CatalogConfig::Reflection)) => {
                return Err(ExecutorError::ConfigurationError(
                    anyhow::anyhow!(
                        "Default catalog is configured as 'reflection', which is not supported \
                         by the DataFusion executor. Reflection catalogs work only with \
                         database-style executors like Trino."
                    )
                    .into(),
                ));
            }
            None => {
                // No default catalog configured — nothing to reflect.
            }
        }

        Ok(hamelin_lib::catalog::Catalog(tables))
    }

    fn stats(&self) -> Option<ExecutorStats> {
        Some(self.stats_inner.snapshot())
    }
}

/// Derive Hamelin columns from a table provider's Arrow schema and populate
/// the per-query Hamelin catalog (used for type-checking).
fn register_resolved(
    provider: &Arc<dyn TableProvider>,
    identifier: &Identifier,
    catalog: &CatalogProvider,
) -> Result<(), ExecutorError> {
    let columns = arrow_schema_to_columns(&provider.schema());

    let cols_map = columns
        .iter()
        .map(|c| {
            let id = c.name.clone().into();
            let typ = c.typ.clone().try_into().map_err(|e| {
                ExecutorError::ConfigurationError(
                    anyhow::anyhow!("Invalid column type: {}", e).into(),
                )
            })?;
            Ok((id, typ))
        })
        .collect::<Result<ordermap::OrderMap<_, _>, ExecutorError>>()?;

    catalog.set(identifier.clone(), cols_map);

    Ok(())
}
