use std::collections::{HashMap, HashSet};

use pg_query::{NodeEnum, protobuf::Node};

use crate::core::schema::{QualifiedName, SchemaModel, TableDef};

/// Per-statement scope built from a single SQL statement's AST.
///
/// Tracks aliases, CTE names, and the FROM-clause table set so that column
/// references and aliased table references can be resolved correctly.
#[derive(Debug, Default)]
pub struct ScopeStack {
    /// Alias name → resolved `QualifiedName`. E.g. `"o"` → `("public", "orders")`.
    pub aliases: HashMap<String, QualifiedName>,
    /// CTE names defined in this statement's `WITH` clause.
    /// These shadow real table names — never emit schema refs for these.
    pub ctes: HashSet<String>,
    /// All tables/views in the FROM scope of this statement.
    /// Used for bare column reference resolution.
    pub from_scope: Vec<QualifiedName>,
}

// collects CTEs first so matching RangeVars are correctly shadowed
pub fn build_scope(root: &NodeEnum, model: &SchemaModel) -> ScopeStack {
    let mut stack = ScopeStack::default();
    collect_ctes(root, &mut stack);
    walk_for_range_vars(root, &mut stack, model);
    stack
}

/// Resolve a bare (unqualified) column name using the FROM scope.
///
/// Returns `None` when there are zero or multiple tables in scope (ambiguous).
pub fn resolve_bare_col(
    stack: &ScopeStack,
    col: &str,
    model: &SchemaModel,
) -> Option<(QualifiedName, String)> {
    if stack.from_scope.len() != 1 {
        tracing::debug!(
            col,
            tables = stack.from_scope.len(),
            "bare column is ambiguous or has no FROM scope, skipping"
        );
        return None;
    }
    let qn = &stack.from_scope[0];
    let table = lookup_table(model, qn)?;
    if table.columns.contains_key(col) {
        Some((qn.clone(), col.to_owned()))
    } else {
        tracing::debug!(col, table = %qn, "bare column not found in table");
        None
    }
}

/// Resolve a two-part `table_or_alias.col` reference.
///
/// Resolves the table prefix via aliases then schema; then looks up the column.
pub fn resolve_qualified_col(
    stack: &ScopeStack,
    table_or_alias: &str,
    col: &str,
    model: &SchemaModel,
) -> Option<(QualifiedName, String)> {
    let qn = resolve_table_or_alias(stack, table_or_alias, model)?;
    let table = lookup_table(model, &qn)?;
    if table.columns.contains_key(col) {
        Some((qn, col.to_owned()))
    } else {
        tracing::debug!(col, table = %qn, "column not found in resolved table");
        None
    }
}

/// Resolve a table-or-alias string to its canonical `QualifiedName`.
pub fn resolve_table_or_alias(
    stack: &ScopeStack,
    name: &str,
    model: &SchemaModel,
) -> Option<QualifiedName> {
    if let Some(qn) = stack.aliases.get(name) {
        return Some(qn.clone());
    }
    let candidate = QualifiedName::new(Some("public"), name);
    if model.tables.contains_key(&candidate) || model.views.contains_key(&candidate) {
        return Some(candidate);
    }
    tracing::debug!(name, "could not resolve table or alias");
    None
}

pub fn lookup_table<'a>(model: &'a SchemaModel, qn: &QualifiedName) -> Option<&'a TableDef> {
    model.tables.get(qn)
}

/// Resolve a `RangeVar`'s schema-qualified or bare relname against the model.
pub fn resolve_range_var(
    rv: &pg_query::protobuf::RangeVar,
    model: &SchemaModel,
) -> Option<QualifiedName> {
    if rv.relname.is_empty() {
        return None;
    }
    let schema = if rv.schemaname.is_empty() {
        None
    } else {
        Some(rv.schemaname.as_str())
    };
    let qn = QualifiedName::new(schema, &rv.relname);
    if model.tables.contains_key(&qn) || model.views.contains_key(&qn) {
        Some(qn)
    } else {
        tracing::debug!(
            rel = rv.relname,
            schema = rv.schemaname,
            "RangeVar not found in schema"
        );
        None
    }
}

fn collect_ctes(root: &NodeEnum, stack: &mut ScopeStack) {
    let with_clause = match root {
        NodeEnum::SelectStmt(s) => s.with_clause.as_ref(),
        NodeEnum::InsertStmt(s) => s.with_clause.as_ref(),
        NodeEnum::UpdateStmt(s) => s.with_clause.as_ref(),
        NodeEnum::DeleteStmt(s) => s.with_clause.as_ref(),
        _ => None,
    };
    let Some(wc) = with_clause else { return };
    for node in &wc.ctes {
        if let Some(NodeEnum::CommonTableExpr(cte)) = node.node.as_ref() {
            stack.ctes.insert(cte.ctename.clone());
        }
    }
}

fn walk_for_range_vars(node: &NodeEnum, stack: &mut ScopeStack, model: &SchemaModel) {
    if let NodeEnum::RangeVar(rv) = node {
        register_range_var(rv, stack, model);
        return;
    }
    walk_range_var_children(node, stack, model);
}

fn walk_range_var_children(node: &NodeEnum, stack: &mut ScopeStack, model: &SchemaModel) {
    match node {
        NodeEnum::SelectStmt(s) => {
            walk_rv_node_slice(&s.from_clause, stack, model);
            if let Some(wc) = &s.with_clause {
                walk_rv_node_slice(&wc.ctes, stack, model);
            }
            if let Some(l) = s.larg.as_ref() {
                walk_for_range_vars(&NodeEnum::SelectStmt(l.clone()), stack, model);
            }
            if let Some(r) = s.rarg.as_ref() {
                walk_for_range_vars(&NodeEnum::SelectStmt(r.clone()), stack, model);
            }
        }
        NodeEnum::JoinExpr(j) => {
            if let Some(l) = j.larg.as_deref() {
                walk_rv_single(l, stack, model);
            }
            if let Some(r) = j.rarg.as_deref() {
                walk_rv_single(r, stack, model);
            }
        }
        NodeEnum::CommonTableExpr(c) => {
            if let Some(q) = c.ctequery.as_deref() {
                walk_rv_single(q, stack, model);
            }
        }
        NodeEnum::RangeSubselect(rs) => {
            if let Some(q) = rs.subquery.as_deref() {
                walk_rv_single(q, stack, model);
            }
        }
        _ => {}
    }
}

pub(super) fn walk_rv_node_slice(nodes: &[Node], stack: &mut ScopeStack, model: &SchemaModel) {
    for n in nodes {
        walk_rv_single(n, stack, model);
    }
}

pub(super) fn walk_rv_single(node: &Node, stack: &mut ScopeStack, model: &SchemaModel) {
    if let Some(inner) = node.node.as_ref() {
        walk_for_range_vars(inner, stack, model);
    }
}

fn register_range_var(
    rv: &pg_query::protobuf::RangeVar,
    stack: &mut ScopeStack,
    model: &SchemaModel,
) {
    if rv.relname.is_empty() || stack.ctes.contains(&rv.relname) {
        return;
    }
    let schema = if rv.schemaname.is_empty() {
        None
    } else {
        Some(rv.schemaname.as_str())
    };
    let qn = QualifiedName::new(schema, &rv.relname);
    if !model.tables.contains_key(&qn) && !model.views.contains_key(&qn) {
        tracing::debug!(
            rel = rv.relname,
            "RangeVar not in schema, skipping from scope"
        );
        return;
    }
    if let Some(alias) = rv.alias.as_ref() {
        if !alias.aliasname.is_empty() {
            stack.aliases.insert(alias.aliasname.clone(), qn.clone());
        }
    }
    stack.from_scope.push(qn);
}
