//! Detect aggregation before JOIN pattern.
//!
//! BFS over the AST. At each level, if any node is a Join, check if any
//! sibling has a GroupBy descendant.

use crate::tagger::types::{TagReport, TagResult};
use polyglot_sql::expressions::Expression;
use polyglot_sql::generator::Generator;
use std::collections::VecDeque;

pub const TAG_NAME: &str = "agg_before_join";

/// Format a Join expression as SQL, since the generator doesn't handle standalone Join.
fn format_join_sql(expr: &Expression) -> String {
    match expr {
        Expression::Join(j) => {
            let kind_str = format!("{:?}", j.kind).to_uppercase();
            let join_keyword = match kind_str.as_str() {
                "INNER" => "JOIN".to_string(),
                "IMPLICIT" => ",".to_string(),
                other => format!("{} JOIN", other),
            };
            let table_sql = Generator::sql(&j.this).unwrap_or_default();
            let on_sql =
                j.on.as_ref()
                    .map(|on| format!("\n  ON {}", Generator::sql(on).unwrap_or_default()))
                    .unwrap_or_default();
            format!("{}\n  {}{}", join_keyword, table_sql, on_sql)
        }
        _ => Generator::pretty_sql(expr).unwrap_or_default(),
    }
}

/// Get ALL direct children of an expression for BFS traversal.
fn get_all_children(expr: &Expression) -> Vec<Expression> {
    match expr {
        Expression::Select(sel) => {
            let mut children = Vec::new();
            for e in &sel.expressions {
                children.push(e.clone());
            }
            if let Some(ref from) = sel.from {
                for e in &from.expressions {
                    children.push(e.clone());
                }
            }
            for join in &sel.joins {
                children.push(Expression::Join(Box::new(join.clone())));
            }
            if let Some(ref w) = sel.where_clause {
                children.push(Expression::Where(Box::new(w.clone())));
            }
            if let Some(ref g) = sel.group_by {
                children.push(Expression::GroupBy(Box::new(g.clone())));
            }
            if let Some(ref h) = sel.having {
                children.push(Expression::Having(Box::new(h.clone())));
            }
            if let Some(ref o) = sel.order_by {
                children.push(Expression::OrderBy(Box::new(o.clone())));
            }
            if let Some(ref l) = sel.limit {
                children.push(Expression::Limit(Box::new(l.clone())));
            }
            if let Some(ref with) = sel.with {
                children.push(Expression::With(Box::new(with.clone())));
            }
            children
        }
        Expression::Union(u) => vec![u.left.clone(), u.right.clone()],
        Expression::Intersect(i) => vec![i.left.clone(), i.right.clone()],
        Expression::Except(e) => vec![e.left.clone(), e.right.clone()],
        Expression::Subquery(s) => vec![s.this.clone()],
        Expression::Paren(p) => vec![p.this.clone()],
        Expression::Alias(a) => vec![a.this.clone()],
        Expression::Join(j) => {
            let mut children = vec![j.this.clone()];
            if let Some(ref on) = j.on {
                children.push(on.clone());
            }
            children
        }
        Expression::Where(w) => vec![w.this.clone()],
        Expression::Having(h) => vec![h.this.clone()],
        Expression::GroupBy(g) => g.expressions.clone(),
        Expression::From(f) => f.expressions.clone(),
        Expression::With(w) => w
            .ctes
            .iter()
            .map(|c| Expression::Cte(Box::new(c.clone())))
            .collect(),
        Expression::Cte(c) => vec![c.this.clone()],
        Expression::And(op)
        | Expression::Or(op)
        | Expression::Eq(op)
        | Expression::Neq(op)
        | Expression::Lt(op)
        | Expression::Lte(op)
        | Expression::Gt(op)
        | Expression::Gte(op)
        | Expression::Add(op)
        | Expression::Sub(op)
        | Expression::Mul(op)
        | Expression::Div(op) => {
            vec![op.left.clone(), op.right.clone()]
        }
        Expression::Not(u) | Expression::Neg(u) => vec![u.this.clone()],
        Expression::Function(f) => f.args.clone(),
        Expression::AggregateFunction(f) => f.args.clone(),
        Expression::In(i) => {
            let mut v = vec![i.this.clone()];
            v.extend(i.expressions.clone());
            v
        }
        Expression::Between(b) => vec![b.this.clone(), b.low.clone(), b.high.clone()],
        Expression::Case(c) => {
            let mut v = Vec::new();
            if let Some(ref op) = c.operand {
                v.push(op.clone());
            }
            for w in &c.whens {
                v.push(w.0.clone());
                v.push(w.1.clone());
            }
            if let Some(ref else_) = c.else_ {
                v.push(else_.clone());
            }
            v
        }
        Expression::IsNull(i) => vec![i.this.clone()],
        Expression::Exists(e) => vec![e.this.clone()],
        Expression::Cast(c) | Expression::TryCast(c) | Expression::SafeCast(c) => {
            vec![c.this.clone()]
        }
        Expression::Upper(u)
        | Expression::Lower(u)
        | Expression::Length(u)
        | Expression::Abs(u) => {
            vec![u.this.clone()]
        }
        Expression::Sum(a) | Expression::Avg(a) | Expression::Min(a) | Expression::Max(a) => {
            vec![a.this.clone()]
        }
        Expression::Count(c) => match &c.this {
            Some(this) => vec![this.clone()],
            None => vec![],
        },
        Expression::Like(op) | Expression::ILike(op) => {
            vec![op.left.clone(), op.right.clone()]
        }
        _ => vec![],
    }
}

