pub mod hll;
pub mod optimizations;
pub mod panic;
mod partition_filter;
mod planning;
use datafusion::execution::runtime_env::RuntimeEnv;
use datafusion::logical_expr::planner::ExprPlanner;
use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
use datafusion_datasource::memory::MemorySourceConfig;
use datafusion_datasource::source::DataSourceExec;
pub use planning::PlanningMeta;
mod check_memory;
mod group_by_limit_aggregate;
pub mod physical_plan_flags;
pub mod pretty_printers;
mod projection_above_limit;
pub mod query_executor;
pub mod serialized_plan;
mod tail_limit;
mod topk;
pub mod trace_data_loaded;
use serialized_plan::PreSerializedPlan;
pub use topk::MIN_TOPK_STREAM_ROWS;
mod filter_by_key_range;
pub mod info_schema;
mod inline_aggregate;
#[cfg(test)]
mod is_not_distinct_from_join_test;
pub mod merge_sort;
pub mod metadata_cache;
pub mod providers;
mod rolling;
#[cfg(test)]
mod test_utils;
pub mod udf_xirr;
pub mod udfs;

use crate::cachestore::CacheStore;
use crate::config::injection::DIService;
use crate::config::ConfigObj;
use crate::metastore::multi_index::MultiPartition;
use crate::metastore::table::{Table, TablePath};
use crate::metastore::{IdRow, MetaStore};
use crate::queryplanner::info_schema::{
    ColumnsInfoSchemaTableDef, RocksDBPropertiesTableDef, SchemataInfoSchemaTableDef,
    SystemCacheTableDef, SystemChunksTableDef, SystemIndexesTableDef, SystemJobsTableDef,
    SystemPartitionsTableDef, SystemQueueResultsTableDef, SystemQueueTableDef,
    SystemReplayHandlesTableDef, SystemSnapshotsTableDef, SystemTablesTableDef,
    TablesInfoSchemaTableDef,
};
use crate::queryplanner::planning::{choose_index_ext, ClusterSendNode};
// TODO upgrade DF
// use crate::queryplanner::projection_above_limit::ProjectionAboveLimit;
use crate::queryplanner::query_executor::{
    batches_to_dataframe, ClusterSendExec, InlineTableProvider,
};
use crate::queryplanner::serialized_plan::SerializedPlan;
use crate::queryplanner::topk::ClusterAggregateTopKLower;

use crate::queryplanner::metadata_cache::MetadataCacheFactory;
use crate::queryplanner::optimizations::is_not_distinct_from_join_keys::IsNotDistinctFromJoinKeysRule;
use crate::queryplanner::optimizations::rolling_optimizer::RollingOptimizerRule;
use crate::queryplanner::pretty_printers::{pp_plan_ext, PPOptions};
use crate::queryplanner::udfs::{registerable_aggregate_udfs_iter, registerable_scalar_udfs_iter};
use crate::sql::cache::SqlResultCache;
use crate::sql::InlineTables;
use crate::store::DataFrame;
use crate::trace::{OpGuard, OpKind};
use crate::{app_metrics, metastore, CubeError};
use async_trait::async_trait;
use core::fmt;
use datafusion::arrow::array::ArrayRef;
use datafusion::arrow::datatypes::{DataType, Field};
use datafusion::arrow::record_batch::RecordBatch;
use datafusion::arrow::{datatypes::Schema, datatypes::SchemaRef};
use datafusion::catalog::Session;
use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion, TreeNodeVisitor};
use datafusion::common::{plan_datafusion_err, TableReference};
use datafusion::config::ConfigOptions;
use datafusion::datasource::{provider_as_source, TableType};
use datafusion::error::DataFusionError;
use datafusion::execution::{SessionState, SessionStateBuilder, TaskContext};
use datafusion::logical_expr::{
    AggregateUDF, Expr, Extension, LogicalPlan, ScalarUDF, TableProviderFilterPushDown,
    TableSource, WindowUDF,
};
use datafusion::physical_expr::EquivalenceProperties;
use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
use datafusion::physical_plan::{
    collect, DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties,
    SendableRecordBatchStream,
};
use datafusion::prelude::{SessionConfig, SessionContext};
use datafusion::sql::parser::Statement;
use datafusion::sql::planner::{ContextProvider, SqlToRel};
use datafusion::{cube_ext, datasource::TableProvider};
use log::{debug, trace};
use mockall::automock;
use serde_derive::{Deserialize, Serialize};
use smallvec::alloc::fmt::Formatter;
use std::any::Any;
use std::collections::{HashMap, HashSet};
use std::fmt::Debug;
use std::hash::{Hash, Hasher};
use std::sync::Arc;
use std::time::SystemTime;

