hamelin_legacy 0.3.9

use std::{rc::Rc, sync::Arc};

use antlr_rust::parser_rule_context::ParserRuleContext;

use hamelin_lib::{
    antlr::{
        completion_interval,
        hamelinparser::{
            FromClauseContextAll, FromClauseContextAttrs, TableAliasContextAttrs,
            TableReferenceContextAttrs,
        },
        interval,
    },
    completion::{Completion, CompletionItem},
    err::{TranslationError, TranslationErrors},
    sql::{
        expression::identifier::{
            HamelinIdentifier, HamelinSimpleIdentifier, Identifier, SimpleIdentifier,
        },
        query::TableReference,
    },
};

use crate::{ast::QueryTranslationContext, env::Environment};

#[derive(Clone)]
pub struct HamelinFromClause {
    ctx: Rc<FromClauseContextAll<'static>>,
    cte: Arc<Environment>,
    context: Rc<QueryTranslationContext>,
}

impl HamelinFromClause {
    pub fn new(
        ctx: Rc<FromClauseContextAll<'static>>,
        cte: Arc<Environment>,
        context: Rc<QueryTranslationContext>,
    ) -> Self {
        HamelinFromClause { ctx, cte, context }
    }
    pub fn table_identifier(&self) -> Result<Identifier, TranslationErrors> {
        let tr = match self.ctx.tableAlias() {
            Some(tac) => TranslationErrors::expect(tac.as_ref(), tac.tableReference())?,
            None => TranslationErrors::expect(self.ctx.as_ref(), self.ctx.tableReference())?,
        };

        HamelinIdentifier::new(TranslationErrors::expect(tr.as_ref(), tr.identifier())?).to_sql()
    }

    pub fn table_alias(&self) -> Option<HamelinSimpleIdentifier> {
        self.ctx
            .tableAlias()
            .and_then(|tac| tac.simpleIdentifier())
            .map(|si| HamelinSimpleIdentifier::new(si))
    }

    pub fn reflect(
        &self,
    ) -> Result<(Environment, TableReference, Option<SimpleIdentifier>), TranslationErrors> {
        let table_identifier = self.table_identifier()?;
        let table_alias = self.table_alias().map(|hsi| hsi.to_sql()).transpose()?;

        let table_reference = TableReference::new(table_identifier.clone());
        let env = match self.cte.lookup(&table_identifier) {
            Ok(typ) => {
                let strct = typ
                    .try_unwrap_struct()
                    .map_err(|e| TranslationError::wrap_box(self.ctx.as_ref(), e.into()))?;
                Environment::new(strct)
            }
            Err(_) => self
                .context
                .provider
                .reflect_columns(table_reference.clone())
                .map(Environment::new)
                .map_err(|e| TranslationError::wrap_box(self.ctx.as_ref(), e.into()))?,
        };

        Ok((env, table_reference, table_alias))
    }

    pub fn append_completions(&self) {
        let table_reference = self
            .ctx
            .as_ref()
            .tableReference()
            .or_else(|| self.ctx.tableAlias().and_then(|ta| ta.tableReference()));

        match (table_reference, self.context.at) {
            // Someone has started typing a table reference
            (Some(cur), Some(at)) if completion_interval(cur.as_ref()).contains(&at) => {
                let mut guard = self.context.completions.borrow_mut();
                self.do_completion_append(cur.as_ref(), &mut guard);
            }

            // We are completing inside a broken from clause (probably because nothing typed yet)
            (None, Some(at)) if completion_interval(self.ctx.as_ref()).contains(&at) => {
                let mut guard = self.context.completions.borrow_mut();
                self.do_completion_append(self.ctx.as_ref(), &mut guard);
            }
            _ => {}
        }
    }

    fn do_completion_append<'a, PRC: ParserRuleContext<'a>>(
        &self,
        cur: &PRC,
        guard: &mut Option<Completion>,
    ) {
        let insert_interval = interval(cur);
        let mut completion = Completion::new(insert_interval);
        for k in self.cte.keys().map(|i| i.clone().into()).chain(
            self.context
                .provider
                .reflect_datasets()
                .unwrap_or_default()
                .into_iter(),
        ) {
            if k.to_hamelin().contains(cur.get_text().as_str()) {
                completion.add_item(CompletionItem::new(k.to_hamelin()));
            }
        }
        *guard = Some(completion);
    }
}