//! Plan walker that converts a DataFusion LogicalPlan into structured PlanSteps.
//!
//! Walks the plan tree bottom-up (children first), producing a flat list
//! of plan steps with sequential numbering.

use super::types::*;
use datafusion::common::JoinType;
use datafusion::logical_expr::{Expr, LogicalPlan};
use std::sync::atomic::{AtomicUsize, Ordering};

/// Walk a DataFusion LogicalPlan tree and produce a list of PlanSteps.
///
/// The walk is bottom-up: child nodes are visited before their parents,
/// so the output reflects the data flow from table scans up to the final
/// projection.
pub fn walk_plan(plan: &LogicalPlan) -> Vec<PlanStep> {
    let counter = AtomicUsize::new(0);
    walk_recursive(plan, &counter)
}

fn walk_recursive(plan: &LogicalPlan, counter: &AtomicUsize) -> Vec<PlanStep> {
    let mut steps = Vec::new();

    // Visit children first (bottom-up)
    // For Join and SubqueryAlias, we handle children specially
    match plan {
        LogicalPlan::Join(join) => {
            // Walk left and right children
            let left_steps = walk_recursive(&join.left, counter);
            let right_steps = walk_recursive(&join.right, counter);
            steps.extend(left_steps);
            steps.extend(right_steps);
        }
        LogicalPlan::SubqueryAlias(alias) => {
            // Walk inner plan separately for subquery wrapping
            let inner_steps = walk_recursive(&alias.input, counter);
            let step_num = counter.fetch_add(1, Ordering::Relaxed) + 1;
            steps.push(PlanStep {
                step: step_num,
                operation: PlanOperation::Subquery {
                    alias: alias.alias.table().to_string(),
                    inner_steps,
                },
                notes: None,
            });
            return steps;
        }
        _ => {
            for child in plan.inputs() {
                let child_steps = walk_recursive(child, counter);
                steps.extend(child_steps);
            }
        }
    }

    // Extract subquery plans embedded in expressions.
    // EXISTS, IN subquery, and scalar subqueries can appear in Filter
    // predicates or Projection expressions, not as child plan nodes.
    match plan {
        LogicalPlan::Filter(filter) => {
            let mut subquery_plans = Vec::new();
            extract_subquery_plans(&filter.predicate, &mut subquery_plans);
            for sub_plan in subquery_plans {
                let sub_steps = walk_recursive(sub_plan, counter);
                steps.extend(sub_steps);
            }
        }
        LogicalPlan::Projection(proj) => {
            let mut subquery_plans = Vec::new();
            for expr in &proj.expr {
                extract_subquery_plans(expr, &mut subquery_plans);
            }
            for sub_plan in subquery_plans {
                let sub_steps = walk_recursive(sub_plan, counter);
                steps.extend(sub_steps);
            }
        }
        _ => {}
    }

    // Now visit the current node
    if let Some(operation) = convert_node(plan) {
        let step_num = counter.fetch_add(1, Ordering::Relaxed) + 1;
        steps.push(PlanStep {
            step: step_num,
            operation,
            notes: None,
        });
    }

    steps
}

