//! Lineage tracking engine.
//!
//! Traces column-level data flow through SQL queries by walking
//! the sqlparser AST. Builds a lineage graph showing which source
//! columns feed into which output columns.

use std::collections::{HashMap, HashSet};

use super::types::*;
use crate::schema::ddl::extract_table_name;
use crate::schema::types::SchemaDefinition;

use sqlparser::ast::{
    Expr, Query, Select, SelectItem, SetExpr, Statement, TableFactor, TableObject, TableWithJoins,
};
use sqlparser::dialect::GenericDialect;
use sqlparser::parser::Parser;

/// Track data lineage across multiple SQL queries.
///
/// For each query, walks the AST to extract source/target tables and
/// column-level lineage edges. Computes a dependency order (topological sort)
/// and impact map showing how changes to source columns propagate.
pub fn track_lineage(queries: &[&str], schema: &SchemaDefinition) -> Result<LineageResult, String> {
    let dialect = GenericDialect {};
    let mut all_query_lineages = Vec::new();

    for (i, sql) in queries.iter().enumerate() {
        let statements = Parser::parse_sql(&dialect, sql)
            .map_err(|e| format!("Failed to parse query {}: {e}", i + 1))?;

        for stmt in &statements {
            let lineage = analyze_statement(stmt, schema);
            all_query_lineages.push(lineage);
        }
    }

    // Build cross-query dependencies: if query N writes to table X and
    // query M reads from table X, link them
    let dependency_order = compute_dependency_order(&all_query_lineages);
    let impact_map = compute_impact_map(&all_query_lineages);

    Ok(LineageResult {
        queries: all_query_lineages,
        dependency_order,
        impact_map,
    })
}

/// Analyze a single SQL statement for lineage information.
fn analyze_statement(stmt: &Statement, schema: &SchemaDefinition) -> QueryLineage {
    match stmt {
        Statement::Query(query) => analyze_query(query, None, schema),

        Statement::Insert(insert) => {
            let target_table = match &insert.table {
                TableObject::TableName(name) => extract_table_name(name),
                TableObject::TableFunction(_) => "<function>".to_string(),
            };

            let mut lineage = if let Some(source) = &insert.source {
                analyze_query(source, Some(&target_table), schema)
            } else {
                QueryLineage {
                    source_tables: Vec::new(),
                    target_tables: Vec::new(),
                    edges: Vec::new(),
                    output_columns: Vec::new(),
                }
            };

            if !lineage.target_tables.contains(&target_table) {
                lineage.target_tables.push(target_table.clone());
            }

            // Map INSERT columns to target table columns
            if !insert.columns.is_empty() {
                let insert_cols: Vec<String> =
                    insert.columns.iter().map(|c| c.value.clone()).collect();

                // Update edges to point to the target table columns
                let mut updated_edges = Vec::new();
                for (idx, edge) in lineage.edges.iter().enumerate() {
                    if idx < insert_cols.len() {
                        updated_edges.push(LineageEdge {
                            source: edge.source.clone(),
                            target: ColumnRef {
                                table: target_table.clone(),
                                column: insert_cols[idx].clone(),
                            },
                            transform_type: edge.transform_type.clone(),
                        });
                    } else {
                        updated_edges.push(edge.clone());
                    }
                }
                lineage.edges = updated_edges;
            }

            lineage
        }

        Statement::CreateTable(create_table) => {
            if let Some(query) = &create_table.query {
                let target_table = extract_table_name(&create_table.name);
                let mut lineage = analyze_query(query, Some(&target_table), schema);
                if !lineage.target_tables.contains(&target_table) {
                    lineage.target_tables.push(target_table);
                }
                lineage
            } else {
                QueryLineage {
                    source_tables: Vec::new(),
                    target_tables: vec![extract_table_name(&create_table.name)],
                    edges: Vec::new(),
                    output_columns: Vec::new(),
                }
            }
        }

        _ => QueryLineage {
            source_tables: Vec::new(),
            target_tables: Vec::new(),
            edges: Vec::new(),
            output_columns: Vec::new(),
        },
    }
}

