//! Query rewriting engine.
//!
//! Analyzes SQL queries and suggests optimizations: subquery-to-JOIN,
//! UNION-to-UNION ALL, SELECT * expansion, redundant DISTINCT removal,
//! EXISTS instead of COUNT, and index suggestions.

use super::types::*;
use crate::schema::types::SchemaDefinition;
use sqlparser::ast::{
    Distinct, Expr, FunctionArg, FunctionArgExpr, FunctionArguments, GroupByExpr, Join,
    JoinConstraint, JoinOperator, ObjectNamePart, OrderByKind, Query, Select, SelectItem, SetExpr,
    SetOperator, SetQuantifier, Statement, TableFactor, TableWithJoins, Value,
};
use sqlparser::dialect::GenericDialect;
use sqlparser::parser::Parser;

/// Analyze a SQL query and suggest optimized rewrites.
pub fn rewrite(sql: &str, schema: &SchemaDefinition) -> RewriteResult {
    let dialect = GenericDialect {};
    let statements = match Parser::parse_sql(&dialect, sql) {
        Ok(stmts) => stmts,
        Err(_) => {
            return RewriteResult {
                original_sql: sql.to_string(),
                suggestions: vec![],
                index_suggestions: vec![],
            };
        }
    };

    let mut suggestions = Vec::new();
    let mut index_suggestions = Vec::new();

    for stmt in &statements {
        if let Statement::Query(query) = stmt {
            check_star_expansion(query, schema, &mut suggestions);
            check_subquery_to_join(query, &mut suggestions);
            check_redundant_distinct(query, &mut suggestions);
            check_exists_instead_of_count(query, &mut suggestions);
            check_union_to_union_all(query, &mut suggestions);
            collect_index_suggestions(query, schema, &mut index_suggestions);
        }
    }

    RewriteResult {
        original_sql: sql.to_string(),
        suggestions,
        index_suggestions,
    }
}

/// Rule: SELECT * -> SELECT col1, col2, ...
fn check_star_expansion(
    query: &Query,
    schema: &SchemaDefinition,
    suggestions: &mut Vec<RewriteSuggestion>,
) {
    let select = match extract_select(query) {
        Some(s) => s,
        None => return,
    };

    let has_wildcard = select.projection.iter().any(|item| {
        matches!(
            item,
            SelectItem::Wildcard(_) | SelectItem::QualifiedWildcard(_, _)
        )
    });

    if !has_wildcard {
        return;
    }

    let table_names = extract_table_names_from_select(select);

    let mut expanded_columns = Vec::new();
    for table_name in &table_names {
        if let Some(cols) = schema.column_names(table_name) {
            for col in cols {
                if table_names.len() > 1 {
                    expanded_columns.push(format!("{table_name}.{col}"));
                } else {
                    expanded_columns.push(col.to_string());
                }
            }
        }
    }

    if expanded_columns.is_empty() {
        return;
    }

    let column_list = expanded_columns.join(", ");
    let original_sql = query.to_string();
    let rewritten = replace_select_star(&original_sql, &column_list);

    suggestions.push(RewriteSuggestion {
        rule: "star_to_columns".to_string(),
        explanation: "Replace SELECT * with explicit column list for clarity and performance"
            .to_string(),
        rewritten_sql: rewritten,
        improvement: "Avoids fetching unnecessary columns, makes schema dependencies explicit"
            .to_string(),
        confidence: 0.9,
    });
}

/// Simple SELECT * replacement in the SQL string.
fn replace_select_star(sql: &str, column_list: &str) -> String {
    let lower = sql.to_lowercase();
    if let Some(pos) = lower.find("select *") {
        let prefix = &sql[..pos + 7];
        let suffix = &sql[pos + 8..];
        format!("{prefix} {column_list}{suffix}")
    } else if let Some(pos) = lower.find("select  *") {
        let prefix = &sql[..pos + 7];
        let suffix = &sql[pos + 9..];
        format!("{prefix} {column_list}{suffix}")
    } else {
        format!("-- Expanded columns: {column_list}\n{sql}")
    }
}

