use std::collections::HashMap;
use std::ops::Range;
use std::rc::Rc;
use std::sync::Arc;

use antlr_rust::parser_rule_context::ParserRuleContext;
use antlr_rust::tree::{ParseTree, TerminalNode};
use hamelin_lib::antlr::completion_interval;
use hamelin_lib::antlr::cuddled_completion_interval;
use hamelin_lib::antlr::fit_cursor;
use hamelin_lib::antlr::hamelinparser::AggCommandContextAttrs;
use hamelin_lib::antlr::hamelinparser::ArrayLiteralContextAttrs;
use hamelin_lib::antlr::hamelinparser::AssignmentClauseContextAll;
use hamelin_lib::antlr::hamelinparser::AssignmentClauseContextAttrs;
use hamelin_lib::antlr::hamelinparser::AssignmentContextAttrs;
use hamelin_lib::antlr::hamelinparser::CastContextAttrs;
use hamelin_lib::antlr::hamelinparser::CommandContextAll;
use hamelin_lib::antlr::hamelinparser::DefRhsContextAll;
use hamelin_lib::antlr::hamelinparser::DefRhsExpressionContextAttrs;
use hamelin_lib::antlr::hamelinparser::DefRhsPipelineContextAttrs;
use hamelin_lib::antlr::hamelinparser::DefStatementContextAttrs;
use hamelin_lib::antlr::hamelinparser::DistinctCommandContextAttrs;
use hamelin_lib::antlr::hamelinparser::DropCommandContextAttrs;
use hamelin_lib::antlr::hamelinparser::ExplodeCommandContextAttrs;
use hamelin_lib::antlr::hamelinparser::ExpressionContextAll;
use hamelin_lib::antlr::hamelinparser::ExpressionEOFContextAttrs;
use hamelin_lib::antlr::hamelinparser::ExpressionQueryMainContextAttrs;
use hamelin_lib::antlr::hamelinparser::FieldReferenceAltContextAttrs;
use hamelin_lib::antlr::hamelinparser::FieldReferenceContextAll;
use hamelin_lib::antlr::hamelinparser::FieldReferenceContextAttrs;
use hamelin_lib::antlr::hamelinparser::FromClauseContextAll;
use hamelin_lib::antlr::hamelinparser::FromClauseContextAttrs;
use hamelin_lib::antlr::hamelinparser::FromCommandContextAttrs;
use hamelin_lib::antlr::hamelinparser::FunctionCallContextAttrs;
use hamelin_lib::antlr::hamelinparser::GroupClauseContextAttrs;
use hamelin_lib::antlr::hamelinparser::HamelinParserContextType;
use hamelin_lib::antlr::hamelinparser::JoinCommandContextAttrs;
use hamelin_lib::antlr::hamelinparser::MatchCommandContextAttrs;
use hamelin_lib::antlr::hamelinparser::NamedArgumentContextAttrs;
use hamelin_lib::antlr::hamelinparser::ParenthesizedExpressionContextAttrs;
use hamelin_lib::antlr::hamelinparser::PipelineAltContextAttrs;
use hamelin_lib::antlr::hamelinparser::PipelineContextAll;
use hamelin_lib::antlr::hamelinparser::PipelineQueryMainContextAttrs;
use hamelin_lib::antlr::hamelinparser::PositionalArgumentContextAttrs;
use hamelin_lib::antlr::hamelinparser::QueryContextAll;
use hamelin_lib::antlr::hamelinparser::QueryEOFContextAll;
use hamelin_lib::antlr::hamelinparser::QueryEOFContextAttrs;
use hamelin_lib::antlr::hamelinparser::QueryMainContextAll;
use hamelin_lib::antlr::hamelinparser::QueryWithDefsContextAttrs;
use hamelin_lib::antlr::hamelinparser::QueryWithoutDefsContextAttrs;
use hamelin_lib::antlr::hamelinparser::SelectCommandContextAttrs;
use hamelin_lib::antlr::hamelinparser::SelectionContextAll;
use hamelin_lib::antlr::hamelinparser::SetCommandContextAttrs;
use hamelin_lib::antlr::hamelinparser::SimpleIdentifierContextAll;
use hamelin_lib::antlr::hamelinparser::SortCommandContextAttrs;
use hamelin_lib::antlr::hamelinparser::SortExpressionContextAttrs;
use hamelin_lib::antlr::hamelinparser::StructLiteralContextAttrs;
use hamelin_lib::antlr::hamelinparser::SuppressCommandContextAttrs;
use hamelin_lib::antlr::hamelinparser::TableAliasContextAttrs;
use hamelin_lib::antlr::hamelinparser::TablePathContextAttrs;
use hamelin_lib::antlr::hamelinparser::TablePathSegmentContextAttrs;
use hamelin_lib::antlr::hamelinparser::TableReferenceContextAll;
use hamelin_lib::antlr::hamelinparser::TableReferenceContextAttrs;
use hamelin_lib::antlr::hamelinparser::TemplatePairContextAttrs;
use hamelin_lib::antlr::hamelinparser::TemplateParameterExprContext;
use hamelin_lib::antlr::hamelinparser::TemplateParameterExprContextAttrs;
use hamelin_lib::antlr::hamelinparser::TsTruncContextAttrs;
use hamelin_lib::antlr::hamelinparser::TupleLiteralContextAttrs;
use hamelin_lib::antlr::hamelinparser::UnaryPostfixOperatorContextAttrs;
use hamelin_lib::antlr::hamelinparser::UnaryPrefixOperatorContextAttrs;
use hamelin_lib::antlr::hamelinparser::UnionCommandContextAttrs;
use hamelin_lib::antlr::hamelinparser::UnnestCommandContextAttrs;
use hamelin_lib::antlr::hamelinparser::WhereCommandContextAttrs;
use hamelin_lib::antlr::hamelinparser::WindowCommandContextAttrs;
use hamelin_lib::antlr::hamelinparser::WithinCommandContextAttrs;
use hamelin_lib::antlr::hamelinparser::{LimitCommandContextAttrs, TrimstringsCommandContextAttrs};
use hamelin_lib::antlr::interval;
use hamelin_lib::antlr::resilient_parse_query;
use hamelin_lib::catalog::Catalog;
use hamelin_lib::catalog::CatalogProvider;
use hamelin_lib::completion::Completion;
use hamelin_lib::completion::CompletionItem;
use hamelin_lib::func::registry::FunctionRegistry;
use hamelin_lib::parser::make_hamelin_parser_from_input;
use hamelin_lib::provider::EnvironmentProvider;
use hamelin_lib::tree::ast::context::FromCst;
use hamelin_lib::tree::ast::context::ParseContext;
use hamelin_lib::tree::ast::expression::Expression as AstExpression;
use hamelin_lib::tree::ast::identifier::{Identifier, ParsedSimpleIdentifier};
use hamelin_lib::tree::ast::query::Query;
use hamelin_lib::tree::ast::ParseWithErrors;
use hamelin_lib::tree::options::{
    ExpressionTypeCheckOptions, TemplateParameterKind, TypeCheckOptions,
};
use hamelin_lib::tree::typed_ast::query::TypedStatement;
use hamelin_lib::type_check_expression;
use hamelin_lib::type_check_with_options;
use hamelin_lib::types::Type;
use wasm_bindgen::prelude::wasm_bindgen;
use wasm_bindgen::JsValue;