/// Recursively extract LogicalPlan references from subquery expressions
/// (EXISTS, IN subquery, scalar subquery) within an Expr tree.
fn extract_subquery_plans<'a>(expr: &'a Expr, plans: &mut Vec<&'a LogicalPlan>) {
    match expr {
        Expr::Exists(exists) => {
            plans.push(&exists.subquery.subquery);
        }
        Expr::InSubquery(in_sub) => {
            plans.push(&in_sub.subquery.subquery);
            extract_subquery_plans(&in_sub.expr, plans);
        }
        Expr::ScalarSubquery(sub) => {
            plans.push(&sub.subquery);
        }
        // Recurse into boolean combinators
        Expr::BinaryExpr(binary) => {
            extract_subquery_plans(&binary.left, plans);
            extract_subquery_plans(&binary.right, plans);
        }
        Expr::Not(inner) => {
            extract_subquery_plans(inner, plans);
        }
        // CASE WHEN ... THEN ... ELSE ... END may contain subqueries
        Expr::Case(case) => {
            if let Some(ref operand) = case.expr {
                extract_subquery_plans(operand, plans);
            }
            for (when_expr, then_expr) in &case.when_then_expr {
                extract_subquery_plans(when_expr, plans);
                extract_subquery_plans(then_expr, plans);
            }
            if let Some(ref else_expr) = case.else_expr {
                extract_subquery_plans(else_expr, plans);
            }
        }
        // Cast and TryCast may wrap subqueries
        Expr::Cast(cast) => {
            extract_subquery_plans(&cast.expr, plans);
        }
        Expr::TryCast(cast) => {
            extract_subquery_plans(&cast.expr, plans);
        }
        // Alias wraps another expression
        Expr::Alias(alias) => {
            extract_subquery_plans(&alias.expr, plans);
        }
        _ => {}
    }
}

/// Convert a single LogicalPlan node to a PlanOperation.
/// Returns None for nodes that should be skipped (e.g., Repartition).
fn convert_node(plan: &LogicalPlan) -> Option<PlanOperation> {
    match plan {
        LogicalPlan::TableScan(scan) => {
            let table = scan.table_name.table().to_string();
            let columns_read: Vec<String> = scan
                .projected_schema
                .fields()
                .iter()
                .map(|f| f.name().to_string())
                .collect();
            let filters: Vec<String> = scan.filters.iter().map(|f| format!("{f}")).collect();
            Some(PlanOperation::TableScan {
                table,
                columns_read,
                filters,
            })
        }

        LogicalPlan::Join(join) => {
            let left_name = extract_table_name_from_plan(&join.left)
                .unwrap_or_else(|| "left_input".to_string());
            let right_name = extract_table_name_from_plan(&join.right)
                .unwrap_or_else(|| "right_input".to_string());

            let join_type = convert_join_type(join.join_type);

            let mut condition_parts: Vec<String> =
                join.on.iter().map(|(l, r)| format!("{l} = {r}")).collect();
            if let Some(ref filter) = join.filter {
                condition_parts.push(format!("{filter}"));
            }
            let condition = if condition_parts.is_empty() {
                "no condition".to_string()
            } else {
                condition_parts.join(" AND ")
            };

            Some(PlanOperation::Join {
                join_type,
                left: left_name,
                right: right_name,
                condition,
            })
        }

        LogicalPlan::Filter(filter) => Some(PlanOperation::Filter {
            predicate: format!("{}", filter.predicate),
        }),

        LogicalPlan::Aggregate(agg) => {
            let group_by: Vec<String> = agg.group_expr.iter().map(|e| format!("{e}")).collect();
            let functions: Vec<String> = agg.aggr_expr.iter().map(|e| format!("{e}")).collect();
            Some(PlanOperation::Aggregate {
                group_by,
                functions,
            })
        }

        LogicalPlan::Sort(sort) => {
            let order_by: Vec<SortKey> = sort
                .expr
                .iter()
                .map(|se| SortKey {
                    column: format!("{}", se.expr),
                    direction: if se.asc {
                        SortDirection::Asc
                    } else {
                        SortDirection::Desc
                    },
                    nulls_first: se.nulls_first,
                })
                .collect();
            Some(PlanOperation::Sort { order_by })
        }

        LogicalPlan::Limit(limit) => {
            let fetch = limit.fetch.as_ref().and_then(|e| extract_literal_usize(e));
            let skip = limit.skip.as_ref().and_then(|e| extract_literal_usize(e));
            Some(PlanOperation::Limit {
                limit: fetch,
                offset: skip,
            })
        }

        LogicalPlan::Projection(proj) => {
            let columns: Vec<String> = proj.expr.iter().map(|e| format!("{e}")).collect();
            Some(PlanOperation::Projection { columns })
        }

        LogicalPlan::Window(win) => {
            // Extract window function details
            let mut function = String::new();
            let mut partition_by = Vec::new();
            let mut order_by = Vec::new();

            for expr in &win.window_expr {
                function = format!("{expr}");
                // Try to extract partition/order from the expression string
                // The full details are in the expression itself
            }
            // Try to parse partition_by and order_by from the window expressions
            // DataFusion's Window node stores them in the expressions
            for expr in &win.window_expr {
                if let datafusion::logical_expr::Expr::WindowFunction(wf) = expr {
                    partition_by = wf
                        .params
                        .partition_by
                        .iter()
                        .map(|e| format!("{e}"))
                        .collect();
                    order_by = wf
                        .params
                        .order_by
                        .iter()
                        .map(|se| format!("{}", se.expr))
                        .collect();
                    function = format!("{expr}");
                    break;
                }
            }

            Some(PlanOperation::Window {
                function,
                partition_by,
                order_by,
            })
        }

        LogicalPlan::Union(union) => Some(PlanOperation::Union {
            all: true, // DataFusion optimizes UNION to UNION ALL + Distinct
            branches: union.inputs.len(),
        }),

        LogicalPlan::Distinct(_) => Some(PlanOperation::Distinct),

        LogicalPlan::EmptyRelation(_) => Some(PlanOperation::EmptyRelation),

        // SubqueryAlias is handled in walk_recursive
        LogicalPlan::SubqueryAlias(_) => None,

        // Skip other node types (Repartition, Extension, etc.)
        _ => None,
    }
}