#[automock]
#[async_trait]
pub trait QueryPlanner: DIService + Send + Sync {
    async fn logical_plan(
        &self,
        statement: Statement,
        inline_tables: &InlineTables,
        trace_obj: Option<String>,
    ) -> Result<QueryPlan, CubeError>;
    async fn execute_meta_plan(&self, plan: LogicalPlan) -> Result<DataFrame, CubeError>;
}

crate::di_service!(MockQueryPlanner, [QueryPlanner]);

pub struct QueryPlannerImpl {
    meta_store: Arc<dyn MetaStore>,
    cache_store: Arc<dyn CacheStore>,
    config: Arc<dyn ConfigObj>,
    cache: Arc<SqlResultCache>,
    metadata_cache_factory: Arc<dyn MetadataCacheFactory>,
}

crate::di_service!(QueryPlannerImpl, [QueryPlanner]);

pub enum QueryPlan {
    Meta(LogicalPlan),
    Select(PreSerializedPlan, /*workers*/ Vec<String>),
}

#[async_trait]
impl QueryPlanner for QueryPlannerImpl {
    async fn logical_plan(
        &self,
        statement: Statement,
        inline_tables: &InlineTables,
        trace_obj: Option<String>,
    ) -> Result<QueryPlan, CubeError> {
        let pre_execution_context_time = SystemTime::now();
        let ec_guard = OpGuard::start(OpKind::Planning, "plan.session_context");
        let ctx = self.execution_context()?;
        drop(ec_guard);

        let post_execution_context_time = SystemTime::now();
        app_metrics::DATA_QUERY_LOGICAL_PLAN_EXECUTION_CONTEXT_TIME_US.report(
            post_execution_context_time
                .duration_since(pre_execution_context_time)?
                .as_micros() as i64,
        );

        let state = Arc::new(ctx.state());
        let schema_provider = MetaStoreSchemaProvider::new(
            self.meta_store.get_tables_with_path(false).await?,
            self.meta_store.clone(),
            self.cache_store.clone(),
            inline_tables,
            self.cache.clone(),
            state.clone(),
        );

        let query_planner = SqlToRel::new_with_options(&schema_provider, sql_to_rel_options());

        let pre_statement_to_plan_time = SystemTime::now();
        let stp_guard = OpGuard::start(OpKind::Planning, "plan.statement_to_plan");
        let mut logical_plan = query_planner.statement_to_plan(statement)?;
        drop(stp_guard);
        let post_statement_to_plan_time = SystemTime::now();
        app_metrics::DATA_QUERY_LOGICAL_PLAN_QUERY_PLANNER_SETUP_TIME_US.report(
            pre_statement_to_plan_time
                .duration_since(post_execution_context_time)?
                .as_micros() as i64,
        );
        app_metrics::DATA_QUERY_LOGICAL_PLAN_STATEMENT_TO_PLAN_TIME_US.report(
            post_statement_to_plan_time
                .duration_since(pre_statement_to_plan_time)?
                .as_micros() as i64,
        );

        // TODO upgrade DF remove
        trace!(
            "Initial Logical Plan: {}",
            pp_plan_ext(
                &logical_plan,
                &PPOptions {
                    show_filters: true,
                    show_sort_by: true,
                    show_aggregations: true,
                    show_output_hints: true,
                    show_check_memory_nodes: false,
                    ..PPOptions::none()
                }
            )
        );

        let logical_plan_optimize_time = SystemTime::now();
        let opt_guard = OpGuard::start(OpKind::Planning, "plan.optimize");
        logical_plan = state.optimize(&logical_plan)?;
        drop(opt_guard);
        let post_optimize_time = SystemTime::now();
        app_metrics::DATA_QUERY_LOGICAL_PLAN_OPTIMIZE_TIME_US.report(
            post_optimize_time
                .duration_since(logical_plan_optimize_time)?
                .as_micros() as i64,
        );
        trace!(
            "Logical Plan: {}",
            pp_plan_ext(
                &logical_plan,
                &PPOptions {
                    show_filters: true,
                    show_sort_by: true,
                    show_aggregations: true,
                    show_output_hints: true,
                    show_check_memory_nodes: false,
                    ..PPOptions::none()
                }
            )
        );

        let post_is_data_select_query_time: SystemTime;
        let plan = if SerializedPlan::is_data_select_query(&logical_plan) {
            let choose_index_ext_start = SystemTime::now();
            post_is_data_select_query_time = choose_index_ext_start;
            let choose_guard = OpGuard::start_wrapper(OpKind::Planning, "plan.choose_index");
            let (logical_plan, meta) = choose_index_ext(
                logical_plan,
                &self.meta_store.as_ref(),
                self.config.enable_topk(),
            )
            .await?;
            let workers = compute_workers(
                self.config.as_ref(),
                &logical_plan,
                &meta.multi_part_subtree,
            )?;
            drop(choose_guard);
            app_metrics::DATA_QUERY_CHOOSE_INDEX_AND_WORKERS_TIME_US
                .report(choose_index_ext_start.elapsed()?.as_micros() as i64);
            QueryPlan::Select(
                PreSerializedPlan::try_new(logical_plan, meta, trace_obj)?,
                workers,
            )
        } else {
            post_is_data_select_query_time = SystemTime::now();
            QueryPlan::Meta(logical_plan)
        };
        app_metrics::DATA_QUERY_LOGICAL_PLAN_IS_DATA_SELECT_QUERY_US.report(
            post_is_data_select_query_time
                .duration_since(post_optimize_time)?
                .as_micros() as i64,
        );

        Ok(plan)
    }

