//! Comprehensive tests for the schema pruning module.
//!
//! Note: The pruner uses polyglot-sql's extract_tables which may not extract
//! table names from simple SELECT FROM clauses (returns empty for many queries).
//! When no table names are extracted but column names are found, the pruner
//! falls back to matching column names against all schema tables. This means
//! columns like "id" that exist in multiple tables will match all those tables.

use altimate_core::pruner::prune_schema;
use altimate_core::schema::types::{
    ColumnDef, ForeignKeyDef, ForeignKeyRef, SchemaDefinition, SqlDialect, TableDef,
};
use std::collections::HashMap;

// ─── Helpers ─────────────────────────────────────────────────────

fn ecommerce_schema() -> SchemaDefinition {
    SchemaDefinition::from_yaml(include_str!(
        "../../../tests/fixtures/schemas/ecommerce.yaml"
    ))
    .expect("Failed to load ecommerce schema fixture")
}

fn make_schema(tables: Vec<(&str, Vec<&str>, Vec<(&str, &str)>)>) -> SchemaDefinition {
    let mut schema_tables = HashMap::new();
    for (table_name, columns, fks) in tables {
        let col_defs = columns
            .into_iter()
            .map(|name| ColumnDef {
                name: name.to_string(),
                data_type: "INT".to_string(),
                nullable: true,
                description: None,
                tags: vec![],
                synonyms: vec![],

                statistics: None,
            })
            .collect();
        let fk_defs = fks
            .into_iter()
            .map(|(col, ref_table)| ForeignKeyDef {
                columns: vec![col.to_string()],
                references: ForeignKeyRef {
                    table: ref_table.to_string(),
                    columns: vec!["id".to_string()],
                },
            })
            .collect();
        schema_tables.insert(
            table_name.to_string(),
            TableDef {
                description: None,
                database: None,
                schema: None,
                columns: col_defs,
                primary_key: Some(vec!["id".to_string()]),
                foreign_keys: fk_defs,
                statistics: None,

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

// ─── Basic pruning ──────────────────────────────────────────────

#[test]
fn test_prune_does_not_error() {
    let schema = ecommerce_schema();
    let result = prune_schema("SELECT id, email FROM users", &schema);
    assert!(result.is_ok(), "Prune should not error: {:?}", result.err());
}

#[test]
fn test_prune_includes_referenced_table() {
    let schema = ecommerce_schema();
    let result = prune_schema("SELECT id, email FROM users", &schema).expect("should prune");
    // users should always be included (via column match at minimum)
    assert!(result.relevant_tables.contains(&"users".to_string()));
}

#[test]
fn test_prune_preserves_all_columns_for_matched_table() {
    let schema = ecommerce_schema();
    let result = prune_schema("SELECT id FROM users", &schema).expect("should prune");
    // Pruned YAML should include all columns for matched table
    assert!(result.pruned_schema_yaml.contains("email"));
    assert!(result.pruned_schema_yaml.contains("name"));
}

#[test]
fn test_prune_no_tables_returns_empty() {
    let schema = ecommerce_schema();
    let result = prune_schema("SELECT 1 + 1", &schema).expect("should prune");
    assert!(result.relevant_tables.is_empty());
    assert_eq!(result.tables_pruned, schema.tables.len());
}

#[test]
fn test_prune_reports_total_tables() {
    let schema = ecommerce_schema();
    let result = prune_schema("SELECT id FROM users", &schema).expect("should prune");
    assert_eq!(result.total_tables, 4);
}

#[test]
fn test_prune_tables_pruned_plus_relevant_equals_total() {
    let schema = ecommerce_schema();
    let result = prune_schema("SELECT id FROM users", &schema).expect("should prune");
    assert_eq!(
        result.relevant_tables.len() + result.tables_pruned,
        result.total_tables
    );
}

// ─── FK expansion ───────────────────────────────────────────────

#[test]
fn test_prune_includes_fk_referenced_tables() {
    let schema = ecommerce_schema();
    let result = prune_schema("SELECT id FROM orders", &schema).expect("should prune");
    assert!(result.relevant_tables.contains(&"orders".to_string()));
    // orders has FK to users; with column-based matching, users is also included
    // due to shared "id" column
    assert!(result.relevant_tables.contains(&"users".to_string()));
}

#[test]
fn test_prune_includes_nested_fk_tables() {
    let schema = ecommerce_schema();
    let result = prune_schema("SELECT id FROM order_items", &schema).expect("should prune");
    assert!(result.relevant_tables.contains(&"order_items".to_string()));
}

// ─── JOIN queries ───────────────────────────────────────────────

#[test]
fn test_prune_join_includes_both_tables() {
    let schema = ecommerce_schema();
    let result = prune_schema(
        "SELECT u.name, o.total_amount FROM users u JOIN orders o ON u.id = o.user_id",
        &schema,
    )
    .expect("should prune");
    assert!(result.relevant_tables.contains(&"users".to_string()));
    assert!(result.relevant_tables.contains(&"orders".to_string()));
}

#[test]
fn test_prune_multi_join() {
    let schema = ecommerce_schema();
    let result = prune_schema(
        "SELECT u.name, p.name, oi.quantity \
         FROM users u \
         JOIN orders o ON u.id = o.user_id \
         JOIN order_items oi ON o.id = oi.order_id \
         JOIN products p ON oi.product_id = p.id",
        &schema,
    )
    .expect("should prune");
    assert_eq!(result.relevant_tables.len(), 4);
}

// ─── YAML output ─────────────────────────────────────────────────

#[test]
fn test_prune_yaml_is_valid() {
    let schema = ecommerce_schema();
    let result = prune_schema("SELECT id FROM users", &schema).expect("should prune");
    // The pruned YAML should be parseable as a SchemaDefinition
    let reparsed: SchemaDefinition =
        serde_yaml::from_str(&result.pruned_schema_yaml).expect("should parse pruned YAML");
    assert!(!reparsed.tables.is_empty());
}

#[test]
fn test_prune_yaml_contains_version() {
    let schema = ecommerce_schema();
    let result = prune_schema("SELECT id FROM users", &schema).expect("should prune");
    assert!(result.pruned_schema_yaml.contains("version"));
}

// ─── Case insensitivity ──────────────────────────────────────────

#[test]
fn test_prune_case_insensitive_table_match() {
    let schema = ecommerce_schema();
    let result = prune_schema("SELECT id FROM USERS", &schema).expect("should prune");
    assert!(result.relevant_tables.contains(&"users".to_string()));
}

// ─── Custom schema with unique columns ──────────────────────────

#[test]
fn test_prune_custom_schema_fk_expansion() {
    // Use unique column names to avoid cross-table column matching
    let schema = make_schema(vec![
        ("a", vec!["a_id", "b_ref"], vec![("b_ref", "b")]),
        ("b", vec!["b_id", "b_name"], vec![]),
        ("c", vec!["c_id", "c_data"], vec![]),
    ]);
    let result = prune_schema("SELECT a_id FROM a", &schema).expect("should prune");
    assert!(result.relevant_tables.contains(&"a".to_string()));
    // FK expansion should pull in "b"
    assert!(result.relevant_tables.contains(&"b".to_string()));
    // "c" should not be included
    assert!(!result.relevant_tables.contains(&"c".to_string()));
}

#[test]
fn test_prune_custom_schema_all_tables_match() {
    let schema = make_schema(vec![
        ("x", vec!["x_id"], vec![]),
        ("y", vec!["y_id"], vec![]),
    ]);
    // Using unique column names ensures only the correct table matches
    let result = prune_schema("SELECT x_id, y_id FROM x, y", &schema).expect("should prune");
    assert!(result.relevant_tables.contains(&"x".to_string()));
    assert!(result.relevant_tables.contains(&"y".to_string()));
}

// ─── Aggregate and complex queries ───────────────────────────────

#[test]
fn test_prune_aggregate_query() {
    let schema = ecommerce_schema();
    let result = prune_schema(
        "SELECT user_id, SUM(total_amount) FROM orders GROUP BY user_id",
        &schema,
    )
    .expect("should prune");
    assert!(result.relevant_tables.contains(&"orders".to_string()));
}

#[test]
fn test_prune_union_query() {
    let schema = ecommerce_schema();
    let result = prune_schema(
        "SELECT id, name FROM users UNION SELECT id, name FROM products",
        &schema,
    )
    .expect("should prune");
    assert!(result.relevant_tables.contains(&"users".to_string()));
    assert!(result.relevant_tables.contains(&"products".to_string()));
}

// ─── Subqueries ─────────────────────────────────────────────────

#[test]
fn test_prune_subquery_with_explicit_columns() {
    let schema = ecommerce_schema();
    // polyglot-sql's column extraction works on SELECT list columns.
    // Use explicit column names from different tables in SELECT lists.
    let result = prune_schema(
        "SELECT email FROM users WHERE id IN (SELECT user_id FROM orders)",
        &schema,
    )
    .expect("should prune");
    // "email" is unique to users, so users should be matched
    assert!(result.relevant_tables.contains(&"users".to_string()));
    // Note: polyglot may not extract "user_id" from the subquery SELECT,
    // so orders may or may not be included depending on the parser.
}