/// Recursively check if the expression has a GroupBy descendant at depth > 0.
fn has_group_by_descendant(expr: &Expression, level: usize) -> bool {
    if matches!(expr, Expression::GroupBy(_)) && level > 0 {
        return true;
    }

    let children = get_all_children(expr);
    for child in &children {
        if has_group_by_descendant(child, level + 1) {
            return true;
        }
    }
    false
}

pub fn check(expr: &Expression) -> TagResult {
    let mut queue: VecDeque<Expression> = VecDeque::new();
    queue.push_back(expr.clone());

    while !queue.is_empty() {
        let level_size = queue.len();
        let mut level_has_join = false;
        let mut join_idx: Option<usize> = None;

        for (i, item) in queue.iter().take(level_size).enumerate() {
            if matches!(item, Expression::Join(_)) {
                level_has_join = true;
                join_idx = Some(i);
            }
        }

        if level_has_join {
            for item in queue.iter().take(level_size) {
                if has_group_by_descendant(item, 0) {
                    let join_expr = &queue[join_idx.unwrap()];
                    let query_span = format_join_sql(join_expr);

                    return TagResult {
                        tag_name: TAG_NAME.to_string(),
                        triggered: true,
                        reports: vec![TagReport {
                            query_span,
                            functions: None,
                            description: None,
                            metadata: None,
                        }],
                    };
                }
            }
            return TagResult::not_triggered(TAG_NAME);
        }

        for _ in 0..level_size {
            let node = queue.pop_front().unwrap();
            for child in get_all_children(&node) {
                queue.push_back(child);
            }
        }
    }

    TagResult::not_triggered(TAG_NAME)
}

#[cfg(test)]
mod tests {
    use super::*;
    use polyglot_sql::dialects::DialectType;

    /// Run test body in a 32MB-stack thread (polyglot-sql parsing uses deep recursion).
    fn with_large_stack<F: FnOnce() + Send + 'static>(f: F) {
        std::thread::Builder::new()
            .stack_size(32 * 1024 * 1024)
            .spawn(f)
            .unwrap()
            .join()
            .unwrap();
    }

    #[test]
    fn triggered_on_agg_before_join() {
        with_large_stack(|| {
            let expr = polyglot_sql::parse_one(
                "SELECT * FROM (SELECT SUM(x) FROM a GROUP BY y) sub JOIN b ON sub.y = b.y",
                DialectType::Snowflake,
            )
            .unwrap();
            let r = check(&expr);
            assert!(r.triggered);
        });
    }

    #[test]
    fn not_triggered_simple_join() {
        with_large_stack(|| {
            let expr = polyglot_sql::parse_one(
                "SELECT * FROM t1 JOIN t2 ON t1.id = t2.id",
                DialectType::Snowflake,
            )
            .unwrap();
            let r = check(&expr);
            assert!(!r.triggered);
        });
    }
}
