// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

//! Test utilities for physical optimizer tests

use std::fmt::{Display, Formatter};
use std::sync::{Arc, LazyLock};

use arrow::array::Int32Array;
use arrow::compute::SortOptions;
use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
use arrow::record_batch::RecordBatch;
use datafusion::datasource::listing::PartitionedFile;
use datafusion::datasource::memory::MemorySourceConfig;
use datafusion::datasource::physical_plan::ParquetSource;
use datafusion::datasource::source::DataSourceExec;
use datafusion_common::config::ConfigOptions;
use datafusion_common::stats::Precision;
use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
use datafusion_common::utils::expr::COUNT_STAR_EXPANSION;
use datafusion_common::{
    ColumnStatistics, JoinType, NullEquality, Result, Statistics, internal_err,
};
use datafusion_datasource::file_scan_config::FileScanConfigBuilder;
use datafusion_execution::object_store::ObjectStoreUrl;
use datafusion_execution::{SendableRecordBatchStream, TaskContext};
use datafusion_expr::{WindowFrame, WindowFunctionDefinition};
use datafusion_functions_aggregate::count::count_udaf;
use datafusion_physical_expr::EquivalenceProperties;
use datafusion_physical_expr::aggregate::{AggregateExprBuilder, AggregateFunctionExpr};
use datafusion_physical_expr::expressions::{self, col};
use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
use datafusion_physical_expr_common::sort_expr::{
    LexOrdering, OrderingRequirements, PhysicalSortExpr,
};
use datafusion_physical_optimizer::PhysicalOptimizerRule;
use datafusion_physical_optimizer::limited_distinct_aggregation::LimitedDistinctAggregation;
use datafusion_physical_plan::aggregates::{
    AggregateExec, AggregateMode, PhysicalGroupBy,
};
use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec;
use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType};
use datafusion_physical_plan::filter::FilterExec;
use datafusion_physical_plan::joins::utils::{JoinFilter, JoinOn};
use datafusion_physical_plan::joins::{HashJoinExec, PartitionMode, SortMergeJoinExec};
use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec};
use datafusion_physical_plan::projection::{ProjectionExec, ProjectionExpr};
use datafusion_physical_plan::repartition::RepartitionExec;
use datafusion_physical_plan::sorts::sort::SortExec;
use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
use datafusion_physical_plan::streaming::{PartitionStream, StreamingTableExec};
use datafusion_physical_plan::tree_node::PlanContext;
use datafusion_physical_plan::union::UnionExec;
use datafusion_physical_plan::windows::{BoundedWindowAggExec, create_window_expr};
use datafusion_physical_plan::{
    DisplayAs, DisplayFormatType, ExecutionPlan, InputOrderMode, Partitioning,
    PlanProperties, SortOrderPushdownResult, displayable,
};

/// Create a non sorted parquet exec
pub fn parquet_exec(schema: SchemaRef) -> Arc<DataSourceExec> {
    let config = FileScanConfigBuilder::new(
        ObjectStoreUrl::parse("test:///").unwrap(),
        Arc::new(ParquetSource::new(schema)),
    )
    .with_file(PartitionedFile::new("x".to_string(), 100))
    .build();

    DataSourceExec::from_data_source(config)
}

/// Create a single parquet file that is sorted
pub(crate) fn parquet_exec_with_sort(
    schema: SchemaRef,
    output_ordering: Vec<LexOrdering>,
) -> Arc<DataSourceExec> {
    let config = FileScanConfigBuilder::new(
        ObjectStoreUrl::parse("test:///").unwrap(),
        Arc::new(ParquetSource::new(schema)),
    )
    .with_file(PartitionedFile::new("x".to_string(), 100))
    .with_output_ordering(output_ordering)
    .build();

    DataSourceExec::from_data_source(config)
}

fn int64_stats() -> ColumnStatistics {
    ColumnStatistics {
        null_count: Precision::Absent,
        sum_value: Precision::Absent,
        max_value: Precision::Exact(1_000_000.into()),
        min_value: Precision::Exact(0.into()),
        distinct_count: Precision::Absent,
        byte_size: Precision::Absent,
    }
}

