hamelin_legacy 0.3.8

use std::rc::Rc;

use crate::ast::expression::HamelinExpression;
use crate::env::Environment;
use hamelin_lib::antlr::hamelinparser::{ExpressionContextAll, SimpleIdentifierContextAll};
use hamelin_lib::antlr::{cuddled_completion_interval, interval};
use hamelin_lib::completion::Completion;
use hamelin_lib::err::{TranslationError, TranslationErrors, UnexpectedType};
use hamelin_lib::parse_identifier;
use hamelin_lib::sql::expression::apply::{FunctionCallApply, Lambda};
use hamelin_lib::sql::expression::identifier::{HamelinIdentifier, HamelinSimpleIdentifier};
use hamelin_lib::sql::expression::identifier::{Identifier, SimpleIdentifier};
use hamelin_lib::sql::expression::literal::{ColumnReference, IntegerLiteral, StringLiteral};
use hamelin_lib::sql::expression::{Dot, IndexLookup};
use hamelin_lib::sql::expression::{Leaf, SQLExpression};
use hamelin_lib::types::array::Array;
use hamelin_lib::types::struct_type::Struct;
use hamelin_lib::types::tuple::Tuple;
use hamelin_lib::types::{Type, UNKNOWN, VARIANT};

use super::ExpressionTranslationContext;

pub struct HamelinDerefApply {
    pub left: Rc<ExpressionContextAll<'static>>,
    pub right: Rc<SimpleIdentifierContextAll<'static>>,
    pub expression_translation_ctx: Rc<ExpressionTranslationContext>,
}

impl HamelinDerefApply {
    pub fn new(
        left: Rc<ExpressionContextAll<'static>>,
        right: Rc<SimpleIdentifierContextAll<'static>>,
        expression_translation_context: Rc<ExpressionTranslationContext>,
    ) -> Self {
        Self {
            left,
            right,
            expression_translation_ctx: expression_translation_context,
        }
    }

    pub fn infer_type(&self) -> Result<Type, TranslationErrors> {
        let expression =
            HamelinExpression::new(self.left.clone(), self.expression_translation_ctx.clone())
                .translate()?;
        let simple_ident = HamelinSimpleIdentifier::new(self.right.clone());
        let simple_ident_sql = simple_ident.to_sql()?;
        let ident = simple_ident_sql.clone().into();

        match expression.typ {
            Type::Struct(struct_type) => {
                let env = Environment::new(struct_type.clone());
                env.lookup(&ident).map_err(|e| {
                    TranslationError::msg(self.right.as_ref(), "field not found in struct")
                        .with_source(e)
                        .into()
                })
            }
            Type::Tuple(tuple) => {
                let index = simple_ident.get_tuple_index()?;
                let len = tuple.elements.len();
                if index >= len {
                    TranslationError::msg(
                        self.right.as_ref(),
                        &format!("cannot deref: tuple index {} out of bounds", index),
                    )
                    .with_context(self.left.as_ref(), &format!("has length {}", len))
                    .single_result()
                } else {
                    Ok((*tuple.elements)[index].clone())
                }
            }
            Type::Variant => Ok(VARIANT),
            Type::Array(array) => match *array.element_type {
                Type::Struct(inner_struct) => {
                    let env = Environment::new(inner_struct.clone());
                    let element_type = env.lookup(&ident).map(|ty| ty.clone()).map_err(|e| {
                        TranslationError::msg(
                            self.right.as_ref(),
                            "field not found nested in the array",
                        )
                        .with_source(e)
                    })?;

                    if let Type::Array(_) = element_type {
                        Ok(element_type)
                    } else {
                        Ok(Array::new(element_type).into())
                    }
                }
                Type::Variant => Ok(Array::new(VARIANT).into()),
                _ => TranslationError::msg(
                    self.right.as_ref(),
                    "can only deref an array if its elements are struct or variant",
                )
                .with_context(
                    self.left.as_ref(),
                    &format!("array element type is {}", array.element_type),
                )
                .single_result(),
            },
            _ => {
                let e = UnexpectedType::new(
                    expression.typ,
                    vec![
                        VARIANT,
                        Struct::default().into(),
                        Tuple::new(vec![UNKNOWN]).into(),
                        Array::new(Struct::default().into()).into(),
                    ],
                );
                Err(
                    TranslationError::msg(self.left.as_ref(), "type not expected on lhs of deref")
                        .with_source(e)
                        .into(),
                )
            }
        }
    }