/// Rule: WHERE x IN (SELECT col FROM table) -> JOIN table ON x = table.col
fn check_subquery_to_join(query: &Query, suggestions: &mut Vec<RewriteSuggestion>) {
    let select = match extract_select(query) {
        Some(s) => s,
        None => return,
    };

    let selection = match &select.selection {
        Some(expr) => expr,
        None => return,
    };

    if let Some((outer_col, sub_col, sub_table)) = extract_in_subquery(selection) {
        let original_sql = query.to_string();
        let rewritten =
            rewrite_in_subquery_to_join(&original_sql, &outer_col, &sub_col, &sub_table);

        suggestions.push(RewriteSuggestion {
            rule: "subquery_to_join".to_string(),
            explanation: format!(
                "Convert IN subquery on {sub_table} to a JOIN for better performance"
            ),
            rewritten_sql: rewritten,
            improvement: "JOINs are generally faster than correlated subqueries".to_string(),
            confidence: 0.8,
        });
    }
}

/// Extract `col IN (SELECT sub_col FROM sub_table)` patterns.
fn extract_in_subquery(expr: &Expr) -> Option<(String, String, String)> {
    match expr {
        Expr::InSubquery {
            expr: outer_expr,
            subquery,
            negated: false,
            ..
        } => {
            let outer_col = expr_to_column_name(outer_expr)?;
            let sub_select = extract_select(subquery)?;

            if sub_select.projection.len() != 1 {
                return None;
            }
            let sub_col = match &sub_select.projection[0] {
                SelectItem::UnnamedExpr(e) => expr_to_column_name(e)?,
                _ => return None,
            };

            if sub_select.from.len() != 1 {
                return None;
            }
            let sub_table = extract_table_name_from_table_with_joins(&sub_select.from[0])?;

            Some((outer_col, sub_col, sub_table))
        }
        Expr::BinaryOp { left, right, .. } => {
            extract_in_subquery(left).or_else(|| extract_in_subquery(right))
        }
        _ => None,
    }
}

/// Rewrite an IN-subquery to a JOIN in the SQL string.
fn rewrite_in_subquery_to_join(
    sql: &str,
    outer_col: &str,
    sub_col: &str,
    sub_table: &str,
) -> String {
    let lower = sql.to_lowercase();

    let in_pattern = format!("{} in", outer_col.to_lowercase());
    if let Some(in_pos) = lower.find(&in_pattern) {
        let after_in = &sql[in_pos + outer_col.len()..];
        let lower_after = after_in.to_lowercase();
        if let Some(paren_start) = lower_after.find('(') {
            let abs_paren = in_pos + outer_col.len() + paren_start;
            if let Some(paren_end) = find_matching_paren(sql, abs_paren) {
                let before_in = sql[..in_pos].trim();
                let after_paren = sql[paren_end + 1..].trim();

                let cleaned_where = clean_where_clause(before_in, after_paren);

                if let Some(where_pos) = lower.find("where") {
                    let before_where = sql[..where_pos].trim();
                    let join_clause = format!(
                        "{before_where} JOIN {sub_table} ON {outer_col} = {sub_table}.{sub_col}"
                    );

                    if cleaned_where.is_empty() {
                        return join_clause;
                    }
                    return format!("{join_clause} WHERE {cleaned_where}");
                }
            }
        }
    }

    format!("-- Suggested: JOIN {sub_table} ON {outer_col} = {sub_table}.{sub_col}\n{sql}")
}

/// Find matching closing parenthesis.
fn find_matching_paren(sql: &str, open_pos: usize) -> Option<usize> {
    let bytes = sql.as_bytes();
    let mut depth = 0;
    for (i, &b) in bytes.iter().enumerate().skip(open_pos) {
        match b {
            b'(' => depth += 1,
            b')' => {
                depth -= 1;
                if depth == 0 {
                    return Some(i);
                }
            }
            _ => {}
        }
    }
    None
}

