//! Statement translation from Hamelin IR to DataFusion LogicalPlan
//!
//! This module translates a lowered Hamelin statement (including CTEs) into a
//! DataFusion LogicalPlan.

use std::collections::HashMap;
use std::sync::Arc;

use datafusion::arrow::datatypes::{DataType, Schema};
use datafusion::common::Column as DfColumn;
use datafusion::datasource::DefaultTableSource;
use datafusion::logical_expr::builder::subquery_alias;
use datafusion::logical_expr::dml::InsertOp;
use datafusion::logical_expr::{ident, Expr, LogicalPlan, LogicalPlanBuilder, TryCast};
use datafusion_functions::core::expr_fn as core_fn;
use hamelin_executor::executor::ExecutorError;
use hamelin_lib::catalog::Column;
use hamelin_lib::err::TranslationError;
use hamelin_lib::tree::ast::dataset_identifier::QualifiedDatasetIdentifier;
use hamelin_lib::tree::ast::identifier::{Identifier, SimpleIdentifier};
use hamelin_translation::{IRSideEffect, IRStatement};

use crate::dataset_table_catalog::DatasetTableCatalog;
use crate::expr::ExprTranslationContext;
use crate::pipeline::translate_pipeline_with_ctes;

/// Result of translating a statement
pub enum TranslatedStatement {
    /// A query that returns data, with its Hamelin output schema
    Query {
        plan: LogicalPlan,
        /// The Hamelin type schema (column names and types from the typed AST)
        output_schema: Vec<Column>,
    },
    /// A DML statement that inserts data into a target table
    Dml {
        /// The source data plan (everything before APPEND)
        source_plan: LogicalPlan,
        /// The target dataset
        target_table: QualifiedDatasetIdentifier,
        /// Columns for deduplication (empty = plain INSERT, non-empty = anti-join + INSERT)
        distinct_by: Vec<Identifier>,
    },
}

impl TranslatedStatement {
    pub fn query(self) -> Option<(LogicalPlan, Vec<Column>)> {
        if let Self::Query {
            plan,
            output_schema,
        } = self
        {
            Some((plan, output_schema))
        } else {
            None
        }
    }

    pub fn dml(self) -> Option<(LogicalPlan, QualifiedDatasetIdentifier, Vec<Identifier>)> {
        if let Self::Dml {
            source_plan,
            target_table,
            distinct_by,
        } = self
        {
            Some((source_plan, target_table, distinct_by))
        } else {
            None
        }
    }

