//! Analyze over-trace and under-trace patterns in column lineage.
//!
//! For each mismatched query, classifies extra (superset) and missing (subset)
//! source refs by their SQL clause context: WHERE, JOIN ON, CASE WHEN, COALESCE,
//! function call, UNION, or unknown.
//!
//! Run: cargo run --example over_trace_patterns --release

use altimate_core::lineage::column_lineage;
use altimate_core::schema::types::SqlDialect;
use serde::Deserialize;
use std::collections::{BTreeSet, HashMap};
use std::io::BufRead;
use std::path::Path;

// ---------------------------------------------------------------------------
// Fixture
// ---------------------------------------------------------------------------

#[derive(Debug, Deserialize)]
struct Case {
    name: String,
    sql: String,
    schema: serde_json::Value,
    expected: HashMap<String, Vec<String>>,
}

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

fn schema_from_json(value: &serde_json::Value) -> HashMap<String, serde_json::Value> {
    match value {
        serde_json::Value::Object(map) => map.iter().map(|(k, v)| (k.clone(), v.clone())).collect(),
        _ => HashMap::new(),
    }
}

fn normalize(map: &HashMap<String, Vec<String>>) -> HashMap<String, BTreeSet<String>> {
    map.iter()
        .map(|(k, v)| (k.clone(), v.iter().cloned().collect()))
        .collect()
}

fn normalize_to_short(map: &HashMap<String, Vec<String>>) -> HashMap<String, BTreeSet<String>> {
    map.iter()
        .map(|(k, v)| {
            let short: BTreeSet<String> = v
                .iter()
                .map(|ref_str| {
                    let parts: Vec<&str> =
                        ref_str.split('.').map(|p| p.trim_matches('"')).collect();
                    if parts.len() >= 3 {
                        let table = parts[parts.len() - 2];
                        let col = parts[parts.len() - 1];
                        format!("\"{table}\".\"{col}\"")
                    } else {
                        ref_str.clone()
                    }
                })
                .collect();
            (k.clone(), short)
        })
        .collect()
}

/// Extract column name from a "TABLE"."COL" ref (uppercased).
fn col_name(ref_str: &str) -> String {
    let parts: Vec<&str> = ref_str.split('.').map(|p| p.trim_matches('"')).collect();
    parts.last().unwrap_or(&"").to_uppercase()
}

/// Extract table name from a "TABLE"."COL" ref (uppercased).
fn table_name(ref_str: &str) -> String {
    let parts: Vec<&str> = ref_str.split('.').map(|p| p.trim_matches('"')).collect();
    if parts.len() >= 2 {
        parts[parts.len() - 2].to_uppercase()
    } else {
        String::new()
    }
}

/// Collect all table names from the schema.
fn schema_tables(schema: &HashMap<String, serde_json::Value>) -> BTreeSet<String> {
    schema.keys().map(|k| k.to_uppercase()).collect()
}

// ---------------------------------------------------------------------------
// Clause context detection (simple string matching, case-insensitive)
// ---------------------------------------------------------------------------

/// Rough clause-based classification for where a column name appears in SQL.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
enum ClauseContext {
    /// In a CASE WHEN ... THEN ... END expression inside SELECT
    CaseWhenInSelect,
    /// In a WHERE clause
    WhereClause,
    /// In a JOIN ON condition
    JoinOnClause,
    /// In a COALESCE(...) in SELECT
    CoalesceInSelect,
    /// In a function call in SELECT (IFF, NVL, CONCAT, etc.)
    FunctionInSelect,
    /// Query has UNION — could be from a different branch
    UnionBranch,
    /// In PARTITION BY / ORDER BY of a window function
    WindowClause,
    /// In GROUP BY / HAVING
    GroupByHaving,
    /// Could not determine context
    Unknown,
}

impl std::fmt::Display for ClauseContext {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::CaseWhenInSelect => write!(f, "CASE WHEN in SELECT"),
            Self::WhereClause => write!(f, "WHERE clause"),
            Self::JoinOnClause => write!(f, "JOIN ON clause"),
            Self::CoalesceInSelect => write!(f, "COALESCE in SELECT"),
            Self::FunctionInSelect => write!(f, "Function in SELECT"),
            Self::UnionBranch => write!(f, "UNION branch"),
            Self::WindowClause => write!(f, "Window PARTITION/ORDER BY"),
            Self::GroupByHaving => write!(f, "GROUP BY / HAVING"),
            Self::Unknown => write!(f, "Unknown context"),
        }
    }
}