/// Analyze a Query (SELECT) for lineage.
fn analyze_query(
    query: &Query,
    target_table: Option<&str>,
    schema: &SchemaDefinition,
) -> QueryLineage {
    let mut source_tables = Vec::new();
    let mut edges = Vec::new();
    let mut output_columns = Vec::new();

    // Table aliases: alias -> real table name
    let mut table_aliases: HashMap<String, String> = HashMap::new();

    // Walk the query body
    extract_from_set_expr(
        &query.body,
        &mut source_tables,
        &mut edges,
        &mut output_columns,
        &mut table_aliases,
        target_table,
        schema,
    );

    // Deduplicate source tables
    let mut seen = HashSet::new();
    source_tables.retain(|t| seen.insert(t.clone()));

    let target_tables = target_table
        .map(|t| vec![t.to_string()])
        .unwrap_or_default();

    QueryLineage {
        source_tables,
        target_tables,
        edges,
        output_columns,
    }
}

/// Extract lineage from a SetExpr (the body of a query).
fn extract_from_set_expr(
    set_expr: &SetExpr,
    source_tables: &mut Vec<String>,
    edges: &mut Vec<LineageEdge>,
    output_columns: &mut Vec<String>,
    table_aliases: &mut HashMap<String, String>,
    target_table: Option<&str>,
    schema: &SchemaDefinition,
) {
    match set_expr {
        SetExpr::Select(select) => {
            extract_from_select(
                select,
                source_tables,
                edges,
                output_columns,
                table_aliases,
                target_table,
                schema,
            );
        }
        SetExpr::Query(query) => {
            extract_from_set_expr(
                &query.body,
                source_tables,
                edges,
                output_columns,
                table_aliases,
                target_table,
                schema,
            );
        }
        SetExpr::SetOperation { left, right, .. } => {
            extract_from_set_expr(
                left,
                source_tables,
                edges,
                output_columns,
                table_aliases,
                target_table,
                schema,
            );
            extract_from_set_expr(
                right,
                source_tables,
                edges,
                &mut Vec::new(),
                table_aliases,
                target_table,
                schema,
            );
        }
        _ => {}
    }
}

/// Extract lineage from a SELECT statement.
fn extract_from_select(
    select: &Select,
    source_tables: &mut Vec<String>,
    edges: &mut Vec<LineageEdge>,
    output_columns: &mut Vec<String>,
    table_aliases: &mut HashMap<String, String>,
    target_table: Option<&str>,
    schema: &SchemaDefinition,
) {
    // 1. Extract source tables from FROM clause
    for table_with_joins in &select.from {
        extract_tables_from_table_with_joins(table_with_joins, source_tables, table_aliases);
    }

    // 2. Extract column-level lineage from projection
    let output_target = target_table.unwrap_or("_output");
    let mut col_idx = 0usize;
    for item in &select.projection {
        match item {
            SelectItem::UnnamedExpr(expr) => {
                let output_name = expr_to_column_name(expr).unwrap_or_else(|| {
                    col_idx += 1;
                    format!("col{col_idx}")
                });
                output_columns.push(output_name.clone());

                let source_refs = extract_column_refs_from_expr(expr, table_aliases);
                let transform = classify_expr_transform(expr);
                for source_ref in source_refs {
                    edges.push(LineageEdge {
                        source: source_ref,
                        target: ColumnRef {
                            table: output_target.to_string(),
                            column: output_name.clone(),
                        },
                        transform_type: transform.clone(),
                    });
                }
            }
            SelectItem::ExprWithAlias { expr, alias } => {
                let output_name = alias.value.clone();
                output_columns.push(output_name.clone());

                let source_refs = extract_column_refs_from_expr(expr, table_aliases);
                let transform = classify_expr_transform(expr);
                for source_ref in source_refs {
                    edges.push(LineageEdge {
                        source: source_ref,
                        target: ColumnRef {
                            table: output_target.to_string(),
                            column: output_name.clone(),
                        },
                        transform_type: transform.clone(),
                    });
                }
            }
            SelectItem::QualifiedWildcard(kind, _) => {
                // e.g. t.*
                let qualifier = match kind {
                    sqlparser::ast::SelectItemQualifiedWildcardKind::ObjectName(name) => {
                        Some(extract_table_name(name))
                    }
                    _ => None,
                };

                if let Some(qual) = qualifier {
                    let real_table = table_aliases.get(&qual).cloned().unwrap_or(qual.clone());
                    // Try to resolve columns from schema
                    if let Some(cols) = schema.column_names(&real_table) {
                        for col in cols {
                            output_columns.push(col.to_string());
                            edges.push(LineageEdge {
                                source: ColumnRef {
                                    table: real_table.clone(),
                                    column: col.to_string(),
                                },
                                target: ColumnRef {
                                    table: output_target.to_string(),
                                    column: col.to_string(),
                                },
                                transform_type: "direct".to_string(),
                            });
                        }
                    } else {
                        output_columns.push(format!("{qual}.*"));
                    }
                }
            }
            SelectItem::Wildcard(_) => {
                // Expand * using schema for all source tables
                for tbl in source_tables.iter() {
                    if let Some(cols) = schema.column_names(tbl) {
                        for col in cols {
                            output_columns.push(col.to_string());
                            edges.push(LineageEdge {
                                source: ColumnRef {
                                    table: tbl.clone(),
                                    column: col.to_string(),
                                },
                                target: ColumnRef {
                                    table: output_target.to_string(),
                                    column: col.to_string(),
                                },
                                transform_type: "direct".to_string(),
                            });
                        }
                    }
                }
                if source_tables.is_empty() {
                    output_columns.push("*".to_string());
                }
            }
        }
    }
}