    /// Build the full INSERT logical plan for a DML statement, including the
    /// anti-join for DISTINCT BY. Returns an error if called on a Query.
    pub async fn into_dml_plan(
        self,
        catalog: &DatasetTableCatalog,
    ) -> Result<LogicalPlan, ExecutorError> {
        let (source_plan, target_table, distinct_by) = self.dml().ok_or_else(|| {
            ExecutorError::QueryError(anyhow::anyhow!("expected DML statement, got query").into())
        })?;

        let canonical = target_table.to_string();
        let scan_name = target_table.to_execution_identifier().to_string();

        let table_provider = catalog.provider(&target_table).map_err(|msg| {
            ExecutorError::QueryError(
                anyhow::anyhow!("Target table '{canonical}' not found: {msg}").into(),
            )
        })?;

        // Coerce source column types to the target's Arrow schema before any
        // further use, but ONLY for the mismatches the parquet writer's
        // per-field compatibility check rejects: the string-variant family
        // (Utf8/Utf8View/LargeUtf8) and the binary-variant family
        // (Binary/BinaryView/LargeBinary), including when those appear inside
        // struct/list/map children. DataFusion's INSERT analyzer already casts
        // ordinary top-level mismatches; other nested mismatches are left
        // untouched so its diagnostics surface instead of being silently
        // papered over.
        let source_plan =
            coerce_source_to_target_schema(source_plan, table_provider.schema().as_ref())?;

        let effective_source = if distinct_by.is_empty() {
            source_plan
        } else {
            let target_scan = LogicalPlanBuilder::scan(
                scan_name.clone(),
                Arc::new(DefaultTableSource::new(table_provider.clone())),
                None,
            )
            .map_err(|e| {
                ExecutorError::QueryError(
                    anyhow::anyhow!("Failed to scan target table for DISTINCT BY: {}", e).into(),
                )
            })?
            .build()
            .map_err(|e| {
                ExecutorError::QueryError(
                    anyhow::anyhow!("Failed to build target scan: {}", e).into(),
                )
            })?;

            let build_field_expr = |segments: &[SimpleIdentifier]| -> Expr {
                let mut expr = ident(segments[0].as_str());
                for seg in &segments[1..] {
                    expr = core_fn::get_field(expr, seg.as_str());
                }
                expr
            };

            let left_keys: Vec<Expr> = distinct_by
                .iter()
                .map(|id| build_field_expr(id.segments()))
                .collect();
            let right_keys: Vec<Expr> = distinct_by
                .iter()
                .map(|id| build_field_expr(id.segments()))
                .collect();

            LogicalPlanBuilder::from(source_plan)
                .join_with_expr_keys(
                    target_scan,
                    datafusion::logical_expr::JoinType::LeftAnti,
                    (left_keys, right_keys),
                    None,
                )
                .map_err(|e| {
                    ExecutorError::QueryError(
                        anyhow::anyhow!("Failed to build anti-join for DISTINCT BY: {}", e).into(),
                    )
                })?
                .build()
                .map_err(|e| {
                    ExecutorError::QueryError(
                        anyhow::anyhow!("Failed to finalize anti-join plan: {}", e).into(),
                    )
                })?
        };

        LogicalPlanBuilder::insert_into(
            effective_source,
            &scan_name,
            Arc::new(DefaultTableSource::new(table_provider)),
            InsertOp::Append,
        )
        .map_err(|e| {
            ExecutorError::QueryError(anyhow::anyhow!("Failed to build INSERT plan: {}", e).into())
        })?
        .build()
        .map_err(|e| {
            ExecutorError::QueryError(
                anyhow::anyhow!("Failed to finalize INSERT plan: {}", e).into(),
            )
        })
    }
}

/// Project the source plan so columns with a parquet-incompatible
/// string/binary-variant mismatch — at the top level or inside struct/list/map
/// children — are cast to the target's Arrow type. All other columns are
/// passed through unchanged so DataFusion's INSERT analyzer handles ordinary
/// top-level casts and its own diagnostics surface for genuine type
/// mismatches.
///
/// If column counts don't match we return the plan unchanged — `insert_into`
/// will then produce a more specific error than we could. If no column needs
/// coercion, we skip the projection entirely so plans without a mismatch
/// incur no extra node.
pub fn coerce_source_to_target_schema(
    source: LogicalPlan,
    target_schema: &Schema,
) -> Result<LogicalPlan, ExecutorError> {
    let source_schema = source.schema();

    if source_schema.fields().len() != target_schema.fields().len() {
        return Ok(source);
    }

    let mut needs_projection = false;
    let mut projections: Vec<Expr> = Vec::with_capacity(source_schema.fields().len());

    for ((qualifier, src_field), tgt_field) in
        source_schema.iter().zip(target_schema.fields().iter())
    {
        let col_expr = Expr::Column(DfColumn::new(qualifier.cloned(), src_field.name()));
        if needs_targeted_coercion(src_field.data_type(), tgt_field.data_type()) {
            needs_projection = true;
            projections.push(
                Expr::TryCast(TryCast::new(
                    Box::new(col_expr),
                    tgt_field.data_type().clone(),
                ))
                .alias(src_field.name()),
            );
        } else {
            projections.push(col_expr);
        }
    }

    if !needs_projection {
        return Ok(source);
    }

    LogicalPlanBuilder::from(source)
        .project(projections)
        .map_err(|e| {
            ExecutorError::QueryError(
                anyhow::anyhow!("Failed to build INSERT coercion projection: {}", e).into(),
            )
        })?
        .build()
        .map_err(|e| {
            ExecutorError::QueryError(
                anyhow::anyhow!("Failed to finalize INSERT coercion projection: {}", e).into(),
            )
        })
}

