//! The query equivalence engine.
//!
//! Compares two SQL queries heuristically by analyzing their DataFusion
//! logical plans. Also provides schema regression checking to detect
//! when schema changes break existing queries.

use super::types::*;
use crate::schema::types::SchemaDefinition;
use crate::validator::engine::validate_with_plan;
use datafusion::logical_expr::LogicalPlan;
use std::collections::BTreeSet;

/// Check whether two SQL queries are semantically equivalent.
///
/// Uses heuristic comparison of their DataFusion logical plans.
/// Compares output columns, referenced tables, filter predicates,
/// join structure, and aggregation patterns.
///
/// # Example
///
/// ```rust,no_run
/// use altimate_core::equivalence::engine::check_equivalence;
/// use altimate_core::SchemaDefinition;
///
/// # async fn example() {
/// let schema = SchemaDefinition::from_yaml_file("schema.yaml").unwrap();
/// let result = check_equivalence(
///     "SELECT id, name FROM users WHERE age > 25",
///     "SELECT id, name FROM users WHERE age > 25 ORDER BY id",
///     &schema,
/// ).await;
/// println!("Equivalent: {} (confidence: {})", result.equivalent, result.confidence);
/// # }
/// ```
pub async fn check_equivalence(
    sql_a: &str,
    sql_b: &str,
    schema: &SchemaDefinition,
) -> EquivalenceResult {
    // 1. Validate and get logical plans for both queries
    let (result_a, plan_a) = validate_with_plan(sql_a, schema).await;
    let (result_b, plan_b) = validate_with_plan(sql_b, schema).await;

    let mut errors = Vec::new();
    if !result_a.valid {
        errors.push(format!(
            "Query A failed validation: {}",
            result_a
                .errors
                .iter()
                .map(|e| e.message.as_str())
                .collect::<Vec<_>>()
                .join("; ")
        ));
    }
    if !result_b.valid {
        errors.push(format!(
            "Query B failed validation: {}",
            result_b
                .errors
                .iter()
                .map(|e| e.message.as_str())
                .collect::<Vec<_>>()
                .join("; ")
        ));
    }
    if !errors.is_empty() {
        return EquivalenceResult::invalid(errors);
    }

    let plan_a = plan_a.expect("plan should be Some when valid");
    let plan_b = plan_b.expect("plan should be Some when valid");

    // 2. Compare the plans
    let mut differences = Vec::new();

    // Compare output columns
    let output_compatible = compare_output_columns(&plan_a, &plan_b, &mut differences);

    // Compare referenced tables
    compare_tables(&plan_a, &plan_b, &mut differences);

    // Compare filter predicates
    compare_filters(&plan_a, &plan_b, &mut differences);

    // Compare join structure
    compare_joins(&plan_a, &plan_b, &mut differences);

    // Compare aggregation structure
    compare_aggregation(&plan_a, &plan_b, &mut differences);

    // Compare sort/limit
    compare_sort_limit(&plan_a, &plan_b, &mut differences);

    // 3. Build result
    EquivalenceResult::with_diffs(differences, output_compatible)
}

/// Check if a query breaks under a new schema version.
///
/// Validates the query against both the old and new schema definitions.
/// Reports which specific schema changes caused the query to break.
///
/// # Example
///
/// ```rust,no_run
/// use altimate_core::equivalence::engine::check_regression;
/// use altimate_core::SchemaDefinition;
///
/// # async fn example() {
/// let old_schema = SchemaDefinition::from_yaml_file("schema_v1.yaml").unwrap();
/// let new_schema = SchemaDefinition::from_yaml_file("schema_v2.yaml").unwrap();
/// let result = check_regression(
///     "SELECT id, email FROM users",
///     &old_schema,
///     &new_schema,
/// ).await;
/// if !result.passes {
///     for change in &result.breaking_changes {
///         println!("Breaking: {} — {}", change.change_type, change.detail);
///     }
/// }
/// # }
/// ```
pub async fn check_regression(
    sql: &str,
    schema_old: &SchemaDefinition,
    schema_new: &SchemaDefinition,
) -> RegressionResult {
    // 1. Validate against old schema
    let (old_result, _old_plan) = validate_with_plan(sql, schema_old).await;
    if !old_result.valid {
        let errors = old_result
            .errors
            .iter()
            .map(|e| e.message.clone())
            .collect();
        return RegressionResult::already_invalid(errors);
    }

    // 2. Validate against new schema
    let (new_result, _new_plan) = validate_with_plan(sql, schema_new).await;

    if new_result.valid {
        // Query still works — check for warnings about schema changes
        let warnings = detect_schema_warnings(sql, schema_old, schema_new);
        return RegressionResult {
            passes: true,
            breaking_changes: vec![],
            warnings,
            validation_errors: vec![],
        };
    }

    // 3. Query broke — figure out what changed
    let breaking_changes = detect_breaking_changes(sql, schema_old, schema_new, &new_result);

    RegressionResult {
        passes: false,
        breaking_changes,
        warnings: vec![],
        validation_errors: new_result
            .errors
            .iter()
            .map(|e| e.message.clone())
            .collect(),
    }
}

