hamelin_legacy 0.4.4

use std::error::Error;
use std::fmt::Display;

use hamelin_lib::completion::CompletionItem;
use ordermap::OrderMap;
use thiserror::Error;

use hamelin_lib::catalog::Column;
use hamelin_lib::err::{NonMergeableTypes, TranslationError, TranslationErrors};
use hamelin_lib::sql::expression::identifier::HamelinIdentifier;
use hamelin_lib::sql::expression::identifier::{Identifier, SimpleIdentifier};
use hamelin_lib::sql::expression::literal::ColumnReference;
use hamelin_lib::sql::query::projection::ColumnProjection;
use hamelin_lib::sql::query::PatternVariable;
use hamelin_lib::types::struct_type::{DropError, Struct};
use hamelin_lib::types::Type;

#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct Environment {
    pub fields: Struct,
    pub pattern_variables: OrderMap<SimpleIdentifier, PatternVariable>,
}

#[derive(Debug, PartialEq)]
pub struct UnboundColumnReference {
    pub environment: Environment,
    pub column_reference: Identifier,
}

impl UnboundColumnReference {
    pub fn new(environment: &Environment, column_reference: &Identifier) -> Self {
        Self {
            environment: environment.clone(),
            column_reference: column_reference.clone(),
        }
    }
}

impl Display for UnboundColumnReference {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        writeln!(
            f,
            "unbound column reference: {}",
            self.column_reference.to_hamelin(),
        )?;
        // look for fields in the environment that are prefixed with the column reference
        let first = self.column_reference.first();
        let candidates: Vec<&SimpleIdentifier> = self
            .environment
            .fields
            .fields
            .keys()
            .filter(|ident| ident.name.contains(first.name.as_str()))
            .collect();

        if !candidates.is_empty() {
            writeln!(f)?;
            writeln!(f, "the following entries in the environment are close:")?;
            for candidate in candidates {
                write!(f, "- {}", candidate.to_hamelin())?;
                if candidate.to_hamelin().starts_with("`") {
                    write!(f, " (you must actually wrap with ``)")?;
                }
                writeln!(f)?;
            }
        } else {
            writeln!(f, "in {}", self.environment)?;
        }

        Ok(())
    }
}

impl Error for UnboundColumnReference {}

#[derive(Error, Debug)]
pub struct NotAPatternVariable {
    pub env: Environment,
    pub ident: SimpleIdentifier,
}

impl NotAPatternVariable {
    pub fn new(env: Environment, ident: SimpleIdentifier) -> Self {
        Self { env, ident }
    }
}

impl Display for NotAPatternVariable {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        writeln!(f, "{} is not a pattern variable.", self.ident.to_hamelin())?;
        writeln!(f, "Pattern variables in the environment:")?;
        for key in self.env.pattern_variables.keys() {
            let ty = self
                .env
                .lookup(&key.clone().into())
                .unwrap_or_else(|_| Type::Unknown);
            writeln!(f, "- {}: {}", key.to_hamelin(), ty)?;
        }

        Ok(())
    }
}

impl Environment {
    pub fn new(fields: Struct) -> Self {
        Self {
            fields,
            pattern_variables: OrderMap::new(),
        }
    }

    pub fn lookup(&self, id: &Identifier) -> Result<Type, UnboundColumnReference> {
        let prefix = id.prefix();
        let mut maybe_current = Some(&self.fields);
        for name in prefix {
            maybe_current = maybe_current.and_then(|f| f.lookup_nested(&name));
        }

        maybe_current
            .and_then(|current| current.lookup(id.last()).cloned())
            .ok_or_else(|| UnboundColumnReference::new(self, id))
    }

    pub fn lookup_pattern_variable(
        &self,
        id: &SimpleIdentifier,
    ) -> Result<PatternVariable, NotAPatternVariable> {
        // First try local pattern variables
        if let Some(pv) = self.pattern_variables.get(id) {
            return Ok(pv.clone());
        }

        // Not found anywhere
        Err(NotAPatternVariable::new(self.clone(), id.clone()))
    }

    pub fn with_binding(mut self, id: Identifier, ty: Type) -> Self {
        self.fields = self.fields.with(id, ty);
        self
    }

    pub fn set_binding(&mut self, id: Identifier, ty: Type) {
        self.fields = self.fields.with(id, ty);
    }