/// Returns true when `src` should be cast to `tgt` before handing the column
/// to DataFusion's INSERT path. This fires only for mismatches in the
/// string-variant family (Utf8/Utf8View/LargeUtf8) and the binary-variant
/// family (Binary/BinaryView/LargeBinary), including when they appear inside
/// struct/list/map children — the shapes parquet-rs's `types_compatible`
/// accepts at the top level but rejects when nested.
fn needs_targeted_coercion(src: &DataType, tgt: &DataType) -> bool {
    if src == tgt {
        return false;
    }

    match (src, tgt) {
        (DataType::Struct(s_fields), DataType::Struct(t_fields))
            if s_fields.len() == t_fields.len() =>
        {
            s_fields
                .iter()
                .zip(t_fields.iter())
                .any(|(s, t)| needs_targeted_coercion(s.data_type(), t.data_type()))
        }
        (DataType::List(s), DataType::List(t))
        | (DataType::LargeList(s), DataType::LargeList(t)) => {
            needs_targeted_coercion(s.data_type(), t.data_type())
        }
        (DataType::FixedSizeList(s, sn), DataType::FixedSizeList(t, tn)) if sn == tn => {
            needs_targeted_coercion(s.data_type(), t.data_type())
        }
        (DataType::Map(s, _), DataType::Map(t, _)) => {
            needs_targeted_coercion(s.data_type(), t.data_type())
        }
        _ => is_parquet_string_or_binary_variant_pair(src, tgt),
    }
}

/// Mirrors parquet-rs's `types_compatible` table — pairs that the parquet
/// writer treats as interchangeable at the top level but not inside nested
/// children. Kept narrowly in sync with the upstream predicate.
fn is_parquet_string_or_binary_variant_pair(a: &DataType, b: &DataType) -> bool {
    match a {
        DataType::Utf8 => matches!(b, DataType::LargeUtf8 | DataType::Utf8View),
        DataType::Utf8View => matches!(b, DataType::LargeUtf8 | DataType::Utf8),
        DataType::LargeUtf8 => matches!(b, DataType::Utf8 | DataType::Utf8View),
        DataType::Binary => matches!(b, DataType::LargeBinary | DataType::BinaryView),
        DataType::BinaryView => matches!(b, DataType::LargeBinary | DataType::Binary),
        DataType::LargeBinary => matches!(b, DataType::Binary | DataType::BinaryView),
        _ => false,
    }
}