fn column_stats() -> Vec<ColumnStatistics> {
    vec![
        int64_stats(), // a
        int64_stats(), // b
        int64_stats(), // c
        ColumnStatistics::default(),
        ColumnStatistics::default(),
    ]
}

/// Create parquet datasource exec using schema from [`schema`].
pub(crate) fn parquet_exec_with_stats(file_size: u64) -> Arc<DataSourceExec> {
    let mut statistics = Statistics::new_unknown(&schema());
    statistics.num_rows = Precision::Inexact(10000);
    statistics.column_statistics = column_stats();

    let config = FileScanConfigBuilder::new(
        ObjectStoreUrl::parse("test:///").unwrap(),
        Arc::new(ParquetSource::new(schema())),
    )
    .with_file(PartitionedFile::new("x".to_string(), file_size))
    .with_statistics(statistics)
    .build();

    assert_eq!(config.statistics().num_rows, Precision::Inexact(10000));
    DataSourceExec::from_data_source(config)
}

pub fn schema() -> SchemaRef {
    static SCHEMA: LazyLock<SchemaRef> = LazyLock::new(|| {
        Arc::new(Schema::new(vec![
            Field::new("a", DataType::Int64, true),
            Field::new("b", DataType::Int64, true),
            Field::new("c", DataType::Int64, true),
            Field::new("d", DataType::Int32, true),
            Field::new("e", DataType::Boolean, true),
        ]))
    });
    Arc::clone(&SCHEMA)
}

pub fn create_test_schema() -> Result<SchemaRef> {
    static SCHEMA: LazyLock<SchemaRef> = LazyLock::new(|| {
        let nullable_column = Field::new("nullable_col", DataType::Int32, true);
        let non_nullable_column = Field::new("non_nullable_col", DataType::Int32, false);
        Arc::new(Schema::new(vec![nullable_column, non_nullable_column]))
    });
    let schema = Arc::clone(&SCHEMA);
    Ok(schema)
}

pub fn create_test_schema2() -> Result<SchemaRef> {
    static SCHEMA: LazyLock<SchemaRef> = LazyLock::new(|| {
        let col_a = Field::new("col_a", DataType::Int32, true);
        let col_b = Field::new("col_b", DataType::Int32, true);
        Arc::new(Schema::new(vec![col_a, col_b]))
    });
    let schema = Arc::clone(&SCHEMA);
    Ok(schema)
}

// Generate a schema which consists of 5 columns (a, b, c, d, e)
pub fn create_test_schema3() -> Result<SchemaRef> {
    static SCHEMA: LazyLock<SchemaRef> = LazyLock::new(|| {
        let a = Field::new("a", DataType::Int32, true);
        let b = Field::new("b", DataType::Int32, false);
        let c = Field::new("c", DataType::Int32, true);
        let d = Field::new("d", DataType::Int32, false);
        let e = Field::new("e", DataType::Int32, false);
        Arc::new(Schema::new(vec![a, b, c, d, e]))
    });
    let schema = Arc::clone(&SCHEMA);
    Ok(schema)
}

pub fn sort_merge_join_exec(
    left: Arc<dyn ExecutionPlan>,
    right: Arc<dyn ExecutionPlan>,
    join_on: &JoinOn,
    join_type: &JoinType,
) -> Arc<dyn ExecutionPlan> {
    Arc::new(
        SortMergeJoinExec::try_new(
            left,
            right,
            join_on.clone(),
            None,
            *join_type,
            vec![SortOptions::default(); join_on.len()],
            NullEquality::NullEqualsNothing,
        )
        .unwrap(),
    )
}

/// make PhysicalSortExpr with default options
pub fn sort_expr(name: &str, schema: &Schema) -> PhysicalSortExpr {
    sort_expr_options(name, schema, SortOptions::default())
}

/// PhysicalSortExpr with specified options
pub fn sort_expr_options(
    name: &str,
    schema: &Schema,
    options: SortOptions,
) -> PhysicalSortExpr {
    PhysicalSortExpr {
        expr: col(name, schema).unwrap(),
        options,
    }
}

pub fn coalesce_partitions_exec(input: Arc<dyn ExecutionPlan>) -> Arc<dyn ExecutionPlan> {
    Arc::new(CoalescePartitionsExec::new(input))
}