    /// Lookup a field only in this environment's local fields (not inherited from base)
    /// This is useful for determining if a command defines/redefines a specific field.
    /// Returns the type if found locally, None otherwise.
    pub fn lookup_local(&self, id: &Identifier) -> Option<Type> {
        let prefix = id.prefix();
        let mut maybe_current = Some(&self.fields);
        for name in prefix {
            maybe_current = maybe_current.and_then(|f| f.lookup_nested(&name));
        }

        // Return the field type if it exists in local fields (not base)
        maybe_current.and_then(|current| current.lookup(id.last()).cloned())
    }

    pub fn with_pattern_variable(mut self, id: SimpleIdentifier, ty: Type) -> Self {
        // Pattern variables in a SQL match expression must be "clean" -- meaning, you can't escape
        // them or use weird characters like you can for field identifiers. I don't want to put
        // that same requirement on our users, because we treat any dataset reference expression
        // as a workable pattern variable. So, here, we need to create a mapping between identifiers
        // (which can be in column reference or table reference position) and the pattern variable
        // it maps to.
        //
        // If the identifier is "clean" it's easy to map. But if it's not, we have to construct
        // something weird (and make sure it doesn't clash with what's already in the environment).
        let pv: PatternVariable = id.clone().try_into().unwrap_or_else(|_| {
            let mut attemptnum = 1;
            let mut pvattempt = format!("pv{}", attemptnum);
            while self.lookup(&pvattempt.parse().unwrap()).is_ok()
                || self
                    .pattern_variables
                    .values()
                    .find(|pv| pv.name == pvattempt)
                    .is_some()
            {
                attemptnum += 1;
                pvattempt = format!("pv{}", attemptnum);
            }

            PatternVariable::new(pvattempt)
        });

        self.pattern_variables.insert(id.clone(), pv);
        self.with_binding(id.into(), ty)
    }

    pub fn drop(self, columns_to_drop: Vec<ColumnReference>) -> Result<Self, DropError> {
        let mut fields = self.fields;
        for column in columns_to_drop.into_iter() {
            fields = fields.drop(column.identifier)?
        }
        Ok(Self::new(fields))
    }

    pub fn replace(&self, column: ColumnReference, typ: Type) -> Self {
        Self::new(self.fields.replace(column.identifier, typ))
    }

    pub fn merge(&self, other: &Environment) -> Result<Self, NonMergeableTypes> {
        self.fields
            .merge(&other.fields)
            .map(|fields| Self::new(fields))
    }

    pub fn merge_prepend_overwrite(self, other: &Environment) -> Result<Self, NonMergeableTypes> {
        self.fields
            .merge_prepend_overwrite(&other.fields)
            .map(|fields| Self::new(fields))
    }

    pub fn prepend_overwrite(&self, other: &Environment) -> Self {
        Self::new(self.fields.prepend_overwrite(&other.fields))
    }

    pub fn remove_substructure(&self) -> Self {
        Self::new(self.fields.remove_substructure())
    }

    pub fn check_column_references(
        &self,
        column_references: &[HamelinIdentifier],
    ) -> Result<(), TranslationErrors> {
        TranslationErrors::from_vec(
            column_references
                .iter()
                .map(|column_reference| {
                    column_reference.to_sql().and_then(|cr| {
                        if let Err(e) = self.lookup(&cr) {
                            TranslationError::wrap(column_reference.ctx.as_ref(), e).single_result()
                        } else {
                            Ok(())
                        }
                    })
                })
                .collect(),
        )?;

        Ok(())
    }

    pub fn get_column_projections(&self) -> Vec<ColumnProjection> {
        self.fields
            .fields
            .keys()
            .map(|k| ColumnProjection::new(k.clone().into()))
            .collect()
    }

    pub fn keys(&self) -> impl Iterator<Item = &SimpleIdentifier> {
        self.fields.keys()
    }

    pub fn autocomplete_suggestions(&self, prepend_dot: bool) -> Vec<CompletionItem> {
        self.fields.autocomplete_suggestions(prepend_dot)
    }

    pub fn nested_autocomplete_suggestions(
        &self,
        ident: &SimpleIdentifier,
        prepend_dot: bool,
    ) -> Vec<CompletionItem> {
        self.fields
            .nested_autocomplete_suggestions(ident, prepend_dot)
    }

    pub fn into_external_columns(self) -> Vec<Column> {
        self.fields
            .fields
            .into_iter()
            .map(|(ident, typ)| Column {
                name: ident.to_hamelin(),
                typ: typ.into(),
            })
            .collect()
    }
}