/// Translate an IRStatement into a DataFusion LogicalPlan
///
/// The statement must have been lowered via `hamelin_translation::lower()` before
/// calling this function. Lowering ensures:
/// - Compound identifiers in AGG/WINDOW are flattened to simple identifiers
/// - SET/DROP commands are fused into SELECT commands
/// - WITHIN is converted to WHERE
/// - PARSE is converted to SET + WHERE
/// - UNNEST (struct) is converted to SELECT
///
/// This function:
/// 1. Translates each CTE pipeline to a LogicalPlan wrapped in SubqueryAlias
/// 2. Passes the CTE plans to the main pipeline translation
/// 3. When a FROM clause references a CTE, the CTE plan is used directly
///
/// Note: `filter()` and `transform()` calls must be lowered before calling this function.
/// Use `hamelin_translation::normalize_with().with_lower_filter().with_lower_transform()`
/// to enable lambda lowering during IR lowering. DataFusionProvidedExecutor does this automatically.
///
/// Tables are resolved from the [`DatasetTableCatalog`] passed by the caller.
pub async fn translate_statement(
    statement: &IRStatement,
    catalog: &DatasetTableCatalog,
) -> Result<TranslatedStatement, Arc<TranslationError>> {
    let mut expr_ctx = ExprTranslationContext::default();
    for (name, ir_expr) in &statement.scalar_defs {
        expr_ctx
            .scalar_def_bindings
            .insert(name.as_str().to_string(), ir_expr.0.clone());
    }

    // First, translate all CTEs
    // Note: CTEs can reference earlier CTEs, so we translate in order and accumulate
    let mut cte_plans: HashMap<Identifier, Arc<LogicalPlan>> = HashMap::new();

    for pd in &statement.pipeline_defs {
        let cte_name = pd.name.as_str().to_string();
        let cte_ident: Identifier = pd.name.clone().into();

        // Translate the CTE pipeline, passing previously translated CTEs
        // This allows later CTEs to reference earlier ones
        let cte_plan =
            translate_pipeline_with_ctes(&pd.pipeline, catalog, &cte_plans, &expr_ctx).await?;

        // Wrap in SubqueryAlias so it has the correct name in the schema
        let aliased_plan = subquery_alias(cte_plan, &cte_name)
            .map_err(|e| Arc::new(TranslationError::wrap(pd, e)))?;

        cte_plans.insert(cte_ident, Arc::new(aliased_plan));
    }

    // Translate the main pipeline (APPEND is no longer in the command list)
    let plan =
        translate_pipeline_with_ctes(&statement.pipeline, catalog, &cte_plans, &expr_ctx).await?;

    // Check if this is a DML statement
    if let IRSideEffect::Append { table, distinct_by } = &statement.side_effect {
        return Ok(TranslatedStatement::Dml {
            source_plan: plan,
            target_table: table.clone(),
            distinct_by: distinct_by.clone(),
        });
    }

    let output_schema: Vec<Column> = statement.pipeline.output_schema.as_struct().into();

    Ok(TranslatedStatement::Query {
        plan,
        output_schema,
    })
}

#[cfg(test)]
mod tests {
    use datafusion::arrow::datatypes::{DataType, Field, Fields};
    use datafusion::logical_expr::EmptyRelation;

    use super::*;

    fn empty_with_schema(fields: Vec<Field>) -> LogicalPlan {
        let schema = Schema::new(fields);
        LogicalPlan::EmptyRelation(EmptyRelation {
            produce_one_row: false,
            schema: Arc::new(schema.try_into().unwrap()),
        })
    }

    /// Helper returns the plan unchanged when every column already matches.
    #[test]
    fn coerce_is_noop_when_types_match() {
        let fields = vec![
            Field::new("a", DataType::Int64, true),
            Field::new("b", DataType::Utf8View, true),
        ];
        let source = empty_with_schema(fields.clone());
        let target = Schema::new(fields);

        let out = coerce_source_to_target_schema(source.clone(), &target).unwrap();
        // No projection node added — plan is structurally unchanged.
        assert!(
            matches!(out, LogicalPlan::EmptyRelation(_)),
            "expected identical plan, got {out:?}"
        );
    }

    /// Top-level Utf8 → Utf8View: a Cast projection is inserted and the output
    /// schema matches the target.
    #[test]
    fn coerce_casts_top_level_utf8_to_utf8view() {
        let source = empty_with_schema(vec![Field::new("payload", DataType::Utf8, true)]);
        let target = Schema::new(vec![Field::new("payload", DataType::Utf8View, true)]);

        let out = coerce_source_to_target_schema(source, &target).unwrap();
        assert!(
            matches!(out, LogicalPlan::Projection(_)),
            "expected a Projection wrapping the cast, got {out:?}"
        );
        assert_eq!(out.schema().field(0).data_type(), &DataType::Utf8View);
    }