/// Extract source tables from a TableWithJoins node.
fn extract_tables_from_table_with_joins(
    twj: &TableWithJoins,
    source_tables: &mut Vec<String>,
    table_aliases: &mut HashMap<String, String>,
) {
    extract_table_from_factor(&twj.relation, source_tables, table_aliases);
    for join in &twj.joins {
        extract_table_from_factor(&join.relation, source_tables, table_aliases);
    }
}

/// Extract table name from a TableFactor.
fn extract_table_from_factor(
    factor: &TableFactor,
    source_tables: &mut Vec<String>,
    table_aliases: &mut HashMap<String, String>,
) {
    match factor {
        TableFactor::Table { name, alias, .. } => {
            let table_name = extract_table_name(name);
            source_tables.push(table_name.clone());
            if let Some(alias) = alias {
                table_aliases.insert(alias.name.value.clone(), table_name);
            }
        }
        TableFactor::Derived {
            subquery, alias, ..
        } => {
            // For subqueries, extract tables recursively
            extract_tables_from_query(subquery, source_tables, table_aliases);
            if let Some(alias) = alias {
                table_aliases.insert(alias.name.value.clone(), "(subquery)".to_string());
            }
        }
        _ => {}
    }
}

/// Recursively extract table names from a Query.
fn extract_tables_from_query(
    query: &Query,
    source_tables: &mut Vec<String>,
    table_aliases: &mut HashMap<String, String>,
) {
    match query.body.as_ref() {
        SetExpr::Select(select) => {
            for twj in &select.from {
                extract_tables_from_table_with_joins(twj, source_tables, table_aliases);
            }
        }
        SetExpr::Query(q) => {
            extract_tables_from_query(q, source_tables, table_aliases);
        }
        SetExpr::SetOperation { left, right, .. } => {
            extract_tables_from_set_expr(left, source_tables, table_aliases);
            extract_tables_from_set_expr(right, source_tables, table_aliases);
        }
        _ => {}
    }
}