impl Display for Environment {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "environment {{\n")?;
        for (i, (k, v)) in self.fields.fields.iter().enumerate() {
            if i >= 10 {
                writeln!(f, "    ...")?;
                break;
            }
            if let Some(pv) = self.pattern_variables.get(k) {
                write!(f, "    {} (pattern variable {}): ", k.to_hamelin(), pv)?;
            } else {
                write!(f, "    {}: ", k.to_hamelin())?;
            }

            v.fmt_indented(f, 4, 5)?;
            writeln!(f)?;
        }
        write!(f, "}}")
    }
}

#[cfg(test)]
mod test {
    use hamelin_lib::types::{BOOLEAN, INT, STRING, TIMESTAMP};

    use super::*;

    #[test]
    pub fn test_empty() {
        let env = Environment::default();

        assert_eq!(
            env.lookup(&"t1".parse().unwrap()),
            Err(super::UnboundColumnReference::new(
                &env,
                &"t1".parse().unwrap()
            ))
        );
    }

    #[test]
    pub fn test_single() {
        let env = Environment::default().with_binding("t1".parse().unwrap(), INT);

        assert_eq!(env.lookup(&"t1".parse().unwrap()), Ok(INT));
        assert_eq!(
            env.lookup(&"t2".parse().unwrap()),
            Err(super::UnboundColumnReference::new(
                &env,
                &"t2".parse().unwrap()
            ))
        );
    }

    #[test]
    pub fn test_complex() {
        let env = Environment::default().with_binding("t1.t2".parse().unwrap(), INT);

        assert_eq!(
            env.lookup(&"t1".parse().unwrap()),
            Ok(Struct::default().with("t2".parse().unwrap(), INT).into())
        );
        assert_eq!(
            env.lookup(&"t2".parse().unwrap()),
            Err(UnboundColumnReference::new(&env, &"t2".parse().unwrap()))
        );
        assert_eq!(env.lookup(&"t1.t2".parse().unwrap()), Ok(INT));
        assert_eq!(
            env.lookup(&"t1.t3".parse().unwrap()),
            Err(UnboundColumnReference::new(&env, &"t1.t3".parse().unwrap()))
        );
    }

    #[test]
    pub fn drop_complex() {
        let env = Environment::default()
            .with_binding("t1.nt11".parse().unwrap(), INT)
            .with_binding("t1.nt12".parse().unwrap(), BOOLEAN)
            .with_binding("t1.nt21".parse().unwrap(), STRING)
            .with_binding("t1.nt22".parse().unwrap(), TIMESTAMP);

        let dropped_env = env
            .drop(vec!["t1.nt11".parse().unwrap(), "t1.nt21".parse().unwrap()])
            .unwrap();

        assert_eq!(
            dropped_env.lookup(&"t1.nt11".parse().unwrap()),
            Err(UnboundColumnReference::new(
                &dropped_env,
                &"t1.nt11".parse().unwrap()
            ))
        );
        assert_eq!(dropped_env.lookup(&"t1.nt12".parse().unwrap()), Ok(BOOLEAN));
        assert_eq!(
            dropped_env.lookup(&"t1.nt21".parse().unwrap()),
            Err(UnboundColumnReference::new(
                &dropped_env,
                &"t1.nt21".parse().unwrap()
            ))
        );
        assert_eq!(
            dropped_env.lookup(&"t1.nt22".parse().unwrap()),
            Ok(TIMESTAMP)
        );
    }

    #[test]
    fn test_pattern_variable() {
        let env = Environment::default()
            .with_binding("pv1".parse().unwrap(), INT)
            .with_pattern_variable("kyle".parse().unwrap(), INT)
            .with_pattern_variable("rachel".parse().unwrap(), STRING)
            .with_pattern_variable("`let`".parse().unwrap(), BOOLEAN)
            .with_pattern_variable("`select`".parse().unwrap(), TIMESTAMP);

        // Test ordinary words as simple identifiers
        assert_eq!(
            env.lookup_pattern_variable(&"kyle".parse().unwrap())
                .unwrap(),
            PatternVariable::new("kyle".to_string())
        );
        assert_eq!(
            env.lookup_pattern_variable(&"rachel".parse().unwrap())
                .unwrap(),
            PatternVariable::new("rachel".to_string())
        );

        // Test reserved words that should map to special strings
        assert_eq!(
            env.lookup_pattern_variable(&"`let`".parse().unwrap())
                .unwrap(),
            // pv1 is already taken by an actual variable.
            PatternVariable::new("pv2".to_string())
        );
        assert_eq!(
            env.lookup_pattern_variable(&"`select`".parse().unwrap())
                .unwrap(),
            PatternVariable::new("pv3".to_string())
        );
    }
}