/// Given a SQL string and a column name, try to determine which clause
/// the column appears in. Returns the most specific match.
///
/// This is intentionally simple — no parsing, just substring search with
/// positional heuristics.
fn detect_clause_context(sql_upper: &str, col_name_upper: &str) -> ClauseContext {
    // Find positions of key clauses
    let where_pos = sql_upper.rfind(" WHERE ");
    let join_on_pos = sql_upper.rfind(" ON ");
    let group_pos = sql_upper.rfind(" GROUP BY ");
    let having_pos = sql_upper.rfind(" HAVING ");

    let has_union = sql_upper.contains(" UNION ");
    let has_case = sql_upper.contains("CASE ");
    let has_coalesce = sql_upper.contains("COALESCE(");
    let has_window = sql_upper.contains(" OVER(") || sql_upper.contains(" OVER (");

    // Find all occurrences of the column name in the SQL
    let mut positions: Vec<usize> = Vec::new();
    let mut search_from = 0;
    while let Some(pos) = sql_upper[search_from..].find(col_name_upper) {
        let abs_pos = search_from + pos;
        // Make sure it's a word boundary (not part of a longer identifier)
        let before_ok = abs_pos == 0
            || !sql_upper.as_bytes()[abs_pos - 1].is_ascii_alphanumeric()
                && sql_upper.as_bytes()[abs_pos - 1] != b'_';
        let after_pos = abs_pos + col_name_upper.len();
        let after_ok = after_pos >= sql_upper.len()
            || !sql_upper.as_bytes()[after_pos].is_ascii_alphanumeric()
                && sql_upper.as_bytes()[after_pos] != b'_';
        if before_ok && after_ok {
            positions.push(abs_pos);
        }
        search_from = abs_pos + 1;
    }

    if positions.is_empty() {
        return ClauseContext::Unknown;
    }

    // Check if any occurrence is inside a WHERE clause
    if let Some(wp) = where_pos {
        for &pos in &positions {
            if pos > wp {
                // Column appears after WHERE — could be in WHERE clause
                // But also check it's not after GROUP BY/HAVING/ORDER BY
                let after_group = group_pos.map_or(false, |gp| pos > gp);
                let after_having = having_pos.map_or(false, |hp| pos > hp);
                if !after_group && !after_having {
                    return ClauseContext::WhereClause;
                }
            }
        }
    }

    // Check if any occurrence is inside JOIN ON
    if let Some(jp) = join_on_pos {
        for &pos in &positions {
            if pos > jp && pos < jp + 200 {
                // Rough: within 200 chars of ON keyword
                if let Some(wp) = where_pos {
                    if pos < wp {
                        return ClauseContext::JoinOnClause;
                    }
                } else {
                    return ClauseContext::JoinOnClause;
                }
            }
        }
    }

    // Check GROUP BY / HAVING
    if let Some(gp) = group_pos {
        for &pos in &positions {
            if pos > gp {
                return ClauseContext::GroupByHaving;
            }
        }
    }
    if let Some(hp) = having_pos {
        for &pos in &positions {
            if pos > hp {
                return ClauseContext::GroupByHaving;
            }
        }
    }

    // Check window function context
    if has_window {
        // Look for col name near PARTITION BY or ORDER BY within OVER(...)
        if sql_upper.contains(&format!("PARTITION BY {col_name_upper}"))
            || sql_upper.contains(&format!("PARTITION BY  {col_name_upper}"))
            || sql_upper.contains(&format!("ORDER BY {col_name_upper}"))
        {
            return ClauseContext::WindowClause;
        }
    }

    // Check SELECT-level constructs (CASE, COALESCE, functions)
    if has_case {
        // Check if col appears after CASE/WHEN but before END
        if sql_upper.contains(&format!("WHEN {col_name_upper}"))
            || sql_upper.contains(&format!("THEN {col_name_upper}"))
            || sql_upper.contains(&format!("ELSE {col_name_upper}"))
        {
            return ClauseContext::CaseWhenInSelect;
        }
    }

    if has_coalesce {
        if sql_upper.contains(&format!("COALESCE({col_name_upper}"))
            || sql_upper.contains(&format!("COALESCE( {col_name_upper}"))
        {
            return ClauseContext::CoalesceInSelect;
        }
    }

    // Check common functions in SELECT
    let fn_patterns = [
        "IFF(",
        "NVL(",
        "NVL2(",
        "IFNULL(",
        "NULLIF(",
        "CONCAT(",
        "DECODE(",
        "GREATEST(",
        "LEAST(",
        "ZEROIFNULL(",
        "TRY_CAST(",
    ];
    for pat in &fn_patterns {
        if sql_upper.contains(pat) {
            // Rough check: does the col appear near this function?
            if let Some(fn_pos) = sql_upper.find(pat) {
                for &pos in &positions {
                    if pos > fn_pos && pos < fn_pos + 300 {
                        return ClauseContext::FunctionInSelect;
                    }
                }
            }
        }
    }

    if has_union {
        return ClauseContext::UnionBranch;
    }

    ClauseContext::Unknown
}