/// Recursively extract table names from a SetExpr.
fn extract_tables_from_set_expr(
    set_expr: &SetExpr,
    source_tables: &mut Vec<String>,
    table_aliases: &mut HashMap<String, String>,
) {
    match set_expr {
        SetExpr::Select(select) => {
            for twj in &select.from {
                extract_tables_from_table_with_joins(twj, source_tables, table_aliases);
            }
        }
        SetExpr::Query(q) => {
            extract_tables_from_query(q, source_tables, table_aliases);
        }
        SetExpr::SetOperation { left, right, .. } => {
            extract_tables_from_set_expr(left, source_tables, table_aliases);
            extract_tables_from_set_expr(right, source_tables, table_aliases);
        }
        _ => {}
    }
}

/// Extract column references from an expression.
fn extract_column_refs_from_expr(
    expr: &Expr,
    table_aliases: &HashMap<String, String>,
) -> Vec<ColumnRef> {
    let mut refs = Vec::new();
    walk_expr_for_columns(expr, table_aliases, &mut refs);
    refs
}

/// Recursively walk an expression to find column references.
fn walk_expr_for_columns(
    expr: &Expr,
    table_aliases: &HashMap<String, String>,
    refs: &mut Vec<ColumnRef>,
) {
    match expr {
        Expr::Identifier(ident) => {
            // Unqualified column reference — table is unknown
            refs.push(ColumnRef {
                table: "_unknown".to_string(),
                column: ident.value.clone(),
            });
        }
        Expr::CompoundIdentifier(parts) => {
            if parts.len() == 2 {
                let table_or_alias = &parts[0].value;
                let column = &parts[1].value;
                let real_table = table_aliases
                    .get(table_or_alias)
                    .cloned()
                    .unwrap_or_else(|| table_or_alias.clone());
                refs.push(ColumnRef {
                    table: real_table,
                    column: column.clone(),
                });
            } else if parts.len() >= 3 {
                // schema.table.column or catalog.schema.table.column
                let column = parts.last().map(|p| p.value.clone()).unwrap_or_default();
                let table = parts
                    .get(parts.len() - 2)
                    .map(|p| p.value.clone())
                    .unwrap_or_default();
                let real_table = table_aliases.get(&table).cloned().unwrap_or(table);
                refs.push(ColumnRef {
                    table: real_table,
                    column,
                });
            }
        }
        Expr::Function(func) => {
            if let sqlparser::ast::FunctionArguments::List(arg_list) = &func.args {
                for arg in &arg_list.args {
                    if let sqlparser::ast::FunctionArg::Unnamed(
                        sqlparser::ast::FunctionArgExpr::Expr(inner),
                    ) = arg
                    {
                        walk_expr_for_columns(inner, table_aliases, refs);
                    }
                }
            }
        }
        Expr::BinaryOp { left, right, .. } => {
            walk_expr_for_columns(left, table_aliases, refs);
            walk_expr_for_columns(right, table_aliases, refs);
        }
        Expr::UnaryOp { expr: inner, .. } => {
            walk_expr_for_columns(inner, table_aliases, refs);
        }
        Expr::Nested(inner) => {
            walk_expr_for_columns(inner, table_aliases, refs);
        }
        Expr::Case {
            operand,
            conditions,
            else_result,
        } => {
            if let Some(op) = operand {
                walk_expr_for_columns(op, table_aliases, refs);
            }
            for case_when in conditions {
                walk_expr_for_columns(&case_when.condition, table_aliases, refs);
                walk_expr_for_columns(&case_when.result, table_aliases, refs);
            }
            if let Some(else_r) = else_result {
                walk_expr_for_columns(else_r, table_aliases, refs);
            }
        }
        Expr::Cast { expr: inner, .. } => {
            walk_expr_for_columns(inner, table_aliases, refs);
        }
        Expr::Subquery(query) => {
            // We don't track into subquery expressions for column refs
            let _ = query;
        }
        _ => {
            // For other expression types, we don't track deeper
        }
    }
}