use crate::wasm_template_params::parse_template_parameters_js;

fn template_parameter_names(
    specs: &Option<Arc<HashMap<String, TemplateParameterKind>>>,
) -> Vec<String> {
    let mut names = specs
        .as_ref()
        .map(|m| m.keys().cloned().collect::<Vec<String>>())
        .unwrap_or_default();
    names.sort();
    names
}

fn table_path_template_parameter_names(
    specs: &Option<Arc<HashMap<String, TemplateParameterKind>>>,
) -> Vec<String> {
    let mut names: Vec<String> = specs
        .as_ref()
        .map(|m| {
            m.iter()
                .filter(|(_, kind)| matches!(kind, TemplateParameterKind::IdentifierFragment(_)))
                .map(|(k, _)| k.clone())
                .collect()
        })
        .unwrap_or_default();
    names.sort();
    names
}

struct TemplateParameterCompletionTarget {
    prefix: String,
    replace_range: Range<usize>,
}

/// `start_char` is a 0-based **character** (Unicode scalar) index into `s`.
fn slice_from_char(s: &str, start_char: usize) -> &str {
    if start_char == 0 {
        return s;
    }
    let mut idx = 0usize;
    for (byte_idx, _) in s.char_indices() {
        if idx == start_char {
            return &s[byte_idx..];
        }
        idx += 1;
    }
    ""
}

/// If `}` is already there, insert the name only; otherwise append `}`.
fn template_param_apply_string(name: &str, close_brace_already_present: bool) -> String {
    if close_brace_already_present {
        name.to_string()
    } else {
        format!("{}{}", name, "}")
    }
}

fn build_template_parameter_completion(
    target: TemplateParameterCompletionTarget,
    param_names: &[String],
    full_query: &str,
) -> Option<Completion> {
    if param_names.is_empty() {
        return None;
    }
    let replace_end_exclusive = target.replace_range.end;
    let close_brace_already_present =
        slice_from_char(full_query, replace_end_exclusive).starts_with('}');
    let mut completion = Completion::new_exclusive(target.replace_range);
    completion.filter(false);
    for name in param_names {
        if target.prefix.is_empty() || name.starts_with(&target.prefix) {
            let apply = template_param_apply_string(name, close_brace_already_present);
            completion.add_item(CompletionItem::new(name.as_str()).with_apply_string(apply));
        }
    }
    if completion.items.is_empty() {
        None
    } else {
        Some(completion)
    }
}

fn template_completion_target_from_cst_parts(
    simple_identifier: Option<Rc<SimpleIdentifierContextAll<'static>>>,
    template_open: Option<Rc<TerminalNode<'static, HamelinParserContextType>>>,
) -> Option<TemplateParameterCompletionTarget> {
    if let Some(si) = simple_identifier {
        let r = interval(si.as_ref());
        return Some(TemplateParameterCompletionTarget {
            prefix: si.get_text(),
            replace_range: *r.start()..*r.end() + 1,
        });
    }
    let op = template_open?;
    let end_after_open = *interval(op.as_ref()).end() + 1;
    Some(TemplateParameterCompletionTarget {
        prefix: String::new(),
        replace_range: end_after_open..end_after_open,
    })
}

