//! Column and table lineage extraction from DataFusion LogicalPlan.
//!
//! Walks the plan tree to determine which tables are read, which columns
//! are used, and how tables are joined together.

use super::plan_walker::{convert_join_type, extract_table_name_from_plan};
use super::types::*;
use datafusion::logical_expr::LogicalPlan;
use std::collections::HashSet;

/// Extract lineage information from a LogicalPlan.
///
/// Walks the entire plan tree collecting:
/// - Tables referenced (from TableScan nodes)
/// - Columns read (from TableScan projected schemas)
/// - Join graph (from Join nodes)
/// - Output columns (from the root plan's schema)
pub fn extract_lineage(plan: &LogicalPlan) -> QueryLineage {
    let mut tables = Vec::new();
    let mut columns_read = Vec::new();
    let mut join_graph = Vec::new();
    let mut seen_tables = HashSet::new();
    let mut seen_columns = HashSet::new();

    collect_lineage(
        plan,
        &mut tables,
        &mut columns_read,
        &mut join_graph,
        &mut seen_tables,
        &mut seen_columns,
    );

    // Extract output columns from the root plan's schema
    let columns_output = extract_output_columns(plan);

    QueryLineage {
        tables,
        columns_read,
        columns_output,
        join_graph,
    }
}

fn collect_lineage(
    plan: &LogicalPlan,
    tables: &mut Vec<String>,
    columns_read: &mut Vec<ColumnRef>,
    join_graph: &mut Vec<JoinEdge>,
    seen_tables: &mut HashSet<String>,
    seen_columns: &mut HashSet<(String, String)>,
) {
    match plan {
        LogicalPlan::TableScan(scan) => {
            let table_name = scan.table_name.table().to_string();
            if seen_tables.insert(table_name.clone()) {
                tables.push(table_name.clone());
            }
            // Collect columns from the projected schema
            for field in scan.projected_schema.fields() {
                let col_key = (table_name.clone(), field.name().to_string());
                if seen_columns.insert(col_key) {
                    columns_read.push(ColumnRef {
                        table: table_name.clone(),
                        column: field.name().to_string(),
                    });
                }
            }
        }

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

            let on_columns: Vec<(String, String)> = join
                .on
                .iter()
                .map(|(l, r)| (format!("{l}"), format!("{r}")))
                .collect();

            join_graph.push(JoinEdge {
                left_table,
                right_table,
                join_type: convert_join_type(join.join_type),
                on_columns,
            });
        }

        _ => {}
    }

    // Recurse into children
    for child in plan.inputs() {
        collect_lineage(
            child,
            tables,
            columns_read,
            join_graph,
            seen_tables,
            seen_columns,
        );
    }
}

/// Extract output columns from the root plan's output schema.
fn extract_output_columns(plan: &LogicalPlan) -> Vec<OutputColumn> {
    let schema = plan.schema();
    schema
        .fields()
        .iter()
        .map(|field| {
            let alias = field.name().to_string();

            // Try to determine the source of this output column
            let source = if let Some(qualifier) = schema
                .iter()
                .find(|(_, f)| f.name() == field.name())
                .and_then(|(q, _)| q)
            {
                ColumnSource::DirectRef {
                    table: qualifier.table().to_string(),
                    column: field.name().to_string(),
                }
            } else {
                ColumnSource::Expression {
                    expr: field.name().to_string(),
                }
            };

            OutputColumn { alias, source }
        })
        .collect()
}