    async fn execute_meta_plan(&self, plan: LogicalPlan) -> Result<DataFrame, CubeError> {
        let ctx = self.execution_context()?;

        let plan_ctx = ctx.clone();
        let plan_to_move = plan.clone();
        let physical_plan = plan_ctx.state().create_physical_plan(&plan_to_move).await?;

        let execution_time = SystemTime::now();
        let results = collect(physical_plan, ctx.task_ctx()).await?;
        let execution_time = execution_time.elapsed()?;
        app_metrics::META_QUERY_TIME_MS.report(execution_time.as_millis() as i64);
        debug!("Meta query data processing time: {:?}", execution_time,);
        let data_frame = cube_ext::spawn_blocking(move || batches_to_dataframe(results)).await??;
        Ok(data_frame)
    }
}

impl QueryPlannerImpl {
    pub fn new(
        meta_store: Arc<dyn MetaStore>,
        cache_store: Arc<dyn CacheStore>,
        config: Arc<dyn ConfigObj>,
        cache: Arc<SqlResultCache>,
        metadata_cache_factory: Arc<dyn MetadataCacheFactory>,
    ) -> Arc<QueryPlannerImpl> {
        Arc::new(QueryPlannerImpl {
            meta_store,
            cache_store,
            config,
            cache,
            metadata_cache_factory,
        })
    }
}

impl QueryPlannerImpl {
    /// Has the user defined functions to define query language behavior, but might exclude Cube
    /// optimizer rules or other parameters affecting execution performance.  This is used by
    /// `QueryPlannerImpl::make_execution_context`.
    pub fn minimal_session_state_from_final_config(config: SessionConfig) -> SessionStateBuilder {
        Self::minimal_session_state_from_final_config_with_runtime(
            config,
            Arc::new(RuntimeEnv::default()),
        )
    }

    pub fn minimal_session_state_from_final_config_with_runtime(
        config: SessionConfig,
        runtime_env: Arc<RuntimeEnv>,
    ) -> SessionStateBuilder {
        let mut state_builder = SessionStateBuilder::new()
            .with_config(config)
            .with_runtime_env(runtime_env)
            .with_default_features();
        state_builder
            .aggregate_functions()
            .get_or_insert_default()
            .extend(registerable_aggregate_udfs_iter().map(Arc::new));
        state_builder
            .scalar_functions()
            .get_or_insert_default()
            .extend(registerable_scalar_udfs_iter().map(Arc::new));
        state_builder
    }

    const EXECUTION_BATCH_SIZE: usize = 4096;

    pub fn make_execution_context(config: SessionConfig) -> SessionContext {
        Self::make_execution_context_with_runtime(config, Arc::new(RuntimeEnv::default()))
    }

    pub fn make_execution_context_with_runtime(
        mut config: SessionConfig,
        runtime_env: Arc<RuntimeEnv>,
    ) -> SessionContext {
        // The config parameter is from metadata_cache_factory (which we need to rename) but doesn't
        // include all necessary configs.
        config
            .options_mut()
            .execution
            .dont_parallelize_sort_preserving_merge_exec_inputs = true;
        config.options_mut().execution.batch_size = Self::EXECUTION_BATCH_SIZE;
        config.options_mut().execution.parquet.split_row_group_reads = false;

        // TODO upgrade DF: build SessionContexts consistently
        let state = Self::minimal_session_state_from_final_config_with_runtime(config, runtime_env)
            .with_optimizer_rule(Arc::new(RollingOptimizerRule {}))
            .with_optimizer_rule(Arc::new(IsNotDistinctFromJoinKeysRule {}))
            .build();

        let context = SessionContext::new_with_state(state);

        // TODO upgrade DF
        // context
        // .add_optimizer_rule(Arc::new(ProjectionAboveLimit {})),
        context
    }

    fn execution_context(&self) -> Result<Arc<SessionContext>, CubeError> {
        Ok(Arc::new(Self::make_execution_context(
            self.metadata_cache_factory.make_session_config(),
        )))
    }
}

struct MetaStoreSchemaProvider {
    /// Keeps the data used by [by_name] alive.
    _data: Arc<Vec<TablePath>>,
    by_name: HashSet<TableKey>,
    meta_store: Arc<dyn MetaStore>,
    cache_store: Arc<dyn CacheStore>,
    inline_tables: InlineTables,
    cache: Arc<SqlResultCache>,
    config_options: ConfigOptions,
    expr_planners: Vec<Arc<dyn ExprPlanner>>, // session_state.expr_planners clone
    session_state: Arc<SessionState>,
}

/// Points into [MetaStoreSchemaProvider::data], never null.
struct TableKey(*const TablePath);
unsafe impl Send for TableKey {}
unsafe impl Sync for TableKey {}

impl TableKey {
    fn qual_name(&self) -> (&str, &str) {
        let s = unsafe { &*self.0 };
        (s.schema_lower_name.as_str(), s.table_lower_name.as_str())
    }
}

impl PartialEq for TableKey {
    fn eq(&self, o: &Self) -> bool {
        self.qual_name() == o.qual_name()
    }
}
impl Eq for TableKey {}
impl Hash for TableKey {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.qual_name().hash(state)
    }
}

impl MetaStoreSchemaProvider {
    pub fn new(
        tables: Arc<Vec<TablePath>>,
        meta_store: Arc<dyn MetaStore>,
        cache_store: Arc<dyn CacheStore>,
        inline_tables: &InlineTables,
        cache: Arc<SqlResultCache>,
        session_state: Arc<SessionState>,
    ) -> Self {
        let by_name = tables.iter().map(|t| TableKey(t)).collect();
        Self {
            _data: tables,
            by_name,
            meta_store,
            cache_store,
            cache,
            inline_tables: (*inline_tables).clone(),
            config_options: ConfigOptions::new(),
            expr_planners: datafusion::execution::FunctionRegistry::expr_planners(
                session_state.as_ref(),
            ),
            session_state,
        }
    }
}

impl ContextProvider for MetaStoreSchemaProvider {
    fn get_table_source(
        &self,
        name: TableReference,
    ) -> Result<Arc<dyn TableSource>, DataFusionError> {
        let (schema, table) = match &name {
            TableReference::Partial { schema, table } => (schema.clone(), table.clone()),
            TableReference::Bare { table } => {
                let table = self
                    .inline_tables
                    .iter()
                    .find(|inline_table| inline_table.name == table.as_ref())
                    .ok_or_else(|| {
                        DataFusionError::Plan(format!("Inline table {} was not found", name))
                    })?;
                return Ok(provider_as_source(Arc::new(InlineTableProvider::new(
                    table.id,
                    table.data.clone(),
                    Vec::new(),
                ))));
            }
            TableReference::Full { .. } => {
                return Err(DataFusionError::Plan(format!(
                    "Catalog table names aren't supported but {} was provided",
                    name
                )))
            }
        };

        // Mock table path for hash set access.
        let table_path = TablePath::new(
            Arc::new(IdRow::new(0, metastore::Schema::new(schema.to_string()))),
            IdRow::new(
                u64::MAX,
                Table::new(
                    table.to_string(),
                    u64::MAX,
                    Vec::new(),
                    None,
                    None,
                    false,
                    None,
                    None,
                    None,
                    None,
                    None,
                    None,
                    Vec::new(),
                    None,
                    None,
                    None,
                ),
            ),
        );

        let res = self
            .by_name
            .get(&TableKey(&table_path))
            .map(|table| -> Arc<dyn TableProvider> {
                let table = unsafe { &*table.0 };
                let schema = Arc::new(Schema::new(
                    table
                        .table
                        .get_row()
                        .get_columns()
                        .iter()
                        .map(|c| c.clone().into())
                        .collect::<Vec<Field>>(),
                ));
                Arc::new(CubeTableLogical {
                    table: table.clone(),
                    schema,
                })
            });
        res.or_else(|| -> Option<Arc<dyn TableProvider>> {
            match (schema.as_ref(), table.as_ref()) {
                ("information_schema", "columns") => Some(Arc::new(InfoSchemaTableProvider::new(
                    self.meta_store.clone(),
                    self.cache_store.clone(),
                    InfoSchemaTable::Columns,
                ))),
                ("information_schema", "tables") => Some(Arc::new(InfoSchemaTableProvider::new(
                    self.meta_store.clone(),
                    self.cache_store.clone(),
                    InfoSchemaTable::Tables,
                ))),
                ("information_schema", "schemata") => Some(Arc::new(InfoSchemaTableProvider::new(
                    self.meta_store.clone(),
                    self.cache_store.clone(),
                    InfoSchemaTable::Schemata,
                ))),
                ("system", "query_cache") => Some(Arc::new(
                    providers::InfoSchemaQueryCacheTableProvider::new(self.cache.clone()),
                )),
                ("system", "cache") => Some(Arc::new(InfoSchemaTableProvider::new(
                    self.meta_store.clone(),
                    self.cache_store.clone(),
                    InfoSchemaTable::SystemCache,
                ))),
                ("system", "tables") => Some(Arc::new(InfoSchemaTableProvider::new(
                    self.meta_store.clone(),
                    self.cache_store.clone(),
                    InfoSchemaTable::SystemTables,
                ))),
                ("system", "indexes") => Some(Arc::new(InfoSchemaTableProvider::new(
                    self.meta_store.clone(),
                    self.cache_store.clone(),
                    InfoSchemaTable::SystemIndexes,
                ))),
                ("system", "partitions") => Some(Arc::new(InfoSchemaTableProvider::new(
                    self.meta_store.clone(),
                    self.cache_store.clone(),
                    InfoSchemaTable::SystemPartitions,
                ))),
                ("system", "chunks") => Some(Arc::new(InfoSchemaTableProvider::new(
                    self.meta_store.clone(),
                    self.cache_store.clone(),
                    InfoSchemaTable::SystemChunks,
                ))),
                ("system", "queue") => Some(Arc::new(InfoSchemaTableProvider::new(
                    self.meta_store.clone(),
                    self.cache_store.clone(),
                    InfoSchemaTable::SystemQueue,
                ))),
                ("system", "queue_results") => Some(Arc::new(InfoSchemaTableProvider::new(
                    self.meta_store.clone(),
                    self.cache_store.clone(),
                    InfoSchemaTable::SystemQueueResults,
                ))),
                ("system", "replay_handles") => Some(Arc::new(InfoSchemaTableProvider::new(
                    self.meta_store.clone(),
                    self.cache_store.clone(),
                    InfoSchemaTable::SystemReplayHandles,
                ))),
                ("system", "jobs") => Some(Arc::new(InfoSchemaTableProvider::new(
                    self.meta_store.clone(),
                    self.cache_store.clone(),
                    InfoSchemaTable::SystemJobs,
                ))),
                ("system", "snapshots") => Some(Arc::new(InfoSchemaTableProvider::new(
                    self.meta_store.clone(),
                    self.cache_store.clone(),
                    InfoSchemaTable::SystemSnapshots,
                ))),
                ("metastore", "rocksdb_properties") => {
                    Some(Arc::new(InfoSchemaTableProvider::new(
                        self.meta_store.clone(),
                        self.cache_store.clone(),
                        InfoSchemaTable::MetastoreRocksDBProperties,
                    )))
                }
                ("cachestore", "rocksdb_properties") => {
                    Some(Arc::new(InfoSchemaTableProvider::new(
                        self.meta_store.clone(),
                        self.cache_store.clone(),
                        InfoSchemaTable::CachestoreRocksDBProperties,
                    )))
                }
                _ => None,
            }
        })
        .map(|p| provider_as_source(p))
        .ok_or_else(|| {
            trace!("Table {} was not found, state: {:?}", name, self._data);

            DataFusionError::Plan(format!("Table {} was not found\n{:?}", name, table_path))
        })
    }

    fn get_table_function_source(
        &self,
        name: &str,
        args: Vec<Expr>,
    ) -> datafusion::common::Result<Arc<dyn TableSource>> {
        let tbl_func = self
            .session_state
            .table_functions()
            .get(name)
            .cloned()
            .ok_or_else(|| plan_datafusion_err!("table function '{name}' not found"))?;
        let provider = tbl_func.create_table_provider(&args)?;

        Ok(provider_as_source(provider))
    }

    fn get_function_meta(&self, name: &str) -> Option<Arc<ScalarUDF>> {
        let name = name.to_ascii_lowercase();
        self.session_state.scalar_functions().get(&name).cloned()
    }

    fn get_aggregate_meta(&self, name_param: &str) -> Option<Arc<AggregateUDF>> {
        let name = name_param.to_ascii_lowercase();
        self.session_state.aggregate_functions().get(&name).cloned()
    }

    fn get_window_meta(&self, name_param: &str) -> Option<Arc<WindowUDF>> {
        let name = name_param.to_ascii_lowercase();
        self.session_state.window_functions().get(&name).cloned()
    }

    fn get_variable_type(&self, _variable_names: &[String]) -> Option<DataType> {
        None
    }

    fn options(&self) -> &ConfigOptions {
        &self.config_options
    }

    fn udf_names(&self) -> Vec<String> {
        self.session_state
            .scalar_functions()
            .keys()
            .cloned()
            .collect()
    }

    fn udaf_names(&self) -> Vec<String> {
        self.session_state
            .aggregate_functions()
            .keys()
            .cloned()
            .collect()
    }

    fn udwf_names(&self) -> Vec<String> {
        self.session_state
            .window_functions()
            .keys()
            .cloned()
            .collect()
    }

    // We implement this for count(*) replacement.
    fn get_expr_planners(&self) -> &[Arc<dyn datafusion::logical_expr::planner::ExprPlanner>] {
        self.expr_planners.as_slice()
    }
}