/// Trailing text after the successfully parsed `query` subtree is often error tokens (for example
/// `FROM a.${name` with no closing `}`). Complete template names when the cursor is there.
fn template_parameter_completion_in_unlexed_suffix(
    full_query: &str,
    at: usize,
    parsed_query_end_char_inclusive: usize,
    param_names: &[String],
) -> Option<Completion> {
    if param_names.is_empty() {
        return None;
    }
    let total_chars = full_query.chars().count();
    if at > total_chars {
        return None;
    }
    if at <= parsed_query_end_char_inclusive || parsed_query_end_char_inclusive + 1 >= total_chars {
        return None;
    }
    let fragment_start = parsed_query_end_char_inclusive + 1;
    let rel_at = at.saturating_sub(fragment_start);
    let fragment: String = full_query.chars().skip(fragment_start).collect();
    let target = template_completion_target_from_fragment(&fragment, rel_at, fragment_start)?;
    build_template_parameter_completion(target, param_names, full_query)
}

fn template_completion_target_from_fragment(
    fragment: &str,
    rel_at: usize,
    fragment_start_char: usize,
) -> Option<TemplateParameterCompletionTarget> {
    let frag_char_len = fragment.chars().count();
    if rel_at > frag_char_len {
        return None;
    }
    let before: String = fragment.chars().take(rel_at).collect();
    let open_byte = before.rfind("${")?;
    let open_char_in_fragment = before[..open_byte].chars().count();
    let after_open: String = before.chars().skip(open_char_in_fragment + 2).collect();
    if after_open.contains('}') || after_open.contains('{') {
        return None;
    }
    let name_start_in_fragment = open_char_in_fragment + 2;
    if rel_at > name_start_in_fragment {
        let prefix: String = fragment
            .chars()
            .skip(name_start_in_fragment)
            .take(rel_at.saturating_sub(name_start_in_fragment))
            .collect();
        if !prefix
            .chars()
            .all(|c| c.is_ascii_alphanumeric() || c == '_')
        {
            return None;
        }
        let abs_start = fragment_start_char + name_start_in_fragment;
        let abs_end = fragment_start_char + (rel_at - 1);
        return Some(TemplateParameterCompletionTarget {
            prefix,
            replace_range: abs_start..abs_end + 1,
        });
    }
    let p = fragment_start_char + name_start_in_fragment;
    Some(TemplateParameterCompletionTarget {
        prefix: String::new(),
        replace_range: p..p,
    })
}

/// Bad `template_parameters` JSON is reported as a `JsValue` exception (matching
/// [`crate::compile::compile_query`]) so callers don't silently lose completions due to a malformed
/// configuration. Catalog and parse failures still resolve to `None`, since completions degrade
/// gracefully when the query itself is incomplete.
#[wasm_bindgen]
pub fn get_semantic_completions(
    query: String,
    at: Option<usize>,
    catalog: Catalog,
    template_parameters: Option<JsValue>,
) -> Result<Option<Completion>, JsValue> {
    let template_specs = parse_template_parameters_js(template_parameters)?;

    let Some(catalog_provider) = CatalogProvider::try_from(catalog).ok().map(Arc::new) else {
        return Ok(None);
    };
    let char_count = query.chars().count();
    let at = at.unwrap_or(char_count).min(char_count);

    let Ok((query_eof, _)) = resilient_parse_query(query.clone()) else {
        return Ok(None);
    };

    let mut parse_ctx = ParseContext::new();
    let Some(query_cst) = query_eof.query() else {
        return Ok(None);
    };
    let ast = Query::from_cst_with_context(query_cst.clone(), &mut parse_ctx);

    let typed_statement = type_check_with_options::<Query>(
        Arc::new(ast),
        TypeCheckOptions::builder()
            .provider(catalog_provider.clone())
            .maybe_template_parameters(template_specs.clone())
            .build(),
    )
    .output;

    let mut walker = CompletionWalker {
        at,
        query: query.clone(),
        typed_statement,
        provider: catalog_provider,
        registry: Arc::new(FunctionRegistry::default()),
        completion: None,
        template_specs,
    };
    walker.traverse_query(query_cst);
    walker.append_eof_trailing_template_param_completions(&query_eof);
    Ok(walker.completion.map(|mut c| {
        c.at = fit_cursor(&c.at, &query);
        c
    }))
}

struct CompletionWalker {
    at: usize,
    query: String,
    typed_statement: TypedStatement,
    provider: Arc<dyn EnvironmentProvider>,
    registry: Arc<FunctionRegistry>,
    completion: Option<Completion>,
    template_specs: Option<Arc<HashMap<String, TemplateParameterKind>>>,
}

impl CompletionWalker {
    fn env(&self) -> Option<hamelin_lib::types::struct_type::Struct> {
        self.typed_statement
            .environment_at_offset(self.at)
            .map(|env| env.as_struct().clone())
    }

    fn append_eof_trailing_template_param_completions(
        &mut self,
        query_eof: &QueryEOFContextAll<'static>,
    ) {
        if self.completion.is_some() {
            return;
        }
        let names = template_parameter_names(&self.template_specs);
        let Some(q) = query_eof.query() else {
            return;
        };
        let parsed_end = *interval(q.as_ref()).end();
        if let Some(c) = template_parameter_completion_in_unlexed_suffix(
            &self.query,
            self.at,
            parsed_end,
            &names,
        ) {
            self.completion = Some(c);
        }
    }