/// Clean up a WHERE clause after removing a condition.
fn clean_where_clause(before: &str, after: &str) -> String {
    let before = before.trim();
    let after = after.trim();

    let before = strip_trailing_logical_op(before);
    let after = strip_leading_logical_op(after);

    let before = before
        .strip_suffix("WHERE")
        .or_else(|| before.strip_suffix("where"))
        .unwrap_or(before)
        .trim();

    match (before.is_empty(), after.is_empty()) {
        (true, true) => String::new(),
        (true, false) => after.to_string(),
        (false, true) => before.to_string(),
        (false, false) => format!("{before} AND {after}"),
    }
}

fn strip_trailing_logical_op(s: &str) -> &str {
    let trimmed = s.trim();
    trimmed
        .strip_suffix("AND")
        .or_else(|| trimmed.strip_suffix("and"))
        .or_else(|| trimmed.strip_suffix("OR"))
        .or_else(|| trimmed.strip_suffix("or"))
        .unwrap_or(trimmed)
        .trim()
}

fn strip_leading_logical_op(s: &str) -> &str {
    let trimmed = s.trim();
    trimmed
        .strip_prefix("AND ")
        .or_else(|| trimmed.strip_prefix("and "))
        .or_else(|| trimmed.strip_prefix("OR "))
        .or_else(|| trimmed.strip_prefix("or "))
        .unwrap_or(trimmed)
        .trim()
}

/// Rule: SELECT DISTINCT ... GROUP BY ... -> remove DISTINCT
fn check_redundant_distinct(query: &Query, suggestions: &mut Vec<RewriteSuggestion>) {
    let select = match extract_select(query) {
        Some(s) => s,
        None => return,
    };

    let has_distinct = matches!(select.distinct, Some(Distinct::Distinct));
    let has_group_by = !matches!(select.group_by, GroupByExpr::All(ref exprs) if exprs.is_empty());

    if has_distinct && has_group_by {
        let original_sql = query.to_string();
        let rewritten = remove_distinct(&original_sql);

        suggestions.push(RewriteSuggestion {
            rule: "redundant_distinct".to_string(),
            explanation: "DISTINCT is redundant when GROUP BY already ensures unique rows"
                .to_string(),
            rewritten_sql: rewritten,
            improvement: "Removes unnecessary deduplication step".to_string(),
            confidence: 0.85,
        });
    }
}

/// Remove DISTINCT from a SQL string.
fn remove_distinct(sql: &str) -> String {
    let lower = sql.to_lowercase();
    if let Some(pos) = lower.find("select distinct") {
        let prefix = &sql[..pos];
        let suffix = &sql[pos + 15..];
        format!("{prefix}SELECT{suffix}")
    } else {
        sql.to_string()
    }
}

/// Rule: WHERE (SELECT COUNT(*) FROM t WHERE ...) > 0 -> WHERE EXISTS (...)
fn check_exists_instead_of_count(query: &Query, suggestions: &mut Vec<RewriteSuggestion>) {
    let select = match extract_select(query) {
        Some(s) => s,
        None => return,
    };

    let selection = match &select.selection {
        Some(expr) => expr,
        None => return,
    };

    if let Some(count_info) = find_count_gt_zero(selection) {
        let original_sql = query.to_string();
        let rewritten = rewrite_count_to_exists(&original_sql, &count_info);

        suggestions.push(RewriteSuggestion {
            rule: "exists_instead_of_count".to_string(),
            explanation: "Use EXISTS instead of COUNT(*) > 0 for existence checks".to_string(),
            rewritten_sql: rewritten,
            improvement: "EXISTS can short-circuit after finding the first matching row"
                .to_string(),
            confidence: 0.9,
        });
    }
}

struct CountGtZeroInfo {
    subquery_sql: String,
}