/// Classify what kind of transform an expression represents.
fn classify_expr_transform(expr: &Expr) -> String {
    match expr {
        Expr::Identifier(_) | Expr::CompoundIdentifier(_) => "direct".to_string(),
        Expr::Function(func) => {
            let name = func.name.to_string().to_uppercase();
            match name.as_str() {
                "SUM" | "AVG" | "COUNT" | "MIN" | "MAX" | "STDDEV" | "VARIANCE" => {
                    "aggregate".to_string()
                }
                _ => "transform".to_string(),
            }
        }
        Expr::BinaryOp { .. } | Expr::UnaryOp { .. } | Expr::Cast { .. } => "transform".to_string(),
        Expr::Case { .. } => "transform".to_string(),
        _ => "transform".to_string(),
    }
}

/// Try to extract a simple column name from an expression.
fn expr_to_column_name(expr: &Expr) -> Option<String> {
    match expr {
        Expr::Identifier(ident) => Some(ident.value.clone()),
        Expr::CompoundIdentifier(parts) => parts.last().map(|p| p.value.clone()),
        _ => None,
    }
}

/// Compute topological dependency order based on which queries write
/// to tables that other queries read from.
fn compute_dependency_order(lineages: &[QueryLineage]) -> Vec<String> {
    // Collect all tables involved
    let mut all_tables = HashSet::new();

    for lineage in lineages {
        for t in &lineage.target_tables {
            all_tables.insert(t.clone());
        }
        for t in &lineage.source_tables {
            all_tables.insert(t.clone());
        }
    }

    // Build dependency graph: table A depends on table B if B is written
    // in a query and A is read from a later query that depends on B
    let mut deps: HashMap<String, HashSet<String>> = HashMap::new();
    for table in &all_tables {
        deps.insert(table.clone(), HashSet::new());
    }

    for lineage in lineages {
        for source in &lineage.source_tables {
            for target in &lineage.target_tables {
                if source != target {
                    // target depends on source
                    deps.entry(target.clone())
                        .or_default()
                        .insert(source.clone());
                }
            }
        }
    }

    // DFS-based topological sort
    let mut result = Vec::new();
    let mut visited = HashSet::new();
    let mut in_progress = HashSet::new();

    fn visit(
        node: &str,
        deps: &HashMap<String, HashSet<String>>,
        visited: &mut HashSet<String>,
        in_progress: &mut HashSet<String>,
        result: &mut Vec<String>,
    ) {
        if visited.contains(node) || in_progress.contains(node) {
            return;
        }
        in_progress.insert(node.to_string());
        if let Some(node_deps) = deps.get(node) {
            for dep in node_deps {
                visit(dep, deps, visited, in_progress, result);
            }
        }
        in_progress.remove(node);
        visited.insert(node.to_string());
        result.push(node.to_string());
    }

    // Sort table names for deterministic ordering
    let mut sorted_tables: Vec<&String> = all_tables.iter().collect();
    sorted_tables.sort();

    for table in sorted_tables {
        visit(table, &deps, &mut visited, &mut in_progress, &mut result);
    }

    result
}