    fn traverse_query(&mut self, query: Rc<QueryContextAll<'static>>) {
        match query.as_ref() {
            QueryContextAll::QueryWithDefsContext(ctx) => {
                for def in ctx.defStatement_all() {
                    self.traverse_def_rhs(def.defRhs());
                }

                if let Some(main) = ctx.queryMain() {
                    self.traverse_query_main(main);
                }
            }
            QueryContextAll::QueryWithoutDefsContext(ctx) => {
                if let Some(main) = ctx.queryMain() {
                    self.traverse_query_main(main);
                }
            }
            QueryContextAll::Error(_) => {}
        }
    }

    fn traverse_query_main(&mut self, main: Rc<QueryMainContextAll<'static>>) {
        match main.as_ref() {
            QueryMainContextAll::PipelineQueryMainContext(ctx) => {
                if let Some(pipeline) = ctx.pipeline() {
                    self.traverse_pipeline(pipeline);
                }
            }
            QueryMainContextAll::ExpressionQueryMainContext(ctx) => {
                if let Some(expression) = ctx.expression() {
                    self.traverse_expression(expression);
                }
            }
            QueryMainContextAll::Error(_) => {}
        }
    }

    fn traverse_def_rhs(&mut self, def_rhs: Option<Rc<DefRhsContextAll<'static>>>) {
        let Some(rhs) = def_rhs else {
            return;
        };
        match rhs.as_ref() {
            DefRhsContextAll::DefRhsPipelineContext(pctx) => {
                if let Some(pipeline) = pctx.pipeline() {
                    self.traverse_pipeline(pipeline);
                }
            }
            DefRhsContextAll::DefRhsExpressionContext(ectx) => {
                if let Some(expression) = ectx.expression() {
                    self.traverse_expression(expression);
                }
            }
            DefRhsContextAll::Error(_) => {}
        }
    }

    fn traverse_pipeline(&mut self, pipeline: Rc<PipelineContextAll<'static>>) {
        if let PipelineContextAll::PipelineAltContext(ctx) = pipeline.as_ref() {
            for command in ctx.command_all() {
                self.traverse_command(command);
            }
        }
    }