/// Find patterns like `(SELECT COUNT(*) FROM ...) > 0`.
fn find_count_gt_zero(expr: &Expr) -> Option<CountGtZeroInfo> {
    match expr {
        Expr::BinaryOp { left, op, right } => {
            use sqlparser::ast::BinaryOperator;
            match op {
                BinaryOperator::Gt => {
                    if is_zero_value(right)
                        && let Some(subquery_sql) = extract_count_subquery(left)
                    {
                        return Some(CountGtZeroInfo { subquery_sql });
                    }
                    None
                }
                BinaryOperator::And | BinaryOperator::Or => {
                    find_count_gt_zero(left).or_else(|| find_count_gt_zero(right))
                }
                _ => None,
            }
        }
        _ => None,
    }
}

fn is_zero_value(expr: &Expr) -> bool {
    matches!(expr, Expr::Value(val) if matches!(&val.value, Value::Number(n, _) if n == "0"))
}

fn extract_count_subquery(expr: &Expr) -> Option<String> {
    if let Expr::Subquery(query) = expr {
        let select = extract_select(query)?;
        if select.projection.len() == 1
            && let SelectItem::UnnamedExpr(Expr::Function(func)) = &select.projection[0]
        {
            let func_name = func.name.to_string().to_lowercase();
            if func_name == "count" {
                return Some(query.to_string());
            }
        }
    }
    None
}

fn rewrite_count_to_exists(sql: &str, info: &CountGtZeroInfo) -> String {
    let lower = sql.to_lowercase();
    let count_lower = info.subquery_sql.to_lowercase();

    if let Some(sub_pos) = lower.find(&count_lower) {
        let paren_start = sql[..sub_pos].rfind('(');
        let after_sub = sub_pos + info.subquery_sql.len();
        let remaining = &sql[after_sub..];
        let gt_zero_pat = remaining
            .find("> 0")
            .or_else(|| remaining.find(">0"))
            .map(|p| p + after_sub);

        if let (Some(paren_start), Some(gt_pos)) = (paren_start, gt_zero_pat) {
            let gt_end = if sql[gt_pos..].starts_with("> 0") {
                gt_pos + 3
            } else {
                gt_pos + 2
            };

            let exists_subquery = replace_count_with_one(&info.subquery_sql);

            let before = &sql[..paren_start];
            let after = &sql[gt_end..];
            return format!("{before}EXISTS ({exists_subquery}){after}");
        }
    }

    format!("-- Suggested: use EXISTS instead of COUNT(*) > 0\n{sql}")
}

fn replace_count_with_one(subquery_sql: &str) -> String {
    let lower = subquery_sql.to_lowercase();
    if let Some(count_pos) = lower.find("count(")
        && let Some(close_paren) = subquery_sql[count_pos..].find(')')
    {
        let before = &subquery_sql[..count_pos];
        let after = &subquery_sql[count_pos + close_paren + 1..];
        return format!("{before}1{after}");
    }
    if let Some(count_pos) = lower.find("count (")
        && let Some(close_paren) = subquery_sql[count_pos..].find(')')
    {
        let before = &subquery_sql[..count_pos];
        let after = &subquery_sql[count_pos + close_paren + 1..];
        return format!("{before}1{after}");
    }
    subquery_sql.to_string()
}

/// Rule: UNION -> UNION ALL when safe
fn check_union_to_union_all(query: &Query, suggestions: &mut Vec<RewriteSuggestion>) {
    if let SetExpr::SetOperation {
        op: SetOperator::Union,
        set_quantifier,
        left,
        right,
    } = query.body.as_ref()
    {
        if matches!(set_quantifier, SetQuantifier::All) {
            return;
        }

        let left_tables = extract_tables_from_set_expr(left);
        let right_tables = extract_tables_from_set_expr(right);

        let has_overlap = left_tables.iter().any(|t| right_tables.contains(t));

        if has_overlap && has_where_clause(left) && has_where_clause(right) {
            let original_sql = query.to_string();
            let rewritten = original_sql.replacen("UNION", "UNION ALL", 1);
            let rewritten = if rewritten == original_sql {
                original_sql.replacen("union", "UNION ALL", 1)
            } else {
                rewritten
            };

            suggestions.push(RewriteSuggestion {
                rule: "union_to_union_all".to_string(),
                explanation: "Consider UNION ALL if duplicate elimination is not needed"
                    .to_string(),
                rewritten_sql: rewritten,
                improvement: "UNION ALL avoids an expensive sort/dedup step".to_string(),
                confidence: 0.6,
            });
        }
    }
}