pub fn memory_exec(schema: &SchemaRef) -> Arc<dyn ExecutionPlan> {
    MemorySourceConfig::try_new_exec(&[vec![]], Arc::clone(schema), None).unwrap()
}

pub fn hash_join_exec(
    left: Arc<dyn ExecutionPlan>,
    right: Arc<dyn ExecutionPlan>,
    on: JoinOn,
    filter: Option<JoinFilter>,
    join_type: &JoinType,
) -> Result<Arc<dyn ExecutionPlan>> {
    Ok(Arc::new(HashJoinExec::try_new(
        left,
        right,
        on,
        filter,
        join_type,
        None,
        PartitionMode::Partitioned,
        NullEquality::NullEqualsNothing,
        false,
    )?))
}

pub fn bounded_window_exec(
    col_name: &str,
    sort_exprs: impl IntoIterator<Item = PhysicalSortExpr>,
    input: Arc<dyn ExecutionPlan>,
) -> Arc<dyn ExecutionPlan> {
    bounded_window_exec_with_partition(col_name, sort_exprs, &[], input)
}

pub fn bounded_window_exec_with_partition(
    col_name: &str,
    sort_exprs: impl IntoIterator<Item = PhysicalSortExpr>,
    partition_by: &[Arc<dyn PhysicalExpr>],
    input: Arc<dyn ExecutionPlan>,
) -> Arc<dyn ExecutionPlan> {
    let sort_exprs = sort_exprs.into_iter().collect::<Vec<_>>();
    let schema = input.schema();
    let window_expr = create_window_expr(
        &WindowFunctionDefinition::AggregateUDF(count_udaf()),
        "count".to_owned(),
        &[col(col_name, &schema).unwrap()],
        partition_by,
        &sort_exprs,
        Arc::new(WindowFrame::new(Some(false))),
        schema,
        false,
        false,
        None,
    )
    .unwrap();

    Arc::new(
        BoundedWindowAggExec::try_new(
            vec![window_expr],
            Arc::clone(&input),
            InputOrderMode::Sorted,
            false,
        )
        .unwrap(),
    )
}

pub fn filter_exec(
    predicate: Arc<dyn PhysicalExpr>,
    input: Arc<dyn ExecutionPlan>,
) -> Arc<dyn ExecutionPlan> {
    Arc::new(FilterExec::try_new(predicate, input).unwrap())
}

pub fn sort_preserving_merge_exec(
    ordering: LexOrdering,
    input: Arc<dyn ExecutionPlan>,
) -> Arc<dyn ExecutionPlan> {
    Arc::new(SortPreservingMergeExec::new(ordering, input))
}

pub fn sort_preserving_merge_exec_with_fetch(
    ordering: LexOrdering,
    input: Arc<dyn ExecutionPlan>,
    fetch: usize,
) -> Arc<dyn ExecutionPlan> {
    Arc::new(SortPreservingMergeExec::new(ordering, input).with_fetch(Some(fetch)))
}

pub fn union_exec(input: Vec<Arc<dyn ExecutionPlan>>) -> Arc<dyn ExecutionPlan> {
    UnionExec::try_new(input).unwrap()
}

pub fn local_limit_exec(
    input: Arc<dyn ExecutionPlan>,
    fetch: usize,
) -> Arc<dyn ExecutionPlan> {
    Arc::new(LocalLimitExec::new(input, fetch))
}

pub fn global_limit_exec(
    input: Arc<dyn ExecutionPlan>,
    skip: usize,
    fetch: Option<usize>,
) -> Arc<dyn ExecutionPlan> {
    Arc::new(GlobalLimitExec::new(input, skip, fetch))
}

pub fn repartition_exec(input: Arc<dyn ExecutionPlan>) -> Arc<dyn ExecutionPlan> {
    Arc::new(RepartitionExec::try_new(input, Partitioning::RoundRobinBatch(10)).unwrap())
}

pub fn spr_repartition_exec(input: Arc<dyn ExecutionPlan>) -> Arc<dyn ExecutionPlan> {
    Arc::new(
        RepartitionExec::try_new(input, Partitioning::RoundRobinBatch(10))
            .unwrap()
            .with_preserve_order(),
    )
}