    fn traverse_command(&mut self, command: Rc<CommandContextAll<'static>>) {
        match command.as_ref() {
            CommandContextAll::FromCommandContext(ctx) => {
                for from_clause in ctx.fromClause_all() {
                    self.append_from_clause_completions(from_clause);
                }
            }
            CommandContextAll::UnionCommandContext(ctx) => {
                for from_clause in ctx.fromClause_all() {
                    self.append_from_clause_completions(from_clause);
                }
            }
            CommandContextAll::JoinCommandContext(ctx) => {
                if let Some(from_clause) = ctx.fromClause() {
                    self.append_from_clause_completions(from_clause);
                }
                if let Some(on) = ctx.on.as_ref() {
                    self.traverse_expression(on.clone());
                }
            }
            CommandContextAll::WhereCommandContext(ctx) => {
                if let Some(expression) = ctx.expression() {
                    self.traverse_expression(expression);
                }
            }
            CommandContextAll::SetCommandContext(ctx) => {
                for assignment in ctx.assignment_all() {
                    if let Some(expression) = assignment.expression() {
                        self.traverse_expression(expression);
                    }
                }
            }
            CommandContextAll::SelectCommandContext(ctx) => {
                for clause in ctx.assignmentClause_all() {
                    self.append_assignment_clause_completions(clause.clone());
                    if let Some(expression) = clause.expression() {
                        self.traverse_expression(expression);
                    }
                    if let Some(assignment) = clause.assignment() {
                        if let Some(expression) = assignment.expression() {
                            self.traverse_expression(expression);
                        }
                    }
                }
            }
            CommandContextAll::AggCommandContext(ctx) => {
                for clause in ctx.assignmentClause_all() {
                    self.append_assignment_clause_completions(clause.clone());
                    if let Some(expression) = clause.expression() {
                        self.traverse_expression(expression);
                    }
                    if let Some(assignment) = clause.assignment() {
                        if let Some(expression) = assignment.expression() {
                            self.traverse_expression(expression);
                        }
                    }
                }
                for group_clause in ctx.groupClause_all() {
                    if let Some(clause) = group_clause.assignmentClause() {
                        self.append_assignment_clause_completions(clause.clone());
                        if let Some(expression) = clause.expression() {
                            self.traverse_expression(expression);
                        }
                        if let Some(assignment) = clause.assignment() {
                            if let Some(expression) = assignment.expression() {
                                self.traverse_expression(expression);
                            }
                        }
                    }
                }
                for sort_expression in ctx.sortExpression_all() {
                    if let Some(expression) = sort_expression.expression() {
                        self.traverse_expression(expression);
                    }
                }
            }
            CommandContextAll::WindowCommandContext(ctx) => {
                for clause in ctx.assignmentClause_all() {
                    self.append_assignment_clause_completions(clause.clone());
                    if let Some(expression) = clause.expression() {
                        self.traverse_expression(expression);
                    }
                    if let Some(assignment) = clause.assignment() {
                        if let Some(expression) = assignment.expression() {
                            self.traverse_expression(expression);
                        }
                    }
                }
                for group_clause in ctx.groupClause_all() {
                    if let Some(clause) = group_clause.assignmentClause() {
                        self.append_assignment_clause_completions(clause.clone());
                        if let Some(expression) = clause.expression() {
                            self.traverse_expression(expression);
                        }
                        if let Some(assignment) = clause.assignment() {
                            if let Some(expression) = assignment.expression() {
                                self.traverse_expression(expression);
                            }
                        }
                    }
                }
                for sort_expression in ctx.sortExpression_all() {
                    if let Some(expression) = sort_expression.expression() {
                        self.traverse_expression(expression);
                    }
                }
                if let Some(within) = ctx.within.as_ref() {
                    self.traverse_expression(within.clone());
                }
            }
            CommandContextAll::SortCommandContext(ctx) => {
                for sort_expression in ctx.sortExpression_all() {
                    if let Some(expression) = sort_expression.expression() {
                        self.traverse_expression(expression);
                    }
                }
            }
            CommandContextAll::LimitCommandContext(ctx) => {
                if let Some(expression) = ctx.expression() {
                    self.traverse_expression(expression);
                }
            }
            CommandContextAll::TrimstringsCommandContext(ctx) => {
                if let Some(expression) = ctx.expression() {
                    self.traverse_expression(expression);
                }
            }
            CommandContextAll::WithinCommandContext(ctx) => {
                if let Some(expression) = ctx.expression() {
                    self.traverse_expression(expression);
                }
            }
            CommandContextAll::UnnestCommandContext(ctx) => {
                if let Some(expression) = ctx.expression() {
                    self.traverse_expression(expression);
                }
            }
            CommandContextAll::ExplodeCommandContext(ctx) => {
                if let Some(clause) = ctx.assignmentClause() {
                    self.append_assignment_clause_completions(clause.clone());
                    if let Some(expression) = clause.expression() {
                        self.traverse_expression(expression);
                    }
                    if let Some(assignment) = clause.assignment() {
                        if let Some(expression) = assignment.expression() {
                            self.traverse_expression(expression);
                        }
                    }
                }
            }
            CommandContextAll::ParseCommandContext(ctx) => {
                if let Some(src) = ctx.src.as_ref() {
                    self.traverse_expression(src.clone());
                }
            }
            CommandContextAll::MatchCommandContext(ctx) => {
                for clause in ctx.assignmentClause_all() {
                    self.append_assignment_clause_completions(clause.clone());
                    if let Some(expression) = clause.expression() {
                        self.traverse_expression(expression);
                    }
                    if let Some(assignment) = clause.assignment() {
                        if let Some(expression) = assignment.expression() {
                            self.traverse_expression(expression);
                        }
                    }
                }
                for group_clause in ctx.groupClause_all() {
                    if let Some(clause) = group_clause.assignmentClause() {
                        self.append_assignment_clause_completions(clause.clone());
                        if let Some(expression) = clause.expression() {
                            self.traverse_expression(expression);
                        }
                        if let Some(assignment) = clause.assignment() {
                            if let Some(expression) = assignment.expression() {
                                self.traverse_expression(expression);
                            }
                        }
                    }
                }
                for sort_expression in ctx.sortExpression_all() {
                    if let Some(expression) = sort_expression.expression() {
                        self.traverse_expression(expression);
                    }
                }
                if let Some(within) = ctx.within.as_ref() {
                    self.traverse_expression(within.clone());
                }
            }
            CommandContextAll::DropCommandContext(ctx) => {
                for selection in ctx.selection_all() {
                    self.append_selection_completions(selection);
                }
            }
            CommandContextAll::DistinctCommandContext(ctx) => {
                for group_clause in ctx.groupClause_all() {
                    if let Some(clause) = group_clause.assignmentClause() {
                        self.append_assignment_clause_completions(clause.clone());
                        if let Some(expression) = clause.expression() {
                            self.traverse_expression(expression);
                        }
                        if let Some(assignment) = clause.assignment() {
                            if let Some(expression) = assignment.expression() {
                                self.traverse_expression(expression);
                            }
                        }
                    }
                }
            }
            CommandContextAll::SuppressCommandContext(ctx) => {
                if let Some(expression) = ctx.expression() {
                    self.traverse_expression(expression);
                }
                for group_clause in ctx.groupClause_all() {
                    if let Some(clause) = group_clause.assignmentClause() {
                        self.append_assignment_clause_completions(clause.clone());
                        if let Some(expression) = clause.expression() {
                            self.traverse_expression(expression);
                        }
                        if let Some(assignment) = clause.assignment() {
                            if let Some(expression) = assignment.expression() {
                                self.traverse_expression(expression);
                            }
                        }
                    }
                }
            }
            CommandContextAll::AppendCommandContext(_)
            | CommandContextAll::NestCommandContext(_)
            | CommandContextAll::Error(_) => {}
        }
    }