/// Collect index suggestions from WHERE, JOIN, and ORDER BY clauses.
fn collect_index_suggestions(
    query: &Query,
    schema: &SchemaDefinition,
    index_suggestions: &mut Vec<IndexSuggestion>,
) {
    let select = match extract_select(query) {
        Some(s) => s,
        None => return,
    };

    let table_names = extract_table_names_from_select(select);
    let mut seen = std::collections::HashSet::new();

    // WHERE clause columns
    if let Some(ref selection) = select.selection {
        let where_cols = extract_columns_from_expr(selection);
        for col in where_cols {
            let table = resolve_column_table(&col, &table_names, schema);
            if let Some(table) = table {
                let key = format!("{table}.{col}");
                if seen.insert(key) {
                    let idx_name = format!("idx_{table}_{col}");
                    index_suggestions.push(IndexSuggestion {
                        table: table.clone(),
                        columns: vec![col.clone()],
                        reason: format!("Column '{col}' is used in WHERE clause"),
                        ddl: format!("CREATE INDEX {idx_name} ON {table} ({col});"),
                    });
                }
            }
        }
    }

    // JOIN ON columns
    for table_with_joins in &select.from {
        for join in &table_with_joins.joins {
            if let Some(cols) = extract_join_columns(join) {
                for (left_col, right_table, right_col) in &cols {
                    let key = format!("{right_table}.{right_col}");
                    if seen.insert(key.clone()) {
                        let idx_name = format!("idx_{right_table}_{right_col}");
                        index_suggestions.push(IndexSuggestion {
                            table: right_table.clone(),
                            columns: vec![right_col.clone()],
                            reason: format!(
                                "Column '{right_col}' is used in JOIN condition with '{left_col}'"
                            ),
                            ddl: format!("CREATE INDEX {idx_name} ON {right_table} ({right_col});"),
                        });
                    }
                }
            }
        }
    }

    // ORDER BY columns
    let order_exprs = query
        .order_by
        .as_ref()
        .and_then(|ob| match &ob.kind {
            OrderByKind::Expressions(exprs) => Some(exprs.as_slice()),
            _ => None,
        })
        .unwrap_or(&[]);
    for order_item in order_exprs {
        if let Some(col) = expr_to_column_name(&order_item.expr) {
            let table = resolve_column_table(&col, &table_names, schema);
            if let Some(table) = table {
                let key = format!("{table}.{col}");
                if seen.insert(key) {
                    let idx_name = format!("idx_{table}_{col}");
                    index_suggestions.push(IndexSuggestion {
                        table: table.clone(),
                        columns: vec![col.clone()],
                        reason: format!("Column '{col}' is used in ORDER BY"),
                        ddl: format!("CREATE INDEX {idx_name} ON {table} ({col});"),
                    });
                }
            }
        }
    }
}

// -- Helper functions --

fn extract_select(query: &Query) -> Option<&Select> {
    if let SetExpr::Select(ref select) = *query.body {
        Some(select)
    } else {
        None
    }
}

fn extract_table_names_from_select(select: &Select) -> Vec<String> {
    let mut names = Vec::new();
    for twj in &select.from {
        if let Some(name) = extract_table_name_from_table_with_joins(twj) {
            names.push(name);
        }
        for join in &twj.joins {
            if let Some(name) = extract_table_name_from_factor(&join.relation) {
                names.push(name);
            }
        }
    }
    names
}

fn extract_table_name_from_table_with_joins(twj: &TableWithJoins) -> Option<String> {
    extract_table_name_from_factor(&twj.relation)
}

fn extract_table_name_from_factor(factor: &TableFactor) -> Option<String> {
    if let TableFactor::Table { name, .. } = factor {
        name.0.last().map(|part| match part {
            ObjectNamePart::Identifier(ident) => ident.value.clone(),
        })
    } else {
        None
    }
}