pub fn aggregate_exec(input: Arc<dyn ExecutionPlan>) -> Arc<dyn ExecutionPlan> {
    let schema = input.schema();
    Arc::new(
        AggregateExec::try_new(
            AggregateMode::Final,
            PhysicalGroupBy::default(),
            vec![],
            vec![],
            input,
            schema,
        )
        .unwrap(),
    )
}

pub fn sort_exec(
    ordering: LexOrdering,
    input: Arc<dyn ExecutionPlan>,
) -> Arc<dyn ExecutionPlan> {
    sort_exec_with_fetch(ordering, None, input)
}

pub fn sort_exec_with_preserve_partitioning(
    ordering: LexOrdering,
    input: Arc<dyn ExecutionPlan>,
) -> Arc<dyn ExecutionPlan> {
    Arc::new(SortExec::new(ordering, input).with_preserve_partitioning(true))
}

pub fn sort_exec_with_fetch(
    ordering: LexOrdering,
    fetch: Option<usize>,
    input: Arc<dyn ExecutionPlan>,
) -> Arc<dyn ExecutionPlan> {
    Arc::new(SortExec::new(ordering, input).with_fetch(fetch))
}

pub fn projection_exec(
    expr: Vec<(Arc<dyn PhysicalExpr>, String)>,
    input: Arc<dyn ExecutionPlan>,
) -> Result<Arc<dyn ExecutionPlan>> {
    let proj_exprs: Vec<ProjectionExpr> = expr
        .into_iter()
        .map(|(expr, alias)| ProjectionExpr { expr, alias })
        .collect();
    Ok(Arc::new(ProjectionExec::try_new(proj_exprs, input)?))
}

/// A test [`ExecutionPlan`] whose requirements can be configured.
#[derive(Debug)]
pub struct RequirementsTestExec {
    required_input_ordering: Option<LexOrdering>,
    maintains_input_order: bool,
    input: Arc<dyn ExecutionPlan>,
}

impl RequirementsTestExec {
    pub fn new(input: Arc<dyn ExecutionPlan>) -> Self {
        Self {
            required_input_ordering: None,
            maintains_input_order: true,
            input,
        }
    }

    /// sets the required input ordering
    pub fn with_required_input_ordering(
        mut self,
        required_input_ordering: Option<LexOrdering>,
    ) -> Self {
        self.required_input_ordering = required_input_ordering;
        self
    }

    /// set the maintains_input_order flag
    pub fn with_maintains_input_order(mut self, maintains_input_order: bool) -> Self {
        self.maintains_input_order = maintains_input_order;
        self
    }

    /// returns this ExecutionPlan as an `Arc<dyn ExecutionPlan>`
    pub fn into_arc(self) -> Arc<dyn ExecutionPlan> {
        Arc::new(self)
    }
}

impl DisplayAs for RequirementsTestExec {
    fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result {
        match t {
            DisplayFormatType::Default | DisplayFormatType::Verbose => {
                write!(f, "RequiredInputOrderingExec")
            }
            DisplayFormatType::TreeRender => {
                // TODO: collect info
                write!(f, "")
            }
        }
    }
}

impl ExecutionPlan for RequirementsTestExec {
    fn name(&self) -> &str {
        "RequiredInputOrderingExec"
    }

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

    fn required_input_ordering(&self) -> Vec<Option<OrderingRequirements>> {
        vec![
            self.required_input_ordering
                .as_ref()
                .map(|ordering| OrderingRequirements::from(ordering.clone())),
        ]
    }

    fn maintains_input_order(&self) -> Vec<bool> {
        vec![self.maintains_input_order]
    }

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

    fn with_new_children(
        self: Arc<Self>,
        children: Vec<Arc<dyn ExecutionPlan>>,
    ) -> Result<Arc<dyn ExecutionPlan>> {
        assert_eq!(children.len(), 1);
        Ok(RequirementsTestExec::new(Arc::clone(&children[0]))
            .with_required_input_ordering(self.required_input_ordering.clone())
            .with_maintains_input_order(self.maintains_input_order)
            .into_arc())
    }

    fn execute(
        &self,
        _partition: usize,
        _context: Arc<TaskContext>,
    ) -> Result<SendableRecordBatchStream> {
        unimplemented!("Test exec does not support execution")
    }
}