    fn traverse_expression(&mut self, expression: Rc<ExpressionContextAll<'static>>) {
        match expression.as_ref() {
            ExpressionContextAll::FieldReferenceAltContext(ctx) => {
                self.append_field_ref_completions(expression.clone(), ctx.fieldReference());
            }
            ExpressionContextAll::FieldLookupContext(ctx) => {
                if let (Some(left), Some(right)) = (ctx.left.clone(), ctx.right.clone()) {
                    self.append_field_lookup_completions(left, right);
                }
            }
            ExpressionContextAll::ParenthesizedExpressionContext(ctx) => {
                if let Some(inner) = ctx.expression() {
                    self.traverse_expression(inner);
                }
            }
            ExpressionContextAll::StructLiteralContext(ctx) => {
                for inner in ctx.expression_all() {
                    self.traverse_expression(inner);
                }
            }
            ExpressionContextAll::ArrayLiteralContext(ctx) => {
                for inner in ctx.expression_all() {
                    self.traverse_expression(inner);
                }
            }
            ExpressionContextAll::TupleLiteralContext(ctx) => {
                for inner in ctx.expression_all() {
                    self.traverse_expression(inner);
                }
            }
            ExpressionContextAll::PairLiteralContext(ctx) => {
                if let Some(left) = ctx.left.clone() {
                    self.traverse_expression(left);
                }
                if let Some(right) = ctx.right.clone() {
                    self.traverse_expression(right);
                }
            }
            ExpressionContextAll::IndexAccessContext(ctx) => {
                if let Some(value) = ctx.value.clone() {
                    self.traverse_expression(value);
                }
                if let Some(index) = ctx.index.clone() {
                    self.traverse_expression(index);
                }
            }
            ExpressionContextAll::TsTruncContext(ctx) => {
                if let Some(inner) = ctx.expression() {
                    self.traverse_expression(inner);
                }
            }
            ExpressionContextAll::CastContext(ctx) => {
                if let Some(inner) = ctx.expression() {
                    self.traverse_expression(inner);
                }
            }
            ExpressionContextAll::BinaryOperatorContext(ctx) => {
                if let Some(left) = ctx.left.clone() {
                    self.traverse_expression(left);
                }
                if let Some(right) = ctx.right.clone() {
                    self.traverse_expression(right);
                }
            }
            ExpressionContextAll::FunctionCallContext(ctx) => {
                for positional in ctx.positionalArgument_all() {
                    if let Some(expression) = positional.expression() {
                        self.traverse_expression(expression);
                    }
                }
                for named in ctx.namedArgument_all() {
                    if let Some(expression) = named.expression() {
                        self.traverse_expression(expression);
                    }
                }
            }
            ExpressionContextAll::UnaryPrefixOperatorContext(ctx) => {
                if let Some(inner) = ctx.expression() {
                    self.traverse_expression(inner);
                }
            }
            ExpressionContextAll::UnaryPostfixOperatorContext(ctx) => {
                if let Some(inner) = ctx.expression() {
                    self.traverse_expression(inner);
                }
            }
            ExpressionContextAll::TemplateParameterExprContext(ctx) => {
                self.append_template_parameter_expr_completions(expression.clone(), ctx);
            }
            ExpressionContextAll::Error(ctx) => {
                self.append_expression_error_completions(ctx);
            }
            _ => {}
        }
    }

    fn append_template_parameter_expr_completions(
        &mut self,
        expression: Rc<ExpressionContextAll<'static>>,
        ctx: &TemplateParameterExprContext<'static>,
    ) {
        if self.completion.is_some() {
            return;
        }
        let range = completion_interval(expression.as_ref());
        if !range.contains(&self.at) {
            return;
        }
        let names = template_parameter_names(&self.template_specs);
        if names.is_empty() {
            return;
        }
        let Some(target) =
            template_completion_target_from_cst_parts(ctx.simpleIdentifier(), ctx.TEMPLATE_OPEN())
        else {
            return;
        };
        if let Some(c) = build_template_parameter_completion(target, &names, &self.query) {
            self.completion = Some(c);
        }
    }

    fn append_from_clause_completions(&mut self, from_clause: Rc<FromClauseContextAll<'static>>) {
        if self.completion.is_some() {
            return;
        }
        let table_reference = from_clause
            .tableReference()
            .or_else(|| from_clause.tableAlias().and_then(|ta| ta.tableReference()));

        if let Some(tr) = &table_reference {
            self.append_table_path_template_param_completions(tr);
        }

        if self.completion.is_some() {
            return;
        }

        match table_reference {
            Some(cur) if completion_interval(cur.as_ref()).contains(&self.at) => {
                self.from_completion_append(cur.as_ref());
            }
            None if completion_interval(from_clause.as_ref()).contains(&self.at) => {
                self.from_completion_append(from_clause.as_ref());
            }
            _ => {}
        }
    }

    fn append_table_path_template_param_completions(
        &mut self,
        tr: &TableReferenceContextAll<'static>,
    ) {
        if self.completion.is_some() {
            return;
        }
        let names = table_path_template_parameter_names(&self.template_specs);
        if names.is_empty() {
            return;
        }
        let Some(tp) = tr.tablePath() else {
            return;
        };
        for seg in tp.tablePathSegment_all() {
            for part in seg.templatePair_all() {
                if part.TEMPLATE_OPEN().is_none() {
                    continue;
                }
                if !completion_interval(part.as_ref()).contains(&self.at) {
                    continue;
                }
                let Some(target) = template_completion_target_from_cst_parts(
                    part.simpleIdentifier(),
                    part.TEMPLATE_OPEN(),
                ) else {
                    continue;
                };
                if let Some(c) = build_template_parameter_completion(target, &names, &self.query) {
                    self.completion = Some(c);
                    return;
                }
            }
        }
    }

    fn from_completion_append<'ctx, PRC: ParserRuleContext<'ctx>>(&mut self, cur: &PRC) {
        if self.completion.is_some() {
            return;
        }

        let insert_interval = interval(cur);
        let mut completion = Completion::new(insert_interval);
        completion.filter(false);
        for dataset in self
            .typed_statement
            .pipeline_defs
            .iter()
            .map(|pd| pd.name.to_string())
            .chain(
                self.provider
                    .reflect_datasets()
                    .unwrap_or_default()
                    .into_iter()
                    .map(|dataset| dataset.to_string()),
            )
        {
            if dataset.contains(cur.get_text().as_str()) {
                completion.add_item(CompletionItem::new(dataset));
            }
        }
        self.completion = Some(completion);
    }

