//! Custom DataFusion plan extensions for Hamelin commands that don't lower to standard SQL.
//!
//! Currently houses `TRIMSTRINGS` (byte-cap of every string column, recursing through
//! structs and arrays). The logical extension node `TrimStringsLogicalNode` is a structural
//! marker in the logical plan; its physical counterpart `TrimStringsExec` runs the kernel
//! and emits the `trim_clamped` metric.
//!
//! ## Embedding Hamelin in your own SessionState
//!
//! Downstream callers (e.g. phantom) that build their own `SessionContext` need the
//! Hamelin extension planners registered so they can lower Hamelin's logical extension
//! nodes (`TrimStringsLogicalNode`, etc.) to physical operators at planning time.
//!
//! Use [`register_extensions`] to apply Hamelin's planner to a `SessionStateBuilder`:
//!
//! ```ignore
//! use datafusion::execution::SessionStateBuilder;
//! use datafusion::prelude::SessionContext;
//!
//! let state = hamelin_datafusion::register_extensions(
//!     SessionStateBuilder::new().with_default_features(),
//! )
//! .build();
//! let ctx = SessionContext::new_with_state(state);
//! ```
//!
//! For UDF/UDAF round-tripping over the wire, see `hamelin_datafusion::HamelinPhysicalCodec`.

mod trim_strings;

use std::fmt::Debug;
use std::sync::Arc;

use async_trait::async_trait;
use datafusion::common::Result as DfResult;
use datafusion::execution::context::QueryPlanner;
use datafusion::execution::{SessionState, SessionStateBuilder};
use datafusion::logical_expr::LogicalPlan;
use datafusion::physical_plan::ExecutionPlan;
use datafusion::physical_planner::{DefaultPhysicalPlanner, ExtensionPlanner, PhysicalPlanner};

pub use trim_strings::{TrimStringsExec, TrimStringsExtensionPlanner, TrimStringsLogicalNode};

/// Register Hamelin's custom DataFusion extensions on a `SessionStateBuilder`.
///
/// Wires both:
/// - The query planner that lowers `LogicalPlan::Extension` nodes (e.g.
///   `TrimStringsLogicalNode` → `TrimStringsExec`) at physical-planning time.
/// - The Hamelin optimizer-rule overrides (replaces upstream
///   `extract_leaf_expressions` / `push_down_leaf_projections` with the
///   Hamelin forks that fix DF 53 schema-validation issues, plus the
///   `JsonRewrite` analyzer when the `executor` feature is on).
///
/// Downstream callers (e.g. phantom) that build their own `SessionContext`
/// only need to call this — there's no separate optimizer setup.
pub fn register_extensions(builder: SessionStateBuilder) -> SessionStateBuilder {
    let planners: Vec<Arc<dyn ExtensionPlanner + Send + Sync>> =
        vec![Arc::new(TrimStringsExtensionPlanner)];
    let inner = DefaultPhysicalPlanner::with_extension_planners(planners);
    let builder = builder.with_query_planner(Arc::new(HamelinQueryPlanner { inner }));
    #[cfg(feature = "executor")]
    let builder = crate::optimizer::with_hamelin_optimizer_rules(builder);
    builder
}

/// Adapter that wraps a `PhysicalPlanner` (where extension planners live) so it can be
/// supplied to `SessionStateBuilder::with_query_planner`, which expects a `QueryPlanner`.
struct HamelinQueryPlanner {
    inner: DefaultPhysicalPlanner,
}

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

#[async_trait]
impl QueryPlanner for HamelinQueryPlanner {
    async fn create_physical_plan(
        &self,
        logical_plan: &LogicalPlan,
        session_state: &SessionState,
    ) -> DfResult<Arc<dyn ExecutionPlan>> {
        self.inner
            .create_physical_plan(logical_plan, session_state)
            .await
    }
}