/// A [`PlanContext`] object is susceptible to being left in an inconsistent state after
/// untested mutable operations. It is crucial that there be no discrepancies between a plan
/// associated with the root node and the plan generated after traversing all nodes
/// within the [`PlanContext`] tree. In addition to verifying the plans resulting from optimizer
/// rules, it is essential to ensure that the overall tree structure corresponds with the plans
/// contained within the node contexts.
/// TODO: Once [`ExecutionPlan`] implements [`PartialEq`], string comparisons should be
/// replaced with direct plan equality checks.
pub fn check_integrity<T: Clone>(context: PlanContext<T>) -> Result<PlanContext<T>> {
    context
        .transform_up(|node| {
            let children_plans = node.plan.children();
            assert_eq!(node.children.len(), children_plans.len());
            for (child_plan, child_node) in
                children_plans.iter().zip(node.children.iter())
            {
                assert_eq!(
                    displayable(child_plan.as_ref()).one_line().to_string(),
                    displayable(child_node.plan.as_ref()).one_line().to_string()
                );
            }
            Ok(Transformed::no(node))
        })
        .data()
}

// construct a stream partition for test purposes
#[derive(Debug)]
pub struct TestStreamPartition {
    pub schema: SchemaRef,
}

impl PartitionStream for TestStreamPartition {
    fn schema(&self) -> &SchemaRef {
        &self.schema
    }
    fn execute(&self, _ctx: Arc<TaskContext>) -> SendableRecordBatchStream {
        unreachable!()
    }
}

/// Create an unbounded stream table without data ordering.
pub fn stream_exec(schema: &SchemaRef) -> Arc<dyn ExecutionPlan> {
    Arc::new(
        StreamingTableExec::try_new(
            Arc::clone(schema),
            vec![Arc::new(TestStreamPartition {
                schema: Arc::clone(schema),
            }) as _],
            None,
            vec![],
            true,
            None,
        )
        .unwrap(),
    )
}

/// Create an unbounded stream table with data ordering.
pub fn stream_exec_ordered(
    schema: &SchemaRef,
    ordering: LexOrdering,
) -> Arc<dyn ExecutionPlan> {
    Arc::new(
        StreamingTableExec::try_new(
            Arc::clone(schema),
            vec![Arc::new(TestStreamPartition {
                schema: Arc::clone(schema),
            }) as _],
            None,
            vec![ordering],
            true,
            None,
        )
        .unwrap(),
    )
}

/// Create an unbounded stream table with data ordering and built-in projection.
pub fn stream_exec_ordered_with_projection(
    schema: &SchemaRef,
    ordering: LexOrdering,
) -> Arc<dyn ExecutionPlan> {
    let projection: Vec<usize> = vec![0, 2, 3];

    Arc::new(
        StreamingTableExec::try_new(
            Arc::clone(schema),
            vec![Arc::new(TestStreamPartition {
                schema: Arc::clone(schema),
            }) as _],
            Some(&projection),
            vec![ordering],
            true,
            None,
        )
        .unwrap(),
    )
}

pub fn mock_data() -> Result<Arc<DataSourceExec>> {
    let schema = Arc::new(Schema::new(vec![
        Field::new("a", DataType::Int32, true),
        Field::new("b", DataType::Int32, true),
    ]));

    let batch = RecordBatch::try_new(
        Arc::clone(&schema),
        vec![
            Arc::new(Int32Array::from(vec![
                Some(1),
                Some(2),
                None,
                Some(1),
                Some(4),
                Some(5),
            ])),
            Arc::new(Int32Array::from(vec![
                Some(1),
                None,
                Some(6),
                Some(2),
                Some(8),
                Some(9),
            ])),
        ],
    )?;

    MemorySourceConfig::try_new_exec(&[vec![batch]], Arc::clone(&schema), None)
}

pub fn build_group_by(input_schema: &SchemaRef, columns: Vec<String>) -> PhysicalGroupBy {
    let mut group_by_expr: Vec<(Arc<dyn PhysicalExpr>, String)> = vec![];
    for column in columns.iter() {
        group_by_expr.push((col(column, input_schema).unwrap(), column.to_string()));
    }
    PhysicalGroupBy::new_single(group_by_expr.clone())
}