/// Extract a table name from a plan by looking for the first TableScan.
pub fn extract_table_name_from_plan(plan: &LogicalPlan) -> Option<String> {
    match plan {
        LogicalPlan::TableScan(scan) => Some(scan.table_name.table().to_string()),
        LogicalPlan::SubqueryAlias(alias) => Some(alias.alias.table().to_string()),
        _ => {
            for child in plan.inputs() {
                if let Some(name) = extract_table_name_from_plan(child) {
                    return Some(name);
                }
            }
            None
        }
    }
}

/// Convert a DataFusion JoinType to our JoinKind enum.
pub fn convert_join_type(jt: JoinType) -> JoinKind {
    match jt {
        JoinType::Inner => JoinKind::Inner,
        JoinType::Left => JoinKind::Left,
        JoinType::Right => JoinKind::Right,
        JoinType::Full => JoinKind::Full,
        JoinType::LeftSemi => JoinKind::LeftSemi,
        JoinType::RightSemi => JoinKind::RightSemi,
        JoinType::LeftAnti => JoinKind::LeftAnti,
        JoinType::RightAnti => JoinKind::RightAnti,
        // LeftMark is a specialized join type; treat as LeftSemi
        _ => JoinKind::Inner,
    }
}

/// Try to extract a literal usize from an Expr (for LIMIT/OFFSET).
fn extract_literal_usize(expr: &datafusion::logical_expr::Expr) -> Option<usize> {
    use datafusion::common::ScalarValue;

    if let datafusion::logical_expr::Expr::Literal(scalar) = expr {
        match scalar {
            ScalarValue::Int8(Some(v)) => Some(*v as usize),
            ScalarValue::Int16(Some(v)) => Some(*v as usize),
            ScalarValue::Int32(Some(v)) => Some(*v as usize),
            ScalarValue::Int64(Some(v)) => Some(*v as usize),
            ScalarValue::UInt8(Some(v)) => Some(*v as usize),
            ScalarValue::UInt16(Some(v)) => Some(*v as usize),
            ScalarValue::UInt32(Some(v)) => Some(*v as usize),
            ScalarValue::UInt64(Some(v)) => Some(*v as usize),
            _ => None,
        }
    } else {
        None
    }
}

/// Count the total number of nodes in a plan tree.
pub fn count_plan_nodes(plan: &LogicalPlan) -> usize {
    let mut count = 1;
    for child in plan.inputs() {
        count += count_plan_nodes(child);
    }
    count
}