// ─── Plan comparison functions ──────────────────────────────────────

/// Compare output columns (names and types) of two plans.
/// Returns whether outputs are compatible.
fn compare_output_columns(
    plan_a: &LogicalPlan,
    plan_b: &LogicalPlan,
    diffs: &mut Vec<EquivalenceDiff>,
) -> bool {
    let schema_a = plan_a.schema();
    let schema_b = plan_b.schema();

    let names_a: Vec<String> = schema_a.fields().iter().map(|f| f.name().clone()).collect();
    let names_b: Vec<String> = schema_b.fields().iter().map(|f| f.name().clone()).collect();

    let types_a: Vec<String> = schema_a
        .fields()
        .iter()
        .map(|f| format!("{:?}", f.data_type()))
        .collect();
    let types_b: Vec<String> = schema_b
        .fields()
        .iter()
        .map(|f| format!("{:?}", f.data_type()))
        .collect();

    let mut compatible = true;

    if names_a.len() != names_b.len() {
        diffs.push(EquivalenceDiff {
            aspect: "columns".to_string(),
            description: format!(
                "Different number of output columns: {} vs {}",
                names_a.len(),
                names_b.len()
            ),
            severity: DiffSeverity::Structural,
        });
        compatible = false;
    } else {
        // Compare column names (ignoring case)
        let name_diffs: Vec<_> = names_a
            .iter()
            .zip(names_b.iter())
            .filter(|(a, b)| a.to_lowercase() != b.to_lowercase())
            .collect();
        if !name_diffs.is_empty() {
            let desc = name_diffs
                .iter()
                .map(|(a, b)| format!("'{a}' vs '{b}'"))
                .collect::<Vec<_>>()
                .join(", ");
            diffs.push(EquivalenceDiff {
                aspect: "columns".to_string(),
                description: format!("Different output column names: {desc}"),
                severity: DiffSeverity::Minor,
            });
        }

        // Compare column types
        let type_diffs: Vec<_> = types_a
            .iter()
            .zip(types_b.iter())
            .enumerate()
            .filter(|(_, (a, b))| a != b)
            .collect();
        if !type_diffs.is_empty() {
            for (i, (ta, tb)) in type_diffs {
                let col_name = names_a.get(i).map(|s| s.as_str()).unwrap_or("?");
                diffs.push(EquivalenceDiff {
                    aspect: "columns".to_string(),
                    description: format!("Column '{col_name}' has different types: {ta} vs {tb}"),
                    severity: DiffSeverity::Semantic,
                });
            }
            compatible = false;
        }
    }

    compatible
}

/// Compare referenced tables between two plans.
fn compare_tables(plan_a: &LogicalPlan, plan_b: &LogicalPlan, diffs: &mut Vec<EquivalenceDiff>) {
    let tables_a = collect_tables(plan_a);
    let tables_b = collect_tables(plan_b);

    let only_a: BTreeSet<_> = tables_a.difference(&tables_b).collect();
    let only_b: BTreeSet<_> = tables_b.difference(&tables_a).collect();

    if !only_a.is_empty() {
        diffs.push(EquivalenceDiff {
            aspect: "tables".to_string(),
            description: format!(
                "Tables only in query A: {}",
                only_a
                    .iter()
                    .map(|s| s.as_str())
                    .collect::<Vec<_>>()
                    .join(", ")
            ),
            severity: DiffSeverity::Structural,
        });
    }

    if !only_b.is_empty() {
        diffs.push(EquivalenceDiff {
            aspect: "tables".to_string(),
            description: format!(
                "Tables only in query B: {}",
                only_b
                    .iter()
                    .map(|s| s.as_str())
                    .collect::<Vec<_>>()
                    .join(", ")
            ),
            severity: DiffSeverity::Structural,
        });
    }
}