fn expr_to_column_name(expr: &Expr) -> Option<String> {
    match expr {
        Expr::Identifier(ident) => Some(ident.value.clone()),
        Expr::CompoundIdentifier(idents) => idents.last().map(|i| i.value.clone()),
        _ => None,
    }
}

fn extract_columns_from_expr(expr: &Expr) -> Vec<String> {
    let mut cols = Vec::new();
    match expr {
        Expr::Identifier(ident) => {
            cols.push(ident.value.clone());
        }
        Expr::CompoundIdentifier(idents) => {
            if let Some(last) = idents.last() {
                cols.push(last.value.clone());
            }
        }
        Expr::BinaryOp { left, right, .. } => {
            cols.extend(extract_columns_from_expr(left));
            cols.extend(extract_columns_from_expr(right));
        }
        Expr::IsNull(e) | Expr::IsNotNull(e) => {
            cols.extend(extract_columns_from_expr(e));
        }
        Expr::Between {
            expr, low, high, ..
        } => {
            cols.extend(extract_columns_from_expr(expr));
            cols.extend(extract_columns_from_expr(low));
            cols.extend(extract_columns_from_expr(high));
        }
        Expr::InList { expr, list, .. } => {
            cols.extend(extract_columns_from_expr(expr));
            for item in list {
                cols.extend(extract_columns_from_expr(item));
            }
        }
        Expr::Function(func) => {
            if let FunctionArguments::List(arg_list) = &func.args {
                for arg in &arg_list.args {
                    if let FunctionArg::Unnamed(FunctionArgExpr::Expr(e)) = arg {
                        cols.extend(extract_columns_from_expr(e));
                    }
                }
            }
        }
        Expr::Nested(inner) => {
            cols.extend(extract_columns_from_expr(inner));
        }
        _ => {}
    }
    cols
}

fn extract_join_columns(join: &Join) -> Option<Vec<(String, String, String)>> {
    let constraint = match &join.join_operator {
        JoinOperator::Inner(c)
        | JoinOperator::LeftOuter(c)
        | JoinOperator::RightOuter(c)
        | JoinOperator::FullOuter(c) => c,
        _ => return None,
    };

    let right_table = extract_table_name_from_factor(&join.relation)?;

    match constraint {
        JoinConstraint::On(expr) => {
            let mut result = Vec::new();
            extract_join_eq_columns(expr, &right_table, &mut result);
            if result.is_empty() {
                None
            } else {
                Some(result)
            }
        }
        _ => None,
    }
}

fn extract_join_eq_columns(
    expr: &Expr,
    right_table: &str,
    result: &mut Vec<(String, String, String)>,
) {
    if let Expr::BinaryOp { left, op, right } = expr {
        use sqlparser::ast::BinaryOperator;
        if matches!(op, BinaryOperator::Eq) {
            if let (Some(left_col), Some(right_col)) =
                (expr_to_column_name(left), expr_to_column_name(right))
            {
                result.push((left_col, right_table.to_string(), right_col));
            }
        } else if matches!(op, BinaryOperator::And) {
            extract_join_eq_columns(left, right_table, result);
            extract_join_eq_columns(right, right_table, result);
        }
    }
}

fn resolve_column_table(
    col: &str,
    table_names: &[String],
    schema: &SchemaDefinition,
) -> Option<String> {
    for table in table_names {
        if let Some(cols) = schema.column_names(table)
            && cols.iter().any(|c| c.eq_ignore_ascii_case(col))
        {
            return Some(table.clone());
        }
    }
    table_names.first().cloned()
}

fn extract_tables_from_set_expr(set_expr: &SetExpr) -> Vec<String> {
    match set_expr {
        SetExpr::Select(select) => extract_table_names_from_select(select),
        _ => vec![],
    }
}

fn has_where_clause(set_expr: &SetExpr) -> bool {
    match set_expr {
        SetExpr::Select(select) => select.selection.is_some(),
        _ => false,
    }
}