pub fn get_optimized_plan(plan: &Arc<dyn ExecutionPlan>) -> Result<String> {
    let config = ConfigOptions::new();

    let optimized =
        LimitedDistinctAggregation::new().optimize(Arc::clone(plan), &config)?;

    let optimized_result = displayable(optimized.as_ref()).indent(true).to_string();

    Ok(optimized_result)
}

/// Describe the type of aggregate being tested
pub enum TestAggregate {
    /// Testing COUNT(*) type aggregates
    CountStar,

    /// Testing for COUNT(column) aggregate
    ColumnA(Arc<Schema>),
}

impl TestAggregate {
    /// Create a new COUNT(*) aggregate
    pub fn new_count_star() -> Self {
        Self::CountStar
    }

    /// Create a new COUNT(column) aggregate
    pub fn new_count_column(schema: &Arc<Schema>) -> Self {
        Self::ColumnA(Arc::clone(schema))
    }

    /// Return appropriate expr depending if COUNT is for col or table (*)
    pub fn count_expr(&self, schema: &Schema) -> AggregateFunctionExpr {
        AggregateExprBuilder::new(count_udaf(), vec![self.column()])
            .schema(Arc::new(schema.clone()))
            .alias(self.column_name())
            .build()
            .unwrap()
    }

    /// what argument would this aggregate need in the plan?
    fn column(&self) -> Arc<dyn PhysicalExpr> {
        match self {
            Self::CountStar => expressions::lit(COUNT_STAR_EXPANSION),
            Self::ColumnA(s) => col("a", s).unwrap(),
        }
    }

    /// What name would this aggregate produce in a plan?
    pub fn column_name(&self) -> &'static str {
        match self {
            Self::CountStar => "COUNT(*)",
            Self::ColumnA(_) => "COUNT(a)",
        }
    }

    /// What is the expected count?
    pub fn expected_count(&self) -> i64 {
        match self {
            TestAggregate::CountStar => 3,
            TestAggregate::ColumnA(_) => 2,
        }
    }
}

/// A harness for testing physical optimizers.
#[derive(Debug)]
pub struct OptimizationTest {
    input: Vec<String>,
    output: Result<Vec<String>, String>,
}

impl OptimizationTest {
    pub fn new<O>(
        input_plan: Arc<dyn ExecutionPlan>,
        opt: O,
        enable_sort_pushdown: bool,
    ) -> Self
    where
        O: PhysicalOptimizerRule,
    {
        let input = format_execution_plan(&input_plan);
        let input_schema = input_plan.schema();

        let mut config = ConfigOptions::new();
        config.optimizer.enable_sort_pushdown = enable_sort_pushdown;
        let output_result = opt.optimize(input_plan, &config);
        let output = output_result
            .and_then(|plan| {
                if opt.schema_check() && (plan.schema() != input_schema) {
                    internal_err!(
                        "Schema mismatch:\n\nBefore:\n{:?}\n\nAfter:\n{:?}",
                        input_schema,
                        plan.schema()
                    )
                } else {
                    Ok(plan)
                }
            })
            .map(|plan| format_execution_plan(&plan))
            .map_err(|e| e.to_string());

        Self { input, output }
    }
}

impl Display for OptimizationTest {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        writeln!(f, "OptimizationTest:")?;
        writeln!(f, "  input:")?;
        for line in &self.input {
            writeln!(f, "    - {line}")?;
        }
        writeln!(f, "  output:")?;
        match &self.output {
            Ok(output) => {
                writeln!(f, "    Ok:")?;
                for line in output {
                    writeln!(f, "      - {line}")?;
                }
            }
            Err(err) => {
                writeln!(f, "    Err: {err}")?;
            }
        }
        Ok(())
    }
}

pub fn format_execution_plan(plan: &Arc<dyn ExecutionPlan>) -> Vec<String> {
    format_lines(&displayable(plan.as_ref()).indent(false).to_string())
}

fn format_lines(s: &str) -> Vec<String> {
    s.trim().split('\n').map(|s| s.to_string()).collect()
}