/// Compare filter predicates between two plans.
fn compare_filters(plan_a: &LogicalPlan, plan_b: &LogicalPlan, diffs: &mut Vec<EquivalenceDiff>) {
    let filters_a = collect_filters(plan_a);
    let filters_b = collect_filters(plan_b);

    // Normalize and compare
    let norm_a: BTreeSet<String> = filters_a.iter().map(|f| normalize_predicate(f)).collect();
    let norm_b: BTreeSet<String> = filters_b.iter().map(|f| normalize_predicate(f)).collect();

    let only_a: BTreeSet<_> = norm_a.difference(&norm_b).collect();
    let only_b: BTreeSet<_> = norm_b.difference(&norm_a).collect();

    if !only_a.is_empty() {
        diffs.push(EquivalenceDiff {
            aspect: "filters".to_string(),
            description: format!(
                "Filters only in query A: {}",
                only_a.into_iter().cloned().collect::<Vec<_>>().join(", ")
            ),
            severity: DiffSeverity::Semantic,
        });
    }

    if !only_b.is_empty() {
        diffs.push(EquivalenceDiff {
            aspect: "filters".to_string(),
            description: format!(
                "Filters only in query B: {}",
                only_b.into_iter().cloned().collect::<Vec<_>>().join(", ")
            ),
            severity: DiffSeverity::Semantic,
        });
    }
}

/// Compare join structure between two plans.
fn compare_joins(plan_a: &LogicalPlan, plan_b: &LogicalPlan, diffs: &mut Vec<EquivalenceDiff>) {
    let joins_a = collect_joins(plan_a);
    let joins_b = collect_joins(plan_b);

    if joins_a.len() != joins_b.len() {
        diffs.push(EquivalenceDiff {
            aspect: "joins".to_string(),
            description: format!(
                "Different number of joins: {} vs {}",
                joins_a.len(),
                joins_b.len()
            ),
            severity: DiffSeverity::Structural,
        });
        return;
    }

    // Compare join types and conditions
    let norm_a: BTreeSet<String> = joins_a.into_iter().collect();
    let norm_b: BTreeSet<String> = joins_b.into_iter().collect();

    let only_a: Vec<_> = norm_a.difference(&norm_b).collect();
    let only_b: Vec<_> = norm_b.difference(&norm_a).collect();

    if !only_a.is_empty() || !only_b.is_empty() {
        let mut desc_parts = Vec::new();
        if !only_a.is_empty() {
            desc_parts.push(format!(
                "Only in A: {}",
                only_a.into_iter().cloned().collect::<Vec<_>>().join("; ")
            ));
        }
        if !only_b.is_empty() {
            desc_parts.push(format!(
                "Only in B: {}",
                only_b.into_iter().cloned().collect::<Vec<_>>().join("; ")
            ));
        }
        diffs.push(EquivalenceDiff {
            aspect: "joins".to_string(),
            description: format!("Different join structure: {}", desc_parts.join(". ")),
            severity: DiffSeverity::Structural,
        });
    }
}

/// Compare aggregation structure between two plans.
fn compare_aggregation(
    plan_a: &LogicalPlan,
    plan_b: &LogicalPlan,
    diffs: &mut Vec<EquivalenceDiff>,
) {
    let agg_a = collect_aggregates(plan_a);
    let agg_b = collect_aggregates(plan_b);

    let has_agg_a = !agg_a.is_empty();
    let has_agg_b = !agg_b.is_empty();

    if has_agg_a != has_agg_b {
        diffs.push(EquivalenceDiff {
            aspect: "aggregation".to_string(),
            description: format!(
                "One query uses aggregation, the other does not (A: {has_agg_a}, B: {has_agg_b})"
            ),
            severity: DiffSeverity::Structural,
        });
        return;
    }

    if !has_agg_a {
        return;
    }

    // Compare aggregation details
    let norm_a: BTreeSet<String> = agg_a.into_iter().collect();
    let norm_b: BTreeSet<String> = agg_b.into_iter().collect();

    if norm_a != norm_b {
        diffs.push(EquivalenceDiff {
            aspect: "aggregation".to_string(),
            description: "Different aggregation structure".to_string(),
            severity: DiffSeverity::Semantic,
        });
    }
}