    /// Struct children with Utf8 leaves are cast to Utf8View recursively — this
    /// is the mismatch shape that parquet-rs's `types_compatible` rejects and
    /// DataFusion's built-in INSERT coercion does not fix.
    #[test]
    fn coerce_casts_struct_children_recursively() {
        let src_struct = DataType::Struct(Fields::from(vec![
            Field::new("original", DataType::Utf8, true),
            Field::new("tag", DataType::Utf8View, true),
        ]));
        let tgt_struct = DataType::Struct(Fields::from(vec![
            Field::new("original", DataType::Utf8View, true),
            Field::new("tag", DataType::Utf8View, true),
        ]));

        let source = empty_with_schema(vec![Field::new("event", src_struct, true)]);
        let target = Schema::new(vec![Field::new("event", tgt_struct.clone(), true)]);

        let out = coerce_source_to_target_schema(source, &target).unwrap();
        assert!(matches!(out, LogicalPlan::Projection(_)));
        assert_eq!(out.schema().field(0).data_type(), &tgt_struct);
    }

    /// Mismatched column counts: return the plan unchanged so `insert_into`
    /// can produce its own descriptive error.
    #[test]
    fn coerce_passes_through_when_column_counts_differ() {
        let source = empty_with_schema(vec![Field::new("a", DataType::Int64, true)]);
        let target = Schema::new(vec![
            Field::new("a", DataType::Int64, true),
            Field::new("b", DataType::Utf8View, true),
        ]);

        let out = coerce_source_to_target_schema(source, &target).unwrap();
        assert!(matches!(out, LogicalPlan::EmptyRelation(_)));
    }

    /// Unrelated top-level mismatch (Int32 → Int64) is left alone so
    /// DataFusion's INSERT analyzer can cast or error with its own diagnostics.
    #[test]
    fn coerce_passes_through_unrelated_top_level_mismatch() {
        let source = empty_with_schema(vec![Field::new("a", DataType::Int32, true)]);
        let target = Schema::new(vec![Field::new("a", DataType::Int64, true)]);

        let out = coerce_source_to_target_schema(source, &target).unwrap();
        assert!(
            matches!(out, LogicalPlan::EmptyRelation(_)),
            "expected pass-through, got {out:?}"
        );
    }

    /// Unrelated nested mismatch (Struct<Int32> → Struct<Int64>) is left alone
    /// for the same reason — we only paper over the parquet string/binary
    /// variant cases.
    #[test]
    fn coerce_passes_through_unrelated_nested_mismatch() {
        let src_struct =
            DataType::Struct(Fields::from(vec![Field::new("n", DataType::Int32, true)]));
        let tgt_struct =
            DataType::Struct(Fields::from(vec![Field::new("n", DataType::Int64, true)]));

        let source = empty_with_schema(vec![Field::new("event", src_struct, true)]);
        let target = Schema::new(vec![Field::new("event", tgt_struct, true)]);

        let out = coerce_source_to_target_schema(source, &target).unwrap();
        assert!(
            matches!(out, LogicalPlan::EmptyRelation(_)),
            "expected pass-through, got {out:?}"
        );
    }

    /// A struct with mixed mismatches — one string-variant child and one
    /// unrelated child — still gets cast because arrow_cast will handle the
    /// whole column recursively. The unrelated part is not our concern to
    /// gate on; if the full cast fails, that's the error we want to surface.
    #[test]
    fn coerce_casts_whole_column_when_any_child_is_variant_mismatch() {
        let src_struct = DataType::Struct(Fields::from(vec![
            Field::new("tag", DataType::Utf8, true),
            Field::new("n", DataType::Int64, true),
        ]));
        let tgt_struct = DataType::Struct(Fields::from(vec![
            Field::new("tag", DataType::Utf8View, true),
            Field::new("n", DataType::Int64, true),
        ]));

        let source = empty_with_schema(vec![Field::new("event", src_struct, true)]);
        let target = Schema::new(vec![Field::new("event", tgt_struct.clone(), true)]);

        let out = coerce_source_to_target_schema(source, &target).unwrap();
        assert!(matches!(out, LogicalPlan::Projection(_)));
        assert_eq!(out.schema().field(0).data_type(), &tgt_struct);
    }
}