// ---------------------------------------------------------------------------
// Missing source classification
// ---------------------------------------------------------------------------

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
enum MissingReason {
    /// The table from ground truth is not in our schema at all
    TableNotInSchema,
    /// The table exists in schema but the column is not found (view resolution?)
    ViewResolution,
    /// Both table and column exist — we failed to trace
    TracingGap,
    /// Unknown
    Unknown,
}

impl std::fmt::Display for MissingReason {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::TableNotInSchema => write!(f, "Table not in schema"),
            Self::ViewResolution => write!(f, "View resolution (col missing from table)"),
            Self::TracingGap => write!(f, "Tracing gap (table+col exist)"),
            Self::Unknown => write!(f, "Unknown"),
        }
    }
}

fn classify_missing(
    ref_str: &str,
    schema_table_set: &BTreeSet<String>,
    schema: &HashMap<String, serde_json::Value>,
) -> MissingReason {
    let tbl = table_name(ref_str);
    let col = col_name(ref_str);

    if tbl.is_empty() || col.is_empty() {
        return MissingReason::Unknown;
    }

    // Check if table exists in schema (case-insensitive)
    if !schema_table_set.contains(&tbl) {
        return MissingReason::TableNotInSchema;
    }

    // Table exists — check if column exists in that table's schema
    for (table_key, cols_val) in schema {
        if table_key.to_uppercase() == tbl {
            if let serde_json::Value::Object(cols_map) = cols_val {
                let col_exists = cols_map.keys().any(|k| k.to_uppercase() == col);
                if col_exists {
                    return MissingReason::TracingGap;
                } else {
                    return MissingReason::ViewResolution;
                }
            }
        }
    }

    MissingReason::Unknown
}

// ---------------------------------------------------------------------------
// Main
// ---------------------------------------------------------------------------