/// Compare sort and limit between two plans.
fn compare_sort_limit(
    plan_a: &LogicalPlan,
    plan_b: &LogicalPlan,
    diffs: &mut Vec<EquivalenceDiff>,
) {
    let has_sort_a = has_sort(plan_a);
    let has_sort_b = has_sort(plan_b);

    if has_sort_a != has_sort_b {
        diffs.push(EquivalenceDiff {
            aspect: "ordering".to_string(),
            description: format!(
                "Sort differs: query A {} ORDER BY, query B {} ORDER BY",
                if has_sort_a { "has" } else { "no" },
                if has_sort_b { "has" } else { "no" }
            ),
            severity: DiffSeverity::Minor,
        });
    }

    let limit_a = extract_limit(plan_a);
    let limit_b = extract_limit(plan_b);

    if limit_a != limit_b {
        diffs.push(EquivalenceDiff {
            aspect: "limit".to_string(),
            description: format!("Different LIMIT: {limit_a:?} vs {limit_b:?}"),
            severity: DiffSeverity::Semantic,
        });
    }
}

// ─── Schema regression helpers ──────────────────────────────────────

/// Detect warnings about schema changes (non-breaking).
fn detect_schema_warnings(
    _sql: &str,
    old: &SchemaDefinition,
    new: &SchemaDefinition,
) -> Vec<String> {
    let mut warnings = Vec::new();

    // Check for type changes in columns that still exist
    for (table_name, old_table) in &old.tables {
        if let Some(new_table) = new.tables.get(table_name) {
            for old_col in &old_table.columns {
                if let Some(new_col) = new_table.columns.iter().find(|c| c.name == old_col.name) {
                    if old_col.data_type != new_col.data_type {
                        warnings.push(format!(
                            "Column {table_name}.{} changed type from {} to {}",
                            old_col.name, old_col.data_type, new_col.data_type
                        ));
                    }
                    if old_col.nullable && !new_col.nullable {
                        warnings.push(format!(
                            "Column {table_name}.{} changed from nullable to NOT NULL",
                            old_col.name
                        ));
                    }
                }
            }
        }
    }

    warnings
}

/// Detect breaking changes by comparing schemas and error messages.
fn detect_breaking_changes(
    _sql: &str,
    old: &SchemaDefinition,
    new: &SchemaDefinition,
    new_result: &crate::errors::types::ValidationResult,
) -> Vec<BreakingChange> {
    let mut changes = Vec::new();

    // Check for removed tables
    for table_name in old.tables.keys() {
        if !new.tables.contains_key(table_name) {
            changes.push(BreakingChange {
                change_type: "table_removed".to_string(),
                affected_element: table_name.clone(),
                detail: format!("Table '{table_name}' was removed from the schema"),
                suggestion: None,
            });
        }
    }

    // Check for removed or renamed columns
    for (table_name, old_table) in &old.tables {
        if let Some(new_table) = new.tables.get(table_name) {
            for old_col in &old_table.columns {
                let exists_in_new = new_table.columns.iter().any(|c| c.name == old_col.name);
                if !exists_in_new {
                    // Check if it might be renamed (fuzzy match)
                    let possible_rename = new_table
                        .columns
                        .iter()
                        .filter(|c| !old_table.columns.iter().any(|oc| oc.name == c.name))
                        .find(|c| strsim::jaro_winkler(&old_col.name, &c.name) > 0.8);

                    let suggestion = possible_rename.map(|new_col| {
                        format!(
                            "Column may have been renamed to '{}'. Update references: {}.{} -> {}.{}",
                            new_col.name, table_name, old_col.name, table_name, new_col.name
                        )
                    });

                    changes.push(BreakingChange {
                        change_type: "column_removed".to_string(),
                        affected_element: format!("{table_name}.{}", old_col.name),
                        detail: format!(
                            "Column '{}.{}' was removed from the schema",
                            table_name, old_col.name
                        ),
                        suggestion,
                    });
                }
            }

            // Check for type changes that might break the query
            for old_col in &old_table.columns {
                if let Some(new_col) = new_table.columns.iter().find(|c| c.name == old_col.name)
                    && old_col.data_type != new_col.data_type
                {
                    // Check if any error references this column
                    let col_referenced_in_error = new_result
                        .errors
                        .iter()
                        .any(|e| e.message.contains(&old_col.name));

                    if col_referenced_in_error {
                        changes.push(BreakingChange {
                            change_type: "type_changed".to_string(),
                            affected_element: format!("{table_name}.{}", old_col.name),
                            detail: format!(
                                "Column '{}.{}' type changed from {} to {}",
                                table_name, old_col.name, old_col.data_type, new_col.data_type
                            ),
                            suggestion: Some(format!(
                                "Add explicit CAST({}.{} AS {})",
                                table_name, old_col.name, old_col.data_type
                            )),
                        });
                    }
                }
            }
        }
    }

    // If we couldn't identify specific breaking changes from schema diff,
    // use the error messages
    if changes.is_empty() {
        for err in &new_result.errors {
            changes.push(BreakingChange {
                change_type: "unknown".to_string(),
                affected_element: "unknown".to_string(),
                detail: err.message.clone(),
                suggestion: None,
            });
        }
    }

    changes
}