    fn append_selection_completions(&mut self, selection: Rc<SelectionContextAll<'static>>) {
        if self.completion.is_some() {
            return;
        }
        let range = completion_interval(selection.as_ref());
        if !range.contains(&self.at) {
            return;
        }
        let Some(env) = self.env() else {
            return;
        };

        let insert_interval = interval(selection.as_ref());
        let mut completion = Completion::new(insert_interval);
        completion.add_items(env.autocomplete_suggestions(false));
        self.completion = Some(completion);
    }

    fn append_field_ref_completions(
        &mut self,
        parent_tree: Rc<ExpressionContextAll<'static>>,
        field_ref_tree: Option<Rc<FieldReferenceContextAll<'static>>>,
    ) {
        if self.completion.is_some() {
            return;
        }

        let range = completion_interval(parent_tree.as_ref());
        let Some(field_ref) = field_ref_tree else {
            return;
        };
        if !range.contains(&self.at) {
            return;
        }

        let Some(env) = self.env() else {
            return;
        };

        let insert_interval = interval(field_ref.as_ref());
        let mut completion = Completion::new(insert_interval);

        if *range.start() == self.at {
            completion.filter(false);
            completion.add_items(env.autocomplete_suggestions(false));
            completion.add_items(self.registry.autocomplete_suggestions());
            self.completion = Some(completion);
            return;
        }

        let mut parse_ctx = ParseContext::new();
        let maybe_curr_ident = field_ref.simpleIdentifier().and_then(|si| {
            ParsedSimpleIdentifier::from_cst_with_context(si.as_ref(), &mut parse_ctx)
                .valid()
                .ok()
        });
        if let Some(curr_ident) = maybe_curr_ident {
            completion.add_items(env.nested_autocomplete_suggestions(&curr_ident, false));
            completion.add_items(self.registry.autocomplete_suggestions());
        }
        self.completion = Some(completion);
    }

    fn append_field_lookup_completions(
        &mut self,
        left: Rc<ExpressionContextAll<'static>>,
        right: Rc<SimpleIdentifierContextAll<'static>>,
    ) {
        if self.completion.is_some() {
            return;
        }
        let range = cuddled_completion_interval(right.as_ref());
        if !range.contains(&self.at) {
            return;
        }

        let Some(expression_type) = self.infer_expression_type(left) else {
            return;
        };

        let right_ident = Identifier::parse(right.get_text()).ok();

        let insert_interval = if right_ident.is_some() {
            let int = interval(right.as_ref());
            int.start().saturating_sub(1)..=*int.end()
        } else {
            self.at.saturating_sub(1)..=self.at.saturating_sub(1)
        };

        let mut completion = Completion::new(insert_interval);
        completion.filter(false);
        match expression_type {
            Type::Struct(struct_type) => {
                if let Some(items) = field_lookup_items_for_type(&struct_type, &right_ident, true) {
                    completion.add_items(items);
                    self.completion = Some(completion);
                }
            }
            Type::Array(array) => {
                if let Type::Struct(struct_type) = array.element_type.as_ref() {
                    if let Some(items) =
                        field_lookup_items_for_type(struct_type, &right_ident, true)
                    {
                        completion.add_items(items);
                        self.completion = Some(completion);
                    }
                }
            }
            _ => {}
        }
    }

    fn append_assignment_clause_completions(
        &mut self,
        clause: Rc<AssignmentClauseContextAll<'static>>,
    ) {
        if self.completion.is_some() {
            return;
        }

        let range = completion_interval(clause.as_ref());
        if !range.contains(&self.at) {
            return;
        }

        let text = clause.get_text();
        let Some(env) = self.env() else {
            return;
        };

        if text.is_empty() {
            let insert_interval = interval(clause.as_ref());
            let mut completion = Completion::new(insert_interval);
            completion.add_items(env.autocomplete_suggestions(false));
            completion.add_items(self.registry.autocomplete_suggestions());
            self.completion = Some(completion);
            return;
        }

        // Mirror legacy behavior for broken assignment clauses by reparsing as an expression.
        let parsed = make_hamelin_parser_from_input(text).0.expressionEOF();
        if let Ok(expression_eof) = parsed {
            if let Some(expression) = expression_eof.expression() {
                let sub_at = self.at.saturating_sub(*range.start());
                if let Some(nested) =
                    self.build_nested_expression_completion(expression, sub_at, env)
                {
                    let mut completion = Completion::new_exclusive(
                        nested.at.start + *range.start()..nested.at.end + *range.start(),
                    );
                    if let Some(filter) = nested.filter {
                        completion.filter(filter);
                    }
                    completion.add_items(nested.items);
                    self.completion = Some(completion);
                }
            }
        }
    }