/// Create a simple ProjectionExec with column indices (simplified version)
pub fn simple_projection_exec(
    input: Arc<dyn ExecutionPlan>,
    columns: Vec<usize>,
) -> Arc<dyn ExecutionPlan> {
    let schema = input.schema();
    let exprs: Vec<(Arc<dyn PhysicalExpr>, String)> = columns
        .iter()
        .map(|&i| {
            let field = schema.field(i);
            (
                Arc::new(expressions::Column::new(field.name(), i))
                    as Arc<dyn PhysicalExpr>,
                field.name().to_string(),
            )
        })
        .collect();

    projection_exec(exprs, input).unwrap()
}

/// Create a ProjectionExec with column aliases
pub fn projection_exec_with_alias(
    input: Arc<dyn ExecutionPlan>,
    columns: Vec<(usize, &str)>,
) -> Arc<dyn ExecutionPlan> {
    let schema = input.schema();
    let exprs: Vec<(Arc<dyn PhysicalExpr>, String)> = columns
        .iter()
        .map(|&(i, alias)| {
            (
                Arc::new(expressions::Column::new(schema.field(i).name(), i))
                    as Arc<dyn PhysicalExpr>,
                alias.to_string(),
            )
        })
        .collect();

    projection_exec(exprs, input).unwrap()
}

/// Create a sort expression with custom name and index
pub fn sort_expr_named(name: &str, index: usize) -> PhysicalSortExpr {
    PhysicalSortExpr {
        expr: Arc::new(expressions::Column::new(name, index)),
        options: SortOptions::default(),
    }
}

/// A test data source that can display any requested ordering.
/// This is useful for testing sort pushdown behavior.
///
/// ## Configuration
///
/// - `exact_pushdown`: if `true`, `try_pushdown_sort` returns `Exact`
///   (source guarantees ordering, SortExec can be removed); if `false`
///   (default), returns `Inexact` (SortExec kept).
/// - `supports_fetch`: if `true`, `with_fetch()` returns `Some` so the
///   optimizer can push a LIMIT into the source; if `false` (default),
///   `with_fetch()` returns `None`, forcing a `GlobalLimitExec` wrapper.
#[derive(Debug, Clone)]
pub struct TestScan {
    schema: SchemaRef,
    output_ordering: Vec<LexOrdering>,
    plan_properties: Arc<PlanProperties>,
    // Store the requested ordering for display
    requested_ordering: Option<LexOrdering>,
    /// If true, `try_pushdown_sort` returns `Exact` instead of `Inexact`.
    exact_pushdown: bool,
    /// If true, `with_fetch()` returns `Some(...)` (source absorbs the limit).
    supports_fetch: bool,
    /// The fetch (LIMIT) value pushed into this scan via `with_fetch()`.
    fetch: Option<usize>,
}

impl TestScan {
    /// Create a new TestScan with the given schema and output ordering
    pub fn new(schema: SchemaRef, output_ordering: Vec<LexOrdering>) -> Self {
        let eq_properties = if !output_ordering.is_empty() {
            // Convert Vec<LexOrdering> to the format expected by new_with_orderings
            // We need to extract the inner Vec<PhysicalSortExpr> from each LexOrdering
            let orderings: Vec<Vec<PhysicalSortExpr>> = output_ordering
                .iter()
                .map(|lex_ordering| {
                    // LexOrdering implements IntoIterator, so we can collect it
                    lex_ordering.iter().cloned().collect()
                })
                .collect();

            EquivalenceProperties::new_with_orderings(Arc::clone(&schema), orderings)
        } else {
            EquivalenceProperties::new(Arc::clone(&schema))
        };

        let plan_properties = PlanProperties::new(
            eq_properties,
            Partitioning::UnknownPartitioning(1),
            EmissionType::Incremental,
            Boundedness::Bounded,
        );

        Self {
            schema,
            output_ordering,
            plan_properties: Arc::new(plan_properties),
            requested_ordering: None,
            exact_pushdown: false,
            supports_fetch: false,
            fetch: None,
        }
    }

    /// Create a TestScan with a single output ordering
    pub fn with_ordering(schema: SchemaRef, ordering: LexOrdering) -> Self {
        Self::new(schema, vec![ordering])
    }

    /// Set whether `try_pushdown_sort` returns `Exact` (true) or `Inexact` (false).
    pub fn with_exact_pushdown(mut self, exact: bool) -> Self {
        self.exact_pushdown = exact;
        self
    }