// ─── Plan extraction helpers ────────────────────────────────────────

/// Collect all table names from a plan.
fn collect_tables(plan: &LogicalPlan) -> BTreeSet<String> {
    let mut tables = BTreeSet::new();
    visit_plan(plan, &mut |node| {
        if let LogicalPlan::TableScan(scan) = node {
            tables.insert(scan.table_name.table().to_string());
        }
    });
    tables
}

/// Collect all filter predicates (normalized) from a plan.
fn collect_filters(plan: &LogicalPlan) -> Vec<String> {
    let mut filters = Vec::new();
    visit_plan(plan, &mut |node| {
        if let LogicalPlan::Filter(filter) = node {
            filters.push(format!("{}", filter.predicate));
        }
    });
    filters
}

/// Collect join descriptions from a plan.
fn collect_joins(plan: &LogicalPlan) -> Vec<String> {
    let mut joins = Vec::new();
    visit_plan(plan, &mut |node| {
        if let LogicalPlan::Join(join) = node {
            let join_type = format!("{:?}", join.join_type);
            let conditions: Vec<String> = join
                .on
                .iter()
                .map(|(l, r)| {
                    let mut parts = [format!("{l}"), format!("{r}")];
                    parts.sort();
                    parts.join(" = ")
                })
                .collect();
            let mut desc = format!("{join_type}({})", conditions.join(", "));
            if let Some(ref filter) = join.filter {
                desc.push_str(&format!(" WHERE {filter}"));
            }
            joins.push(desc);
        }
        // In DataFusion 46, cross joins are Join with empty on/filter
    });
    joins
}

/// Collect aggregation descriptions from a plan.
fn collect_aggregates(plan: &LogicalPlan) -> Vec<String> {
    let mut aggs = Vec::new();
    visit_plan(plan, &mut |node| {
        if let LogicalPlan::Aggregate(agg) = node {
            let group_by: Vec<String> = agg.group_expr.iter().map(|e| format!("{e}")).collect();
            let funcs: Vec<String> = agg.aggr_expr.iter().map(|e| format!("{e}")).collect();
            aggs.push(format!(
                "GROUP BY [{}] AGG [{}]",
                group_by.join(", "),
                funcs.join(", ")
            ));
        }
    });
    aggs
}

/// Check if a plan contains a Sort node.
fn has_sort(plan: &LogicalPlan) -> bool {
    let mut found = false;
    visit_plan(plan, &mut |node| {
        if matches!(node, LogicalPlan::Sort(_)) {
            found = true;
        }
    });
    found
}

/// Extract LIMIT value from a plan as a string for comparison.
fn extract_limit(plan: &LogicalPlan) -> Option<String> {
    let mut limit_val = None;
    visit_plan(plan, &mut |node| {
        if let LogicalPlan::Limit(limit) = node
            && let Some(ref fetch) = limit.fetch
        {
            limit_val = Some(format!("{fetch}"));
        }
    });
    limit_val
}

/// Normalize a predicate string for comparison.
fn normalize_predicate(pred: &str) -> String {
    pred.to_lowercase().replace("  ", " ").trim().to_string()
}

/// Visit every node in a LogicalPlan tree.
fn visit_plan<F>(plan: &LogicalPlan, visitor: &mut F)
where
    F: FnMut(&LogicalPlan),
{
    visitor(plan);
    for child in plan.inputs() {
        visit_plan(child, visitor);
    }
}