    fn append_expression_error_completions<'ctx, PRC: ParserRuleContext<'ctx>>(
        &mut self,
        ctx: &PRC,
    ) {
        if self.completion.is_some() {
            return;
        }
        let range = completion_interval(ctx);
        if !range.contains(&self.at) {
            return;
        }
        let Some(env) = self.env() else {
            return;
        };

        let insert_interval = interval(ctx);
        let mut completion = Completion::new(insert_interval);
        completion.filter(false);
        let mut suggestions = env.autocomplete_suggestions(false);
        if suggestions.is_empty() {
            if let Ok(datasets) = self.provider.reflect_datasets() {
                for dataset in datasets {
                    if let Ok(struct_type) = self.provider.reflect_columns(&dataset) {
                        suggestions.extend(struct_type.autocomplete_suggestions(false));
                    }
                }
            }
        }
        completion.add_items(suggestions);
        completion.add_items(self.registry.autocomplete_suggestions());
        self.completion = Some(completion);
    }

    fn build_nested_expression_completion(
        &self,
        expression: Rc<ExpressionContextAll<'static>>,
        at: usize,
        env: hamelin_lib::types::struct_type::Struct,
    ) -> Option<Completion> {
        // Keep this small and focused: nested broken assignment mostly needs column / field lookup completions.
        match expression.as_ref() {
            ExpressionContextAll::FieldReferenceAltContext(ctx) => {
                let parent_range = completion_interval(expression.as_ref());
                let field_ref = ctx.fieldReference()?;
                if !parent_range.contains(&at) {
                    return None;
                }
                let insert_interval = interval(field_ref.as_ref());
                let mut completion = Completion::new(insert_interval);
                if *parent_range.start() == at {
                    completion.filter(false);
                    completion.add_items(env.autocomplete_suggestions(false));
                    completion.add_items(self.registry.autocomplete_suggestions());
                    return Some(completion);
                }
                let mut parse_ctx = ParseContext::new();
                let maybe_curr_ident = field_ref.simpleIdentifier().and_then(|si| {
                    ParsedSimpleIdentifier::from_cst_with_context(si.as_ref(), &mut parse_ctx)
                        .valid()
                        .ok()
                });
                if let Some(curr_ident) = maybe_curr_ident {
                    completion.add_items(env.nested_autocomplete_suggestions(&curr_ident, false));
                    completion.add_items(self.registry.autocomplete_suggestions());
                    return Some(completion);
                }
                None
            }
            ExpressionContextAll::FieldLookupContext(ctx) => {
                let right = ctx.right.clone()?;
                let left = ctx.left.clone()?;
                let range = cuddled_completion_interval(right.as_ref());
                if !range.contains(&at) {
                    return None;
                }
                let expression_type = self.infer_expression_type(left)?;
                let right_ident = Identifier::parse(right.get_text()).ok();
                let insert_interval = if right_ident.is_some() {
                    let int = interval(right.as_ref());
                    int.start().saturating_sub(1)..=*int.end()
                } else {
                    at.saturating_sub(1)..=at.saturating_sub(1)
                };
                let mut completion = Completion::new(insert_interval);
                completion.filter(false);
                match expression_type {
                    Type::Struct(struct_type) => completion.add_items(field_lookup_items_for_type(
                        &struct_type,
                        &right_ident,
                        true,
                    )?),
                    Type::Array(array) => {
                        if let Type::Struct(struct_type) = array.element_type.as_ref() {
                            completion.add_items(field_lookup_items_for_type(
                                struct_type,
                                &right_ident,
                                true,
                            )?)
                        }
                    }
                    _ => return None,
                }
                Some(completion)
            }
            ExpressionContextAll::Error(ctx) => {
                let range = completion_interval(ctx);
                if !range.contains(&at) {
                    return None;
                }
                let mut completion = Completion::new(interval(ctx));
                completion.filter(false);
                completion.add_items(env.autocomplete_suggestions(false));
                completion.add_items(self.registry.autocomplete_suggestions());
                Some(completion)
            }
            _ => None,
        }
    }

    fn infer_expression_type(&self, expression: Rc<ExpressionContextAll<'static>>) -> Option<Type> {
        let bindings = self.typed_statement.environment_at_offset(self.at)?;
        let (ast_expression, _) = AstExpression::parse_with_errors(expression.get_text());
        let typed = type_check_expression(
            ast_expression,
            ExpressionTypeCheckOptions::builder()
                .registry(self.registry.clone())
                .bindings(bindings)
                .maybe_template_parameters(self.template_specs.clone())
                .build(),
        )
        .output;
        Some(typed.resolved_type.as_ref().clone())
    }
}

fn field_lookup_items_for_type(
    struct_type: &hamelin_lib::types::struct_type::Struct,
    right_ident: &Option<Identifier>,
    prepend_dot: bool,
) -> Option<Vec<CompletionItem>> {
    match right_ident {
        Some(Identifier::Simple(si)) => {
            Some(struct_type.nested_autocomplete_suggestions(si, prepend_dot))
        }
        _ => Some(struct_type.autocomplete_suggestions(prepend_dot)),
    }
}

#[cfg(test)]
mod template_param_apply_tests {
    use super::slice_from_char;
    use super::template_param_apply_string;

    #[test]
    fn apply_adds_close_when_missing() {
        assert_eq!(template_param_apply_string("tenant", false), "tenant}");
    }

    #[test]
    fn apply_skips_when_close_already_there() {
        assert_eq!(template_param_apply_string("tenant", true), "tenant");
    }

    #[test]
    fn slice_from_char_respects_unicode_scalar_indices() {
        let s = "a🦀b";
        assert_eq!(slice_from_char(s, 0), s);
        assert_eq!(slice_from_char(s, 1), "🦀b");
        assert_eq!(slice_from_char(s, 2), "b");
        assert_eq!(slice_from_char(s, 3), "");
    }
}