    pub fn to_sql(&self) -> Result<SQLExpression, TranslationErrors> {
        let expression =
            HamelinExpression::new(self.left.clone(), self.expression_translation_ctx.clone())
                .translate()?;
        let simple_ident = HamelinSimpleIdentifier::new(self.right.clone());
        let simple_ident_sql = simple_ident.to_sql()?;
        let ident_sql: Identifier = simple_ident_sql.clone().into();

        match expression.typ {
            Type::Struct(_) => Ok(Dot::new(expression.sql, ident_sql.into()).into()),
            Type::Tuple(_) => {
                let index = simple_ident.get_tuple_index()?;
                Ok(IndexLookup::new(
                    expression.sql,
                    IntegerLiteral::from_int((index + 1) as i64).to_sql_expression(),
                )
                .to_sql_expression())
            }
            Type::Variant => {
                if let SQLExpression::FunctionCallApply(json_extract) = &expression.sql {
                    if json_extract.function_name == "json_extract"
                        && json_extract.arguments.len() == 2
                    {
                        let path = json_extract.arguments.last().unwrap().clone();
                        if let SQLExpression::Leaf(Leaf::StringLiteral(StringLiteral { value })) =
                            path
                        {
                            return Ok(FunctionCallApply::with_two(
                                &json_extract.function_name,
                                json_extract.arguments.first().unwrap().clone(),
                                StringLiteral::new(&format!(
                                    "{}[\"{}\"]",
                                    value, simple_ident_sql.name
                                ))
                                .into(),
                            )
                            .into());
                        }
                    }
                }

                return Ok(FunctionCallApply::with_two(
                    "json_extract",
                    expression.sql,
                    StringLiteral::new(&format!("$[\"{}\"]", simple_ident_sql.name)).into(),
                )
                .into());
            }
            Type::Array(array_type) => match *array_type.element_type {
                Type::Struct(struct_type) => {
                    let env = Environment::new(struct_type.clone());
                    let typ = env.lookup(&ident_sql.clone().into()).map_err(|e| {
                        TranslationError::msg(
                            self.right.as_ref(),
                            "field not found nested in the array",
                        )
                        .with_source(e)
                    })?;
                    let e = SimpleIdentifier::new("e");
                    let cr: SQLExpression = ColumnReference::new(e.clone().into()).into();
                    let transform: SQLExpression = FunctionCallApply::with_two(
                        "transform",
                        expression.sql,
                        Lambda::from_single_argument(e, cr.dot(ident_sql.clone().into())).into(),
                    )
                    .into();
                    return if let Type::Array(_) = typ {
                        Ok(FunctionCallApply::with_one("flatten", transform.into()).into())
                    } else {
                        Ok(transform.into())
                    };
                }
                Type::Variant => {
                    let e = SimpleIdentifier::new("e");
                    let cr: SQLExpression = ColumnReference::new(e.clone().into()).into();
                    let transform: SQLExpression = FunctionCallApply::with_two(
                        "transform",
                        expression.sql,
                        Lambda::from_single_argument(
                            e,
                            FunctionCallApply::with_two(
                                "json_extract",
                                cr,
                                StringLiteral::new(&format!("$[{}]", ident_sql)).into(),
                            )
                            .into(),
                        )
                        .into(),
                    )
                    .into();

                    Ok(transform.into())
                }
                _ => TranslationError::msg(
                    self.left.as_ref(),
                    "can only deref an array if its elements are struct or variant",
                )
                .with_context(
                    self.left.as_ref(),
                    &format!("array element type is {}", array_type.element_type),
                )
                .single_result(),
            },
            _ => {
                let e = UnexpectedType::new(
                    expression.typ,
                    vec![
                        VARIANT,
                        Struct::default().into(),
                        Tuple::new(vec![UNKNOWN]).into(),
                        Array::new(Struct::default().into()).into(),
                    ],
                );
                TranslationError::msg(
                    self.left.as_ref(),
                    "type not expected on lhs of deref".into(),
                )
                .with_source(e)
                .single_result()
            }
        }
    }

    pub fn append_completions(&self) {
        let maybe_guard = self.expression_translation_ctx.completions.try_borrow_mut();

        // Make the decision about whether to trigger completion from the original "right."
        let range = cuddled_completion_interval(self.right.as_ref());

        // Try to infer the type of the left-hand side.
        // Only complete if we can do that.
        let maybe_exp = HamelinExpression::new(
            self.left.clone(),
            self.expression_translation_ctx.clone().without_completion(),
        )
        .translate();

        match (self.expression_translation_ctx.at, maybe_guard, maybe_exp) {
            // - the cursor is cuddled up to the right hand side of an identifier in a deref
            // - the cursor is on top of a deref
            // - and no completion is set yet
            (Some(at), Ok(mut guard), Ok(expression)) if range.contains(&at) && guard.is_none() => {
                // ANTLR error recovery is making shitty identifiers out of the tree way too much.
                // So, re-parse the right to make sure we trust it. If we don't trust it, just ignore it.
                let right = parse_identifier(self.right.get_text());
                let right_ident = right
                    .clone()
                    .and_then(|id| HamelinIdentifier::new(id.clone()).to_sql());

                let insert_interval = if right_ident.is_ok() {
                    let int = interval(self.right.as_ref());
                    int.start() - 1..=int.end().clone()
                } else {
                    at - 1..=at - 1
                };

                let mut completion = Completion::new(insert_interval);
                completion.filter(false);
                match expression.typ {
                    Type::Struct(ref struct_type) => {
                        let env = Environment::new(struct_type.clone());
                        if let Ok(Identifier::Simple(si)) = right_ident {
                            completion.add_items(env.nested_autocomplete_suggestions(&si, true));
                        } else {
                            completion.add_items(env.autocomplete_suggestions(true));
                        }
                    }
                    Type::Array(ref array) => {
                        if let Type::Struct(ref struct_type) = *array.element_type {
                            let env = Environment::new(struct_type.clone());
                            if let Ok(Identifier::Simple(si)) = right_ident {
                                completion
                                    .add_items(env.nested_autocomplete_suggestions(&si, true));
                            } else {
                                completion.add_items(env.autocomplete_suggestions(true));
                            }
                        }
                    }
                    _ => {}
                }
                *guard = Some(completion);
            }
            _ => {}
        }
    }
}