hamelin_legacy 0.4.3

use crate::ast::expression::HamelinExpression;
use hamelin_lib::err::{TranslationError, TranslationErrors, UnexpectedType};
use hamelin_lib::sql::expression::apply::FunctionCallApply;
use hamelin_lib::sql::expression::literal::StringLiteral;
use hamelin_lib::sql::expression::Leaf;
use hamelin_lib::sql::expression::SQLExpression;
use hamelin_lib::types::array::Array;
use hamelin_lib::types::map::Map;
use hamelin_lib::types::{Type, UNKNOWN, VARIANT};
use hamelin_sql::utils::hamelin_array_index_to_sql_with_negatives;

use thiserror::Error;

pub struct HamelinIndexAccess {
    pub value: HamelinExpression,
    pub index: HamelinExpression,
}

impl HamelinIndexAccess {
    pub fn new(value: HamelinExpression, index: HamelinExpression) -> Self {
        Self { value, index }
    }

    pub fn infer_type(&self) -> Result<Type, TranslationErrors> {
        match self.value.translate()?.typ {
            Type::Array(Array { element_type }) => Ok(*element_type),
            Type::Map(Map {
                key_type: _key_type,
                value_type,
            }) => Ok(*value_type),
            Type::Variant => Ok(VARIANT),
            t => TranslationError::wrap(
                self.value.tree(),
                UnexpectedType::new(
                    t,
                    vec![
                        Array::new(UNKNOWN).into(),
                        Map::new(UNKNOWN, UNKNOWN).into(),
                        VARIANT,
                    ],
                ),
            )
            .single_result(),
        }
    }

    pub fn to_sql(&self) -> Result<SQLExpression, TranslationErrors> {
        let value = self.value.translate()?;
        let value_type = value.typ;
        let value_sql = value.sql;
        let index = self.index.translate()?;
        let index_type = index.typ;
        let index_sql = index.sql;
        let rendered_index_sql = index_sql.to_string();

        let expected_types = vec![
            Array::new(UNKNOWN).into(),
            Map::new(UNKNOWN, UNKNOWN).into(),
            VARIANT,
        ];
        match (value_type, index_type) {
            (Type::Array(_), Type::Int) => Ok(FunctionCallApply::with_two(
                "element_at",
                value_sql.clone(),
                hamelin_array_index_to_sql_with_negatives(index_sql, value_sql),
            )
            .into()),
            (Type::Map(map), index_type) if *(map.key_type) == index_type => {
                Ok(FunctionCallApply::with_two("element_at", value_sql, index_sql).into())
            }
            (Type::Variant, Type::Int) => {
                if let SQLExpression::FunctionCallApply(json_extract) = value_sql.clone() {
                    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.clone(),
                                json_extract.arguments.first().unwrap().clone(),
                                StringLiteral::new(&format!("{}[{}]", value, rendered_index_sql))
                                    .into(),
                            )
                            .into());
                        }
                    }
                }

                return Ok(FunctionCallApply::with_two(
                    "json_extract",
                    value_sql,
                    StringLiteral::new(&format!("$[{}]", rendered_index_sql)).into(),
                )
                .into());
            }
            (Type::Array(_), t) => {
                let inner =
                    BadIndexAccess::new(t, "Index must be an integer for array index access.");
                TranslationError::wrap(self.index.tree(), inner).single_result()
            }
            (
                Type::Map(Map {
                    key_type,
                    value_type: _value_type,
                }),
                _,
            ) => {
                let inner = BadIndexAccess::new(
                    *_value_type,
                    &format!(
                        "Index must match key type {} for map index access.",
                        key_type
                    ),
                );
                TranslationError::wrap(self.index.tree(), inner).single_result()
            }
            (Type::Variant, t) => {
                let inner =
                    BadIndexAccess::new(t, "Index must be integer for variant index access.");
                TranslationError::wrap(self.index.tree(), inner).single_result()
            }
            (vt, _) => {
                let inner = UnexpectedType::new(vt, expected_types);
                TranslationError::wrap(self.value.tree(), inner).single_result()
            }
        }
    }
}

#[derive(Error, Debug)]
#[error("Cannot perform index access with {typ}: {reason}")]
pub struct BadIndexAccess {
    pub typ: Type,
    pub reason: String,
}

impl BadIndexAccess {
    pub fn new(typ: Type, reason: &str) -> Self {
        Self {
            typ: typ,
            reason: reason.to_string(),
        }
    }
}