/// Enables our options used with `SqlToRel`.  Sets `enable_ident_normalization` to false.  See also
/// `normalize_for_column_name` and its doc-comment, and similar functions, which must be kept in
/// sync with changes to the `enable_ident_normalization` option set here.
pub fn sql_to_rel_options() -> datafusion::sql::planner::ParserOptions {
    // not to be confused with sql_parser's ParserOptions
    datafusion::sql::planner::ParserOptions {
        enable_ident_normalization: false,
        ..Default::default()
    }
}

#[derive(Clone, Debug)]
pub enum InfoSchemaTable {
    Columns,
    Tables,
    Schemata,
    SystemJobs,
    SystemTables,
    SystemIndexes,
    SystemPartitions,
    SystemChunks,
    SystemQueue,
    SystemQueueResults,
    SystemReplayHandles,
    SystemCache,
    SystemSnapshots,
    CachestoreRocksDBProperties,
    MetastoreRocksDBProperties,
}

pub struct InfoSchemaTableDefContext {
    meta_store: Arc<dyn MetaStore>,
    cache_store: Arc<dyn CacheStore>,
}

#[async_trait]
pub trait InfoSchemaTableDef {
    type T: Send + Sync;

    async fn rows(
        &self,
        ctx: InfoSchemaTableDefContext,
        limit: Option<usize>,
    ) -> Result<Vec<Self::T>, CubeError>;

    fn columns(&self, rows: Vec<Self::T>) -> Vec<ArrayRef>;

    fn schema(&self) -> Vec<Field>;
}

#[async_trait]
pub trait BaseInfoSchemaTableDef {
    fn schema_ref(&self) -> SchemaRef;

    async fn scan(
        &self,
        ctx: InfoSchemaTableDefContext,
        limit: Option<usize>,
    ) -> Result<RecordBatch, CubeError>;
}

#[macro_export]
macro_rules! base_info_schema_table_def {
    ($table: ty) => {
        #[async_trait]
        impl crate::queryplanner::BaseInfoSchemaTableDef for $table {
            fn schema_ref(&self) -> datafusion::arrow::datatypes::SchemaRef {
                Arc::new(datafusion::arrow::datatypes::Schema::new(self.schema()))
            }

            async fn scan(
                &self,
                ctx: crate::queryplanner::InfoSchemaTableDefContext,
                limit: Option<usize>,
            ) -> Result<datafusion::arrow::record_batch::RecordBatch, crate::CubeError> {
                let rows = self.rows(ctx, limit).await?;
                let schema = self.schema_ref();
                let columns = self.columns(rows);
                Ok(datafusion::arrow::record_batch::RecordBatch::try_new(
                    schema, columns,
                )?)
            }
        }
    };
}

impl InfoSchemaTable {
    fn table_def(&self) -> Box<dyn BaseInfoSchemaTableDef + Send + Sync> {
        match self {
            InfoSchemaTable::Columns => Box::new(ColumnsInfoSchemaTableDef),
            InfoSchemaTable::Tables => Box::new(TablesInfoSchemaTableDef),
            InfoSchemaTable::Schemata => Box::new(SchemataInfoSchemaTableDef),
            InfoSchemaTable::SystemTables => Box::new(SystemTablesTableDef),
            InfoSchemaTable::SystemIndexes => Box::new(SystemIndexesTableDef),
            InfoSchemaTable::SystemChunks => Box::new(SystemChunksTableDef),
            InfoSchemaTable::SystemQueue => Box::new(SystemQueueTableDef),
            InfoSchemaTable::SystemQueueResults => Box::new(SystemQueueResultsTableDef),
            InfoSchemaTable::SystemReplayHandles => Box::new(SystemReplayHandlesTableDef),
            InfoSchemaTable::SystemPartitions => Box::new(SystemPartitionsTableDef),
            InfoSchemaTable::SystemJobs => Box::new(SystemJobsTableDef),
            InfoSchemaTable::SystemCache => Box::new(SystemCacheTableDef),
            InfoSchemaTable::SystemSnapshots => Box::new(SystemSnapshotsTableDef),
            InfoSchemaTable::CachestoreRocksDBProperties => {
                Box::new(RocksDBPropertiesTableDef::new_cachestore())
            }
            InfoSchemaTable::MetastoreRocksDBProperties => {
                Box::new(RocksDBPropertiesTableDef::new_metastore())
            }
        }
    }

    fn schema(&self) -> SchemaRef {
        self.table_def().schema_ref()
    }

    async fn scan(
        &self,
        ctx: InfoSchemaTableDefContext,
        limit: Option<usize>,
    ) -> Result<RecordBatch, CubeError> {
        self.table_def().scan(ctx, limit).await
    }
}

pub struct InfoSchemaTableProvider {
    meta_store: Arc<dyn MetaStore>,
    cache_store: Arc<dyn CacheStore>,
    table: InfoSchemaTable,
}

impl InfoSchemaTableProvider {
    pub fn new(
        meta_store: Arc<dyn MetaStore>,
        cache_store: Arc<dyn CacheStore>,
        table: InfoSchemaTable,
    ) -> Self {
        Self {
            meta_store,
            cache_store,
            table,
        }
    }
}

impl Debug for InfoSchemaTableProvider {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(f, "InfoSchemaTableProvider")
    }
}

#[async_trait]
impl TableProvider for InfoSchemaTableProvider {
    fn as_any(&self) -> &dyn Any {
        self
    }

    fn schema(&self) -> SchemaRef {
        self.table.schema()
    }

    fn table_type(&self) -> TableType {
        TableType::Base
    }

    async fn scan(
        &self,
        _state: &dyn Session,
        projection: Option<&Vec<usize>>,
        _filters: &[Expr],
        limit: Option<usize>,
    ) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
        let schema = project_schema(&self.schema(), projection.cloned().as_deref());
        let exec = InfoSchemaTableExec {
            meta_store: self.meta_store.clone(),
            cache_store: self.cache_store.clone(),
            table: self.table.clone(),
            projection: projection.cloned(),
            projected_schema: schema.clone(),
            limit,
            properties: PlanProperties::new(
                EquivalenceProperties::new(schema),
                Partitioning::UnknownPartitioning(1),
                EmissionType::Final,
                Boundedness::Bounded,
            ),
        };
        Ok(Arc::new(exec))
    }
}

fn project_schema(s: &Schema, projection: Option<&[usize]>) -> SchemaRef {
    let projection = match projection {
        None => return Arc::new(s.clone()),
        Some(p) => p,
    };
    let mut fields = Vec::with_capacity(projection.len());
    for &i in projection {
        fields.push(s.field(i).clone())
    }
    Arc::new(Schema::new(fields))
}

#[derive(Clone)]
pub struct InfoSchemaTableExec {
    meta_store: Arc<dyn MetaStore>,
    cache_store: Arc<dyn CacheStore>,
    table: InfoSchemaTable,
    projected_schema: SchemaRef,
    projection: Option<Vec<usize>>,
    limit: Option<usize>,
    properties: PlanProperties,
}

impl fmt::Debug for InfoSchemaTableExec {
    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), fmt::Error> {
        f.write_fmt(format_args!("{:?}", self.table))
    }
}

impl DisplayAs for InfoSchemaTableExec {
    fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result {
        write!(f, "InfoSchemaTableExec")
    }
}

#[async_trait]
impl ExecutionPlan for InfoSchemaTableExec {
    fn as_any(&self) -> &dyn Any {
        self
    }

    fn schema(&self) -> SchemaRef {
        self.projected_schema.clone()
    }

    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
        vec![]
    }

    fn with_new_children(
        self: Arc<Self>,
        _children: Vec<Arc<dyn ExecutionPlan>>,
    ) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
        Ok(self.clone())
    }

    fn execute(
        &self,
        partition: usize,
        _context: Arc<TaskContext>,
    ) -> Result<SendableRecordBatchStream, DataFusionError> {
        if partition != 0 {
            return datafusion::common::internal_err!(
                "invalid partition {} for InfoSchemaTableExec",
                partition
            );
        }
        let table_def = InfoSchemaTableDefContext {
            meta_store: self.meta_store.clone(),
            cache_store: self.cache_store.clone(),
        };
        let table = self.table.clone();
        let limit = self.limit.clone();
        let projection = self.projection.clone();
        let batch = async move {
            let mut batch = table
                .scan(table_def, limit)
                .await
                .map_err(|e| DataFusionError::Execution(e.to_string()))?;
            if let Some(projection) = projection {
                batch = batch.project(projection.as_slice())?;
            }
            Ok(batch)
        };

        let stream = futures::stream::once(batch);
        Ok(Box::pin(RecordBatchStreamAdapter::new(
            self.projected_schema.clone(),
            stream,
        )))
    }

    fn name(&self) -> &str {
        "InfoSchemaTableExec"
    }

    fn properties(&self) -> &PlanProperties {
        &self.properties
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CubeTableLogical {
    pub table: TablePath,
    schema: SchemaRef,
}

#[async_trait]
impl TableProvider for CubeTableLogical {
    fn as_any(&self) -> &dyn Any {
        self
    }

    fn schema(&self) -> SchemaRef {
        self.schema.clone()
    }

    fn table_type(&self) -> TableType {
        TableType::Base
    }

    async fn scan(
        &self,
        _state: &dyn Session,
        _projection: Option<&Vec<usize>>,
        _filters: &[Expr],
        _limit: Option<usize>,
    ) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
        panic!("scan has been called on CubeTableLogical: serialized plan wasn't preprocessed for select");
    }

    fn supports_filters_pushdown(
        &self,
        filters: &[&Expr],
    ) -> datafusion::common::Result<Vec<TableProviderFilterPushDown>> {
        Ok(vec![TableProviderFilterPushDown::Inexact; filters.len()])
    }
}

fn compute_workers(
    config: &dyn ConfigObj,
    p: &LogicalPlan,
    tree: &HashMap<u64, MultiPartition>,
) -> Result<Vec<String>, CubeError> {
    struct Visitor<'a> {
        config: &'a dyn ConfigObj,
        tree: &'a HashMap<u64, MultiPartition>,
        workers: Vec<String>,
    }
    impl<'a> TreeNodeVisitor<'a> for Visitor<'a> {
        type Node = LogicalPlan;

        fn f_down(&mut self, plan: &LogicalPlan) -> Result<TreeNodeRecursion, DataFusionError> {
            match plan {
                LogicalPlan::Extension(Extension { node }) => {
                    let snapshots =
                        if let Some(cs) = node.as_any().downcast_ref::<ClusterSendNode>() {
                            &cs.snapshots
                        } else if let Some(cs) =
                            node.as_any().downcast_ref::<ClusterAggregateTopKLower>()
                        {
                            &cs.snapshots
                        } else {
                            return Ok(TreeNodeRecursion::Continue);
                        };

                    let workers = ClusterSendExec::distribute_to_workers(
                        self.config,
                        snapshots.as_slice(),
                        self.tree,
                    )?;
                    self.workers = workers.into_iter().map(|w| w.0).collect();
                    Ok(TreeNodeRecursion::Stop)
                }
                _ => Ok(TreeNodeRecursion::Continue),
            }
        }
    }

    let mut v = Visitor {
        config,
        tree,
        workers: Vec::new(),
    };
    match p.visit(&mut v) {
        Ok(TreeNodeRecursion::Stop) => Ok(v.workers),
        Ok(TreeNodeRecursion::Continue) | Ok(TreeNodeRecursion::Jump) => Err(CubeError::internal(
            "no cluster send node found in plan".to_string(),
        )),
        Err(e) => Err(CubeError::internal(e.to_string())),
    }
}

/// Creates a [`DataSourceExec`] with a [`MemorySourceConfig`], i.e. the alternative to the
/// deprecated `MemoryExec`.  Useful when the [`MemorySourceConfig`] doesn't need sorting
/// information.
pub fn try_make_memory_data_source(
    partitions: &[Vec<RecordBatch>],
    schema: SchemaRef,
    projection: Option<Vec<usize>>,
) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
    Ok(Arc::new(DataSourceExec::new(Arc::new(
        MemorySourceConfig::try_new(partitions, schema, projection)?,
    ))))
}

#[cfg(test)]
pub mod tests {
    use super::*;

    use crate::queryplanner::serialized_plan::SerializedPlan;
    use crate::sql::parser::{CubeStoreParser, Statement};

    use datafusion::sql::parser::Statement as DFStatement;
    use datafusion::sql::planner::SqlToRel;
    use pretty_assertions::assert_eq;

    fn initial_plan(s: &str, ctx: MetaStoreSchemaProvider) -> LogicalPlan {
        let statement = match CubeStoreParser::new(s, None)
            .unwrap()
            .parse_statement()
            .unwrap()
        {
            Statement::Statement(s) => s,
            other => panic!("not a statement, actual {:?}", other),
        };

        let plan = SqlToRel::new_with_options(&ctx, sql_to_rel_options())
            .statement_to_plan(DFStatement::Statement(Box::new(statement)))
            .unwrap();
        SessionContext::new().state().optimize(&plan).unwrap()
    }

    fn get_test_execution_ctx() -> MetaStoreSchemaProvider {
        MetaStoreSchemaProvider::new(
            Arc::new(vec![]),
            Arc::new(test_utils::MetaStoreMock {}),
            Arc::new(test_utils::CacheStoreMock {}),
            &vec![],
            Arc::new(SqlResultCache::new(1 << 20, None, 10000, None)),
            Arc::new(SessionContext::new().state()),
        )
    }

    #[tokio::test]
    pub async fn test_is_data_select_query() {
        let plan = initial_plan(
            "SELECT * FROM information_schema.columns",
            get_test_execution_ctx(),
        );
        assert_eq!(SerializedPlan::is_data_select_query(&plan), false);

        let plan = initial_plan(
            "SELECT * FROM information_schema.columns as r",
            get_test_execution_ctx(),
        );
        assert_eq!(SerializedPlan::is_data_select_query(&plan), false);

        let plan = initial_plan("select * from system.query_cache", get_test_execution_ctx());
        assert_eq!(SerializedPlan::is_data_select_query(&plan), false);

        let plan = initial_plan("SELECT * FROM system.cache", get_test_execution_ctx());
        assert_eq!(SerializedPlan::is_data_select_query(&plan), false);

        // NOW is no longer a UDF.
        let plan = initial_plan("SELECT NOW()", get_test_execution_ctx());
        assert_eq!(SerializedPlan::is_data_select_query(&plan), false);
    }
}