    /// Set whether `with_fetch()` returns `Some` (true) or `None` (false).
    pub fn with_supports_fetch(mut self, supports: bool) -> Self {
        self.supports_fetch = supports;
        self
    }
}

impl DisplayAs for TestScan {
    fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result {
        match t {
            DisplayFormatType::Default | DisplayFormatType::Verbose => {
                write!(f, "TestScan")?;
                let mut sep = ": ";
                if !self.output_ordering.is_empty() {
                    write!(f, "{sep}output_ordering=[")?;
                    for (i, sort_expr) in self.output_ordering[0].iter().enumerate() {
                        if i > 0 {
                            write!(f, ", ")?;
                        }
                        write!(f, "{sort_expr}")?;
                    }
                    write!(f, "]")?;
                    sep = ", ";
                }
                if let Some(ref req) = self.requested_ordering {
                    write!(f, "{sep}requested_ordering=[")?;
                    for (i, sort_expr) in req.iter().enumerate() {
                        if i > 0 {
                            write!(f, ", ")?;
                        }
                        write!(f, "{sort_expr}")?;
                    }
                    write!(f, "]")?;
                    sep = ", ";
                }
                if let Some(fetch) = self.fetch {
                    write!(f, "{sep}fetch={fetch}")?;
                }
                Ok(())
            }
            DisplayFormatType::TreeRender => {
                write!(f, "TestScan")
            }
        }
    }
}

impl ExecutionPlan for TestScan {
    fn name(&self) -> &str {
        "TestScan"
    }

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

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

    fn with_new_children(
        self: Arc<Self>,
        children: Vec<Arc<dyn ExecutionPlan>>,
    ) -> Result<Arc<dyn ExecutionPlan>> {
        if children.is_empty() {
            Ok(self)
        } else {
            internal_err!("TestScan should have no children")
        }
    }

    fn execute(
        &self,
        _partition: usize,
        _context: Arc<TaskContext>,
    ) -> Result<SendableRecordBatchStream> {
        internal_err!("TestScan is for testing optimizer only, not for execution")
    }

    fn partition_statistics(&self, _partition: Option<usize>) -> Result<Arc<Statistics>> {
        Ok(Arc::new(Statistics::new_unknown(&self.schema)))
    }

    fn with_fetch(&self, fetch: Option<usize>) -> Option<Arc<dyn ExecutionPlan>> {
        if self.supports_fetch {
            let mut new_scan = self.clone();
            new_scan.fetch = fetch;
            Some(Arc::new(new_scan))
        } else {
            None
        }
    }

    fn fetch(&self) -> Option<usize> {
        self.fetch
    }

    // This is the key method - implement sort pushdown
    fn try_pushdown_sort(
        &self,
        order: &[PhysicalSortExpr],
    ) -> Result<SortOrderPushdownResult<Arc<dyn ExecutionPlan>>> {
        // For testing purposes, accept ANY ordering request
        // and create a new TestScan that shows what was requested
        let requested_ordering = LexOrdering::new(order.to_vec());

        let mut new_scan = self.clone();
        new_scan.requested_ordering = requested_ordering;

        if self.exact_pushdown {
            // Update plan properties to reflect the guaranteed ordering
            let orderings: Vec<Vec<PhysicalSortExpr>> = vec![order.to_vec()];
            let eq_properties = EquivalenceProperties::new_with_orderings(
                Arc::clone(&self.schema),
                orderings,
            );
            new_scan.plan_properties = Arc::new(PlanProperties::new(
                eq_properties,
                Partitioning::UnknownPartitioning(1),
                EmissionType::Incremental,
                Boundedness::Bounded,
            ));
            Ok(SortOrderPushdownResult::Exact {
                inner: Arc::new(new_scan),
            })
        } else {
            Ok(SortOrderPushdownResult::Inexact {
                inner: Arc::new(new_scan),
            })
        }
    }
}

/// Helper function to create a TestScan with ordering
pub fn test_scan_with_ordering(
    schema: SchemaRef,
    ordering: LexOrdering,
) -> Arc<dyn ExecutionPlan> {
    Arc::new(TestScan::with_ordering(schema, ordering))
}