/// Compute impact map: for each source column, find all downstream
/// columns that would be affected by a change.
fn compute_impact_map(lineages: &[QueryLineage]) -> Vec<ImpactEntry> {
    // Collect all edges across all queries
    let mut all_edges: Vec<&LineageEdge> = Vec::new();
    for lineage in lineages {
        for edge in &lineage.edges {
            all_edges.push(edge);
        }
    }

    // Build source -> targets map
    let mut source_to_targets: HashMap<ColumnRef, HashSet<ColumnRef>> = HashMap::new();
    for edge in &all_edges {
        source_to_targets
            .entry(edge.source.clone())
            .or_default()
            .insert(edge.target.clone());
    }

    // Compute transitive closure for each source
    let mut impact_map = Vec::new();
    let mut sources: Vec<&ColumnRef> = source_to_targets.keys().collect();
    sources.sort_by(|a, b| (&a.table, &a.column).cmp(&(&b.table, &b.column)));

    for source in sources {
        let mut affected = HashSet::new();
        let mut queue: Vec<ColumnRef> = vec![source.clone()];
        let mut visited = HashSet::new();

        while let Some(current) = queue.pop() {
            if !visited.insert(current.clone()) {
                continue;
            }
            if let Some(targets) = source_to_targets.get(&current) {
                for target in targets {
                    affected.insert(target.clone());
                    queue.push(target.clone());
                }
            }
        }

        if !affected.is_empty() {
            let mut affected_vec: Vec<ColumnRef> = affected.into_iter().collect();
            affected_vec.sort_by(|a, b| (&a.table, &a.column).cmp(&(&b.table, &b.column)));
            impact_map.push(ImpactEntry {
                source: source.clone(),
                affected: affected_vec,
            });
        }
    }

    impact_map
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::schema::types::{
        ColumnDef as SchemaColumnDef, SchemaDefinition, SqlDialect, TableDef,
    };
    use std::collections::HashMap;

    fn test_schema() -> SchemaDefinition {
        let mut tables = HashMap::new();
        tables.insert(
            "users".to_string(),
            TableDef {
                description: None,
                database: None,
                schema: None,
                columns: vec![
                    SchemaColumnDef {
                        name: "id".to_string(),
                        data_type: "INT".to_string(),
                        nullable: false,
                        description: None,
                        tags: vec![],
                        synonyms: vec![],

                        statistics: None,
                    },
                    SchemaColumnDef {
                        name: "name".to_string(),
                        data_type: "VARCHAR".to_string(),
                        nullable: true,
                        description: None,
                        tags: vec![],
                        synonyms: vec![],

                        statistics: None,
                    },
                    SchemaColumnDef {
                        name: "email".to_string(),
                        data_type: "VARCHAR".to_string(),
                        nullable: true,
                        description: None,
                        tags: vec![],
                        synonyms: vec![],

                        statistics: None,
                    },
                ],
                primary_key: Some(vec!["id".to_string()]),
                foreign_keys: vec![],
                statistics: None,
                clustering_keys: vec![],
            },
        );
        tables.insert(
            "orders".to_string(),
            TableDef {
                description: None,
                database: None,
                schema: None,
                columns: vec![
                    SchemaColumnDef {
                        name: "id".to_string(),
                        data_type: "INT".to_string(),
                        nullable: false,
                        description: None,
                        tags: vec![],
                        synonyms: vec![],

                        statistics: None,
                    },
                    SchemaColumnDef {
                        name: "user_id".to_string(),
                        data_type: "INT".to_string(),
                        nullable: false,
                        description: None,
                        tags: vec![],
                        synonyms: vec![],

                        statistics: None,
                    },
                    SchemaColumnDef {
                        name: "amount".to_string(),
                        data_type: "DECIMAL(10,2)".to_string(),
                        nullable: true,
                        description: None,
                        tags: vec![],
                        synonyms: vec![],

                        statistics: None,
                    },
                ],
                primary_key: Some(vec!["id".to_string()]),
                foreign_keys: vec![],
                statistics: None,
                clustering_keys: vec![],
            },
        );
        SchemaDefinition {
            version: "1".to_string(),
            dialect: SqlDialect::Generic,
            tables,
            database: None,
            schema_name: None,
            glossary: None,
        }
    }

    fn empty_schema() -> SchemaDefinition {
        SchemaDefinition {
            version: "1".to_string(),
            dialect: SqlDialect::Generic,
            tables: HashMap::new(),
            database: None,
            schema_name: None,
            glossary: None,
        }
    }

    #[test]
    fn test_simple_select() {
        let result =
            track_lineage(&["SELECT id, name FROM users"], &test_schema()).expect("should succeed");
        assert_eq!(result.queries.len(), 1);
        assert_eq!(result.queries[0].source_tables, vec!["users"]);
        assert!(result.queries[0].target_tables.is_empty());
        assert_eq!(result.queries[0].output_columns, vec!["id", "name"]);
    }

    #[test]
    fn test_select_with_alias() {
        let result = track_lineage(
            &["SELECT u.id, u.name AS user_name FROM users u"],
            &test_schema(),
        )
        .expect("should succeed");
        assert_eq!(result.queries[0].source_tables, vec!["users"]);
        assert_eq!(result.queries[0].output_columns, vec!["id", "user_name"]);
        // Check edges resolve alias
        assert!(
            result.queries[0]
                .edges
                .iter()
                .any(|e| e.source.table == "users" && e.source.column == "id")
        );
        assert!(
            result.queries[0]
                .edges
                .iter()
                .any(|e| e.source.table == "users" && e.source.column == "name")
        );
    }

    #[test]
    fn test_join_query() {
        let result = track_lineage(
            &["SELECT u.name, o.amount FROM users u JOIN orders o ON u.id = o.user_id"],
            &test_schema(),
        )
        .expect("should succeed");
        let q = &result.queries[0];
        assert!(q.source_tables.contains(&"users".to_string()));
        assert!(q.source_tables.contains(&"orders".to_string()));
        assert_eq!(q.output_columns, vec!["name", "amount"]);
    }

    #[test]
    fn test_aggregate_transform() {
        let result = track_lineage(
            &["SELECT user_id, SUM(amount) AS total FROM orders GROUP BY user_id"],
            &test_schema(),
        )
        .expect("should succeed");
        let q = &result.queries[0];
        // SUM(amount) should be classified as "aggregate"
        assert!(
            q.edges
                .iter()
                .any(|e| e.transform_type == "aggregate" && e.target.column == "total")
        );
        // user_id should be "direct"
        assert!(
            q.edges
                .iter()
                .any(|e| e.transform_type == "direct" && e.target.column == "user_id")
        );
    }

    #[test]
    fn test_insert_into_lineage() {
        let result = track_lineage(
            &["INSERT INTO orders (id, user_id, amount) SELECT id, id, 100 FROM users"],
            &test_schema(),
        )
        .expect("should succeed");
        let q = &result.queries[0];
        assert!(q.source_tables.contains(&"users".to_string()));
        assert!(q.target_tables.contains(&"orders".to_string()));
    }

    #[test]
    fn test_multi_query_dependency() {
        let queries = &[
            "SELECT id, name FROM users",
            "INSERT INTO orders (user_id) SELECT id FROM users",
        ];
        let result = track_lineage(queries, &test_schema()).expect("should succeed");
        assert_eq!(result.queries.len(), 2);
        // dependency_order should include users and orders
        assert!(!result.dependency_order.is_empty());
    }

    #[test]
    fn test_wildcard_expansion() {
        let result =
            track_lineage(&["SELECT * FROM users"], &test_schema()).expect("should succeed");
        let q = &result.queries[0];
        // Wildcard should expand to all columns in users
        assert!(q.output_columns.contains(&"id".to_string()));
        assert!(q.output_columns.contains(&"name".to_string()));
        assert!(q.output_columns.contains(&"email".to_string()));
    }

    #[test]
    fn test_create_table_as_select() {
        let result = track_lineage(
            &["CREATE TABLE user_summary AS SELECT id, name FROM users"],
            &test_schema(),
        )
        .expect("should succeed");
        let q = &result.queries[0];
        assert!(q.source_tables.contains(&"users".to_string()));
        assert!(q.target_tables.contains(&"user_summary".to_string()));
    }

    #[test]
    fn test_impact_map() {
        let result = track_lineage(
            &["SELECT u.id, u.name AS user_name FROM users u"],
            &test_schema(),
        )
        .expect("should succeed");
        // The impact map should show that users.id affects _output.id
        assert!(!result.impact_map.is_empty());
        assert!(
            result
                .impact_map
                .iter()
                .any(|entry| entry.source.table == "users" && entry.source.column == "id")
        );
    }

    #[test]
    fn test_parse_error() {
        let result = track_lineage(&["THIS IS NOT SQL AT ALL"], &empty_schema());
        assert!(result.is_err());
    }

    #[test]
    fn test_expression_transform() {
        let result = track_lineage(&["SELECT id + 1 AS next_id FROM users"], &test_schema())
            .expect("should succeed");
        let q = &result.queries[0];
        assert!(
            q.edges
                .iter()
                .any(|e| e.transform_type == "transform" && e.target.column == "next_id")
        );
    }
}
