//! Comprehensive tests for the lineage tracking module.

use altimate_core::lineage::tracker::track_lineage;
use altimate_core::schema::types::{ColumnDef, SchemaDefinition, SqlDialect, TableDef};
use std::collections::HashMap;

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

fn test_schema() -> SchemaDefinition {
    let mut tables = HashMap::new();
    tables.insert(
        "users".to_string(),
        TableDef {
            description: None,
            database: None,
            schema: None,
            columns: vec![
                col("id", "INT", false),
                col("name", "VARCHAR", true),
                col("email", "VARCHAR", true),
            ],
            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![
                col("id", "INT", false),
                col("user_id", "INT", false),
                col("amount", "DECIMAL(10,2)", true),
                col("status", "VARCHAR", true),
            ],
            primary_key: Some(vec!["id".to_string()]),
            foreign_keys: vec![],
            statistics: None,

            clustering_keys: vec![],
        },
    );
    tables.insert(
        "products".to_string(),
        TableDef {
            description: None,
            database: None,
            schema: None,
            columns: vec![
                col("id", "INT", false),
                col("name", "VARCHAR", false),
                col("price", "DECIMAL(10,2)", true),
            ],
            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,
    }
}

fn col(name: &str, dt: &str, nullable: bool) -> ColumnDef {
    ColumnDef {
        name: name.to_string(),
        data_type: dt.to_string(),
        nullable,
        description: None,
        tags: vec![],
        synonyms: vec![],

        statistics: None,
    }
}

// ─── Simple SELECT ───────────────────────────────────────────────

#[test]
fn test_simple_select_source_tables() {
    let r = track_lineage(&["SELECT id, name FROM users"], &test_schema()).unwrap();
    assert_eq!(r.queries.len(), 1);
    assert_eq!(r.queries[0].source_tables, vec!["users"]);
}

#[test]
fn test_simple_select_output_columns() {
    let r = track_lineage(&["SELECT id, name FROM users"], &test_schema()).unwrap();
    assert_eq!(r.queries[0].output_columns, vec!["id", "name"]);
}

#[test]
fn test_simple_select_no_target_tables() {
    let r = track_lineage(&["SELECT id FROM users"], &test_schema()).unwrap();
    assert!(r.queries[0].target_tables.is_empty());
}

#[test]
fn test_simple_select_edges() {
    let r = track_lineage(&["SELECT id, name FROM users"], &test_schema()).unwrap();
    assert!(r.queries[0].edges.len() >= 2);
}

// ─── SELECT with alias ──────────────────────────────────────────

#[test]
fn test_alias_resolves_to_table() {
    let r = track_lineage(
        &["SELECT u.id, u.name AS user_name FROM users u"],
        &test_schema(),
    )
    .unwrap();
    assert_eq!(r.queries[0].source_tables, vec!["users"]);
    assert!(
        r.queries[0]
            .edges
            .iter()
            .any(|e| e.source.table == "users" && e.source.column == "id")
    );
}

#[test]
fn test_alias_output_column_name() {
    let r = track_lineage(&["SELECT u.name AS user_name FROM users u"], &test_schema()).unwrap();
    assert_eq!(r.queries[0].output_columns, vec!["user_name"]);
}

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

#[test]
fn test_join_sources_both_tables() {
    let r = track_lineage(
        &["SELECT u.name, o.amount FROM users u JOIN orders o ON u.id = o.user_id"],
        &test_schema(),
    )
    .unwrap();
    let sources = &r.queries[0].source_tables;
    assert!(sources.contains(&"users".to_string()));
    assert!(sources.contains(&"orders".to_string()));
}

#[test]
fn test_join_output_columns() {
    let r = track_lineage(
        &["SELECT u.name, o.amount FROM users u JOIN orders o ON u.id = o.user_id"],
        &test_schema(),
    )
    .unwrap();
    assert_eq!(r.queries[0].output_columns, vec!["name", "amount"]);
}

#[test]
fn test_three_way_join() {
    let r = track_lineage(
        &["SELECT u.name, o.amount, p.price \
           FROM users u \
           JOIN orders o ON u.id = o.user_id \
           JOIN products p ON p.id = o.id"],
        &test_schema(),
    )
    .unwrap();
    assert_eq!(r.queries[0].source_tables.len(), 3);
}

// ─── Wildcard expansion ─────────────────────────────────────────

#[test]
fn test_wildcard_expands_all_columns() {
    let r = track_lineage(&["SELECT * FROM users"], &test_schema()).unwrap();
    let cols = &r.queries[0].output_columns;
    assert!(cols.contains(&"id".to_string()));
    assert!(cols.contains(&"name".to_string()));
    assert!(cols.contains(&"email".to_string()));
}

#[test]
fn test_wildcard_creates_direct_edges() {
    let r = track_lineage(&["SELECT * FROM users"], &test_schema()).unwrap();
    for edge in &r.queries[0].edges {
        assert_eq!(edge.transform_type, "direct");
    }
}

// ─── Aggregate transforms ────────────────────────────────────────

#[test]
fn test_aggregate_classified_correctly() {
    let r = track_lineage(
        &["SELECT user_id, SUM(amount) AS total FROM orders GROUP BY user_id"],
        &test_schema(),
    )
    .unwrap();
    assert!(
        r.queries[0]
            .edges
            .iter()
            .any(|e| e.transform_type == "aggregate" && e.target.column == "total")
    );
}

#[test]
fn test_direct_column_in_aggregate_query() {
    let r = track_lineage(
        &["SELECT user_id, SUM(amount) AS total FROM orders GROUP BY user_id"],
        &test_schema(),
    )
    .unwrap();
    assert!(
        r.queries[0]
            .edges
            .iter()
            .any(|e| e.transform_type == "direct" && e.target.column == "user_id")
    );
}

#[test]
fn test_count_is_aggregate() {
    let r = track_lineage(&["SELECT COUNT(id) AS cnt FROM users"], &test_schema()).unwrap();
    assert!(
        r.queries[0]
            .edges
            .iter()
            .any(|e| e.transform_type == "aggregate")
    );
}

// ─── Expression transforms ───────────────────────────────────────

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

#[test]
fn test_case_is_transform() {
    let r = track_lineage(
        &["SELECT CASE WHEN amount > 100 THEN 'high' ELSE 'low' END AS tier FROM orders"],
        &test_schema(),
    )
    .unwrap();
    assert!(
        r.queries[0]
            .edges
            .iter()
            .any(|e| e.transform_type == "transform")
    );
}

// ─── INSERT INTO ... SELECT ──────────────────────────────────────

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

#[test]
fn test_insert_select_edges_map_to_target_columns() {
    let r = track_lineage(
        &["INSERT INTO orders (id, user_id, amount) SELECT id, id, 100 FROM users"],
        &test_schema(),
    )
    .unwrap();
    let q = &r.queries[0];
    assert!(q.edges.iter().any(|e| e.target.table == "orders"));
}

// ─── CREATE TABLE AS SELECT ──────────────────────────────────────

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

// ─── Multi-query dependencies ────────────────────────────────────

#[test]
fn test_multi_query_dependency_order() {
    let r = track_lineage(
        &[
            "SELECT id, name FROM users",
            "INSERT INTO orders (user_id) SELECT id FROM users",
        ],
        &test_schema(),
    )
    .unwrap();
    assert_eq!(r.queries.len(), 2);
    assert!(!r.dependency_order.is_empty());
}

#[test]
fn test_dependency_order_sources_before_targets() {
    let r = track_lineage(
        &["INSERT INTO orders (user_id) SELECT id FROM users"],
        &test_schema(),
    )
    .unwrap();
    let users_pos = r.dependency_order.iter().position(|t| t == "users");
    let orders_pos = r.dependency_order.iter().position(|t| t == "orders");
    if let (Some(u), Some(o)) = (users_pos, orders_pos) {
        assert!(u < o, "users should come before orders in dependency order");
    }
}

// ─── Impact map ──────────────────────────────────────────────────

#[test]
fn test_impact_map_not_empty() {
    let r = track_lineage(
        &["SELECT u.id, u.name AS user_name FROM users u"],
        &test_schema(),
    )
    .unwrap();
    assert!(!r.impact_map.is_empty());
}

#[test]
fn test_impact_map_source_column_tracked() {
    let r = track_lineage(
        &["SELECT u.id, u.name AS user_name FROM users u"],
        &test_schema(),
    )
    .unwrap();
    assert!(
        r.impact_map
            .iter()
            .any(|entry| entry.source.table == "users" && entry.source.column == "id")
    );
}

#[test]
fn test_impact_map_shows_affected_targets() {
    let r = track_lineage(&["SELECT u.id FROM users u"], &test_schema()).unwrap();
    let entry = r
        .impact_map
        .iter()
        .find(|e| e.source.column == "id" && e.source.table == "users");
    assert!(entry.is_some());
    assert!(!entry.unwrap().affected.is_empty());
}

// ─── Parse errors ────────────────────────────────────────────────

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

#[test]
fn test_error_message_includes_query_number() {
    let r = track_lineage(&["SELECT 1", "THIS IS NOT SQL"], &empty_schema());
    assert!(r.is_err());
    assert!(r.unwrap_err().contains("query 2"));
}

// ─── Empty queries ───────────────────────────────────────────────

#[test]
fn test_empty_query_list() {
    let r = track_lineage(&[], &test_schema()).unwrap();
    assert!(r.queries.is_empty());
    assert!(r.dependency_order.is_empty());
}

// ─── Non-SELECT statements ──────────────────────────────────────

#[test]
fn test_create_table_no_as_select() {
    let r = track_lineage(&["CREATE TABLE t (id INT, name TEXT)"], &test_schema()).unwrap();
    let q = &r.queries[0];
    assert!(q.source_tables.is_empty());
    assert!(q.target_tables.contains(&"t".to_string()));
}

// ─── UNION queries ───────────────────────────────────────────────

#[test]
fn test_union_includes_both_source_tables() {
    let r = track_lineage(
        &["SELECT id, name FROM users UNION SELECT id, name FROM products"],
        &test_schema(),
    )
    .unwrap();
    let sources = &r.queries[0].source_tables;
    assert!(sources.contains(&"users".to_string()));
    assert!(sources.contains(&"products".to_string()));
}

// ─── Subquery sources ────────────────────────────────────────────

#[test]
fn test_subquery_in_from_extracts_tables() {
    let r = track_lineage(
        &["SELECT sub.id FROM (SELECT id FROM users) sub"],
        &test_schema(),
    )
    .unwrap();
    assert!(r.queries[0].source_tables.contains(&"users".to_string()));
}