fn main() {
    let fixture_path = Path::new(env!("CARGO_MANIFEST_DIR"))
        .join("../../tests/fixtures/lineage/snowflake_ground_truth.jsonl");

    if !fixture_path.exists() {
        eprintln!("Fixture not found: {}", fixture_path.display());
        return;
    }

    let file = std::fs::File::open(&fixture_path).expect("failed to open fixture");
    let reader = std::io::BufReader::with_capacity(1024 * 1024, file);
    let dialect = SqlDialect::Snowflake;

    let start = std::time::Instant::now();

    // -----------------------------------------------------------------------
    // Counters
    // -----------------------------------------------------------------------

    let mut total = 0u64;
    let mut exact = 0u64;
    let mut parse_fail = 0u64;

    // --- Extra source tracking (superset) ---
    let mut extra_clause_counts: HashMap<ClauseContext, u64> = HashMap::new();
    let mut extra_ref_freq: HashMap<String, u64> = HashMap::new(); // "TABLE.COL" -> count
    let mut queries_with_extras = 0u64;
    let mut total_extra_refs = 0u64;

    // --- Missing source tracking (subset) ---
    let mut missing_reason_counts: HashMap<MissingReason, u64> = HashMap::new();
    let mut missing_ref_freq: HashMap<String, u64> = HashMap::new();
    let mut queries_with_missing = 0u64;
    let mut total_missing_refs = 0u64;

    // --- Sample storage for printing ---
    struct ExtraSample {
        name: String,
        sql_prefix: String,
        col: String,
        extra_refs: Vec<String>,
        contexts: Vec<ClauseContext>,
    }

    struct MissingSample {
        name: String,
        sql_prefix: String,
        col: String,
        missing_refs: Vec<String>,
        reasons: Vec<MissingReason>,
    }

    let mut extra_samples: Vec<ExtraSample> = Vec::new();
    let mut missing_samples: Vec<MissingSample> = Vec::new();

    // -----------------------------------------------------------------------
    // Process
    // -----------------------------------------------------------------------

    for line in reader.lines() {
        let line = line.expect("read line");
        let line = line.trim();
        if line.is_empty() {
            continue;
        }

        let case: Case = match serde_json::from_str(line) {
            Ok(c) => c,
            Err(_) => continue,
        };
        total += 1;

        let schema = schema_from_json(&case.schema);
        let schema_tbl_set = schema_tables(&schema);

        let result = match column_lineage(&case.sql, dialect, &schema, None) {
            Ok(r) => r,
            Err(_) => {
                parse_fail += 1;
                continue;
            }
        };

        let actual = normalize_to_short(&result.column_dict);
        let expected = normalize(&case.expected);

        if actual == expected {
            exact += 1;
            continue;
        }

        let sql_upper = case.sql.to_uppercase();
        let sql_prefix: String = case
            .sql
            .chars()
            .filter(|c| !c.is_ascii_control() || *c == ' ')
            .take(300)
            .collect();

        let mut query_has_extras = false;
        let mut query_has_missing = false;

        // Analyze each expected column
        for (col, exp_srcs) in &expected {
            let act_srcs = match actual.get(col) {
                Some(s) => s,
                None => continue, // Missing output column — different problem
            };

            // Find EXTRA sources (in actual but not in expected)
            let extras: Vec<&String> = act_srcs.difference(exp_srcs).collect();
            if !extras.is_empty() {
                query_has_extras = true;
                for extra_ref in &extras {
                    total_extra_refs += 1;
                    let c = col_name(extra_ref);
                    let ctx = detect_clause_context(&sql_upper, &c);
                    *extra_clause_counts.entry(ctx).or_default() += 1;

                    // Track frequency of this specific ref
                    let key = format!("{}.{}", table_name(extra_ref), c);
                    *extra_ref_freq.entry(key).or_default() += 1;
                }

                // Collect sample
                if extra_samples.len() < 30 {
                    let contexts: Vec<ClauseContext> = extras
                        .iter()
                        .map(|r| detect_clause_context(&sql_upper, &col_name(r)))
                        .collect();
                    extra_samples.push(ExtraSample {
                        name: case.name.clone(),
                        sql_prefix: sql_prefix.clone(),
                        col: col.clone(),
                        extra_refs: extras.iter().map(|s| s.to_string()).collect(),
                        contexts,
                    });
                }
            }

            // Find MISSING sources (in expected but not in actual)
            let missing: Vec<&String> = exp_srcs.difference(act_srcs).collect();
            if !missing.is_empty() {
                query_has_missing = true;
                for miss_ref in &missing {
                    total_missing_refs += 1;
                    let reason = classify_missing(miss_ref, &schema_tbl_set, &schema);
                    *missing_reason_counts.entry(reason).or_default() += 1;

                    let key = format!("{}.{}", table_name(miss_ref), col_name(miss_ref));
                    *missing_ref_freq.entry(key).or_default() += 1;
                }

                // Collect sample
                if missing_samples.len() < 30 {
                    let reasons: Vec<MissingReason> = missing
                        .iter()
                        .map(|r| classify_missing(r, &schema_tbl_set, &schema))
                        .collect();
                    missing_samples.push(MissingSample {
                        name: case.name.clone(),
                        sql_prefix: sql_prefix.clone(),
                        col: col.clone(),
                        missing_refs: missing.iter().map(|s| s.to_string()).collect(),
                        reasons,
                    });
                }
            }
        }

        if query_has_extras {
            queries_with_extras += 1;
        }
        if query_has_missing {
            queries_with_missing += 1;
        }
    }

    let elapsed = start.elapsed();

    // -----------------------------------------------------------------------
    // Report
    // -----------------------------------------------------------------------

    println!();
    println!("{}", "=".repeat(78));
    println!("  OVER-TRACE & UNDER-TRACE PATTERN ANALYSIS");
    println!("{}", "=".repeat(78));
    println!();
    println!("Total queries:      {total}");
    println!(
        "Exact match:        {exact} ({:.1}%)",
        exact as f64 / total as f64 * 100.0
    );
    println!("Parse fail:         {parse_fail}");
    println!("Elapsed:            {:.1}s", elapsed.as_secs_f64());

    // ===== EXTRA SOURCES (SUPERSET) =====
    println!();
    println!("{}", "=".repeat(78));
    println!("  EXTRA SOURCES (SUPERSET) — We trace too many source columns");
    println!("{}", "=".repeat(78));
    println!();
    println!("Queries with extras:   {queries_with_extras}");
    println!("Total extra refs:      {total_extra_refs}");
    if queries_with_extras > 0 {
        println!(
            "Avg extras/query:      {:.1}",
            total_extra_refs as f64 / queries_with_extras as f64
        );
    }

    println!();
    println!("--- Extra Source Clause Context (where extra columns actually appear) ---");
    let mut sorted_ctx: Vec<_> = extra_clause_counts.iter().collect();
    sorted_ctx.sort_by(|a, b| b.1.cmp(a.1));
    println!("  {:<35} {:>8} {:>8}", "Context", "Count", "%");
    println!("  {}", "-".repeat(53));
    for (ctx, count) in &sorted_ctx {
        println!(
            "  {:<35} {:>8} {:>7.1}%",
            ctx.to_string(),
            count,
            **count as f64 / total_extra_refs as f64 * 100.0
        );
    }

    println!();
    println!("--- Top 20 Most Common EXTRA Source Refs (TABLE.COL) ---");
    let mut extra_vec: Vec<_> = extra_ref_freq.iter().collect();
    extra_vec.sort_by(|a, b| b.1.cmp(a.1).then_with(|| a.0.cmp(b.0)));
    println!("  {:<3} {:<55} {:>7}", "#", "TABLE.COLUMN", "Count");
    println!("  {}", "-".repeat(67));
    for (i, (ref_key, count)) in extra_vec.iter().take(20).enumerate() {
        println!("  {:>2}. {:<55} {:>7}", i + 1, ref_key, count);
    }

    // ===== MISSING SOURCES (SUBSET) =====
    println!();
    println!("{}", "=".repeat(78));
    println!("  MISSING SOURCES (SUBSET) — We fail to trace some source columns");
    println!("{}", "=".repeat(78));
    println!();
    println!("Queries with missing:  {queries_with_missing}");
    println!("Total missing refs:    {total_missing_refs}");
    if queries_with_missing > 0 {
        println!(
            "Avg missing/query:     {:.1}",
            total_missing_refs as f64 / queries_with_missing as f64
        );
    }

    println!();
    println!("--- Missing Source Reason Breakdown ---");
    let mut sorted_reasons: Vec<_> = missing_reason_counts.iter().collect();
    sorted_reasons.sort_by(|a, b| b.1.cmp(a.1));
    println!("  {:<45} {:>8} {:>8}", "Reason", "Count", "%");
    println!("  {}", "-".repeat(63));
    for (reason, count) in &sorted_reasons {
        let pct = if total_missing_refs > 0 {
            **count as f64 / total_missing_refs as f64 * 100.0
        } else {
            0.0
        };
        println!("  {:<45} {:>8} {:>7.1}%", reason.to_string(), count, pct);
    }

    println!();
    println!("--- Top 20 Most Common MISSING Source Refs (TABLE.COL) ---");
    let mut missing_vec: Vec<_> = missing_ref_freq.iter().collect();
    missing_vec.sort_by(|a, b| b.1.cmp(a.1).then_with(|| a.0.cmp(b.0)));
    println!("  {:<3} {:<55} {:>7}", "#", "TABLE.COLUMN", "Count");
    println!("  {}", "-".repeat(67));
    for (i, (ref_key, count)) in missing_vec.iter().take(20).enumerate() {
        println!("  {:>2}. {:<55} {:>7}", i + 1, ref_key, count);
    }

    // ===== SAMPLE EXTRA CASES =====
    println!();
    println!("{}", "=".repeat(78));
    println!("  SAMPLE: 10 EXTRA SOURCE CASES (with clause context)");
    println!("{}", "=".repeat(78));

    for (i, sample) in extra_samples.iter().take(10).enumerate() {
        println!();
        println!("--- Extra #{} ---", i + 1);
        println!("Query:   {}", sample.name);
        println!("Column:  {}", sample.col);
        println!("SQL:     {}...", sample.sql_prefix);
        for (j, (ref_str, ctx)) in sample
            .extra_refs
            .iter()
            .zip(sample.contexts.iter())
            .enumerate()
        {
            println!("  Extra[{j}]: {ref_str}  =>  {ctx}");
        }
    }

    // ===== SAMPLE MISSING CASES =====
    println!();
    println!("{}", "=".repeat(78));
    println!("  SAMPLE: 10 MISSING SOURCE CASES (with reason)");
    println!("{}", "=".repeat(78));

    for (i, sample) in missing_samples.iter().take(10).enumerate() {
        println!();
        println!("--- Missing #{} ---", i + 1);
        println!("Query:   {}", sample.name);
        println!("Column:  {}", sample.col);
        println!("SQL:     {}...", sample.sql_prefix);
        for (j, (ref_str, reason)) in sample
            .missing_refs
            .iter()
            .zip(sample.reasons.iter())
            .enumerate()
        {
            println!("  Missing[{j}]: {ref_str}  =>  {reason}");
        }
    }

    println!();
    println!("=== Done ({:.1}s) ===", elapsed.as_secs_f64());
}
