hamelin_legacy 0.3.9

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

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 base: Option<Arc<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 {
            base: None,
            fields,
            pattern_variables: OrderMap::new(),
        }
    }

    /// Create a new environment extending an existing one
    pub fn extend(base: Arc<Environment>) -> Self {
        Self {
            base: Some(base),
            fields: Struct::default(),
            pattern_variables: OrderMap::new(),
        }
    }

    /// Flatten this environment by recursively collecting all fields from the base chain
    ///
    /// This is used when an environment needs to be stored as a field value (e.g., in FROM aliasing).
    /// The result is a flat Struct containing all fields from the environment and its base chain,
    /// with local fields overriding base fields of the same name.
    ///
    /// Example (pseudocode):
    /// ```text
    /// Environment {
    ///     base: Some(Environment { fields: {x: int, y: string} }),
    ///     fields: {y: int, z: boolean}  // y overrides base's y
    /// }
    /// .flatten() => Struct {x: int, y: int, z: boolean}
    /// ```
    pub fn flatten(&self) -> Struct {
        let mut all_fields = Struct::default();

        // Collect fields from base chain first (so local fields can override)
        if let Some(base) = &self.base {
            all_fields = base.flatten();
        }

        // Merge in local fields (will override base fields with same name)
        for (key, value) in self.fields.fields.iter() {
            all_fields = all_fields.with(key.clone().into(), value.clone());
        }

        all_fields
    }

    /// Nest all fields from this environment (including base chain) into a struct under a field name
    ///
    /// This creates a new flat Environment with:
    /// - All fields from this environment at the top level
    /// - PLUS a nested struct containing all those same fields under the given name
    ///
    /// This is used for aliased FROM/JOIN clauses where you want both top-level access
    /// and nested access through the alias.
    ///
    /// Example (pseudocode):
    /// ```text
    /// Environment { fields: {x: int, y: string} }
    /// .nest_into("foo")
    /// => Environment {
    ///     fields: {
    ///         x: int,
    ///         y: string,
    ///         foo: Struct {x: int, y: string}
    ///     }
    /// }
    /// ```
    pub fn nest_into(&self, field_name: impl Into<Identifier>) -> Self {
        let flattened = self.flatten();
        Environment::new(flattened.clone()).with_binding(field_name.into(), flattened.into())
    }

    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));
        }

        // First try local fields
        let local_type = maybe_current.and_then(|current| current.lookup(id.last()).cloned());

        // Also check base environment
        let base_type = self.base.as_ref().and_then(|base| base.lookup(id).ok());

        match (local_type, base_type) {
            // Both exist and are structs - merge them (local fields override base)
            (Some(Type::Struct(local_struct)), Some(Type::Struct(base_struct))) => {
                let mut merged = base_struct.clone();
                for (key, value) in local_struct.fields.iter() {
                    merged.fields.insert(key.clone(), value.clone());
                }
                Ok(Type::Struct(merged))
            }
            // Local exists (not a struct, or base doesn't have it) - use local
            (Some(typ), _) => Ok(typ),
            // Only base exists - use base
            (None, Some(typ)) => Ok(typ),
            // Not found anywhere
            (None, None) => Err(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());
        }

        // If not found locally, try the base environment
        if let Some(base) = &self.base {
            return base.lookup_pattern_variable(id);
        }

        // 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())
        );
    }

    #[test]
    fn test_environment_extension() {
        use std::sync::Arc;

        // Create a base environment with some fields
        let base = Environment::default()
            .with_binding("id".parse().unwrap(), INT)
            .with_binding("name".parse().unwrap(), STRING)
            .with_binding("active".parse().unwrap(), BOOLEAN);

        // Extend the base environment
        let extended = Environment::extend(Arc::new(base));

        // Should be able to look up fields from base
        assert_eq!(extended.lookup(&"id".parse().unwrap()), Ok(INT));
        assert_eq!(extended.lookup(&"name".parse().unwrap()), Ok(STRING));
        assert_eq!(extended.lookup(&"active".parse().unwrap()), Ok(BOOLEAN));

        // Should error on non-existent field
        assert!(extended.lookup(&"missing".parse().unwrap()).is_err());
    }

    #[test]
    fn test_environment_extension_with_local_fields() {
        use std::sync::Arc;

        // Create a base environment
        let base = Environment::default()
            .with_binding("id".parse().unwrap(), INT)
            .with_binding("name".parse().unwrap(), STRING);

        // Extend and add local fields
        let extended = Environment::extend(Arc::new(base))
            .with_binding("score".parse().unwrap(), INT)
            .with_binding("timestamp".parse().unwrap(), TIMESTAMP);

        // Should find local fields
        assert_eq!(extended.lookup(&"score".parse().unwrap()), Ok(INT));
        assert_eq!(
            extended.lookup(&"timestamp".parse().unwrap()),
            Ok(TIMESTAMP)
        );

        // Should find base fields
        assert_eq!(extended.lookup(&"id".parse().unwrap()), Ok(INT));
        assert_eq!(extended.lookup(&"name".parse().unwrap()), Ok(STRING));
    }

    #[test]
    fn test_environment_extension_shadowing() {
        use std::sync::Arc;

        // Create a base environment
        let base = Environment::default().with_binding("value".parse().unwrap(), INT);

        // Extend and shadow the field with a different type
        let extended =
            Environment::extend(Arc::new(base)).with_binding("value".parse().unwrap(), STRING);

        // Local field should shadow base field
        assert_eq!(extended.lookup(&"value".parse().unwrap()), Ok(STRING));
    }

    #[test]
    fn test_environment_extension_pattern_variables() {
        use std::sync::Arc;

        // Create a base environment with pattern variables
        let base = Environment::default()
            .with_pattern_variable("user".parse().unwrap(), STRING)
            .with_pattern_variable("host".parse().unwrap(), STRING);

        // Extend the environment
        let extended =
            Environment::extend(Arc::new(base)).with_pattern_variable("port".parse().unwrap(), INT);

        // Should find local pattern variable
        assert_eq!(
            extended
                .lookup_pattern_variable(&"port".parse().unwrap())
                .unwrap(),
            PatternVariable::new("port".to_string())
        );

        // Should find base pattern variables
        assert_eq!(
            extended
                .lookup_pattern_variable(&"user".parse().unwrap())
                .unwrap(),
            PatternVariable::new("user".to_string())
        );
        assert_eq!(
            extended
                .lookup_pattern_variable(&"host".parse().unwrap())
                .unwrap(),
            PatternVariable::new("host".to_string())
        );
    }

    #[test]
    fn test_environment_multi_level_extension() {
        use std::sync::Arc;

        // Create a three-level chain
        let base = Environment::default().with_binding("a".parse().unwrap(), INT);

        let middle = Environment::extend(Arc::new(base)).with_binding("b".parse().unwrap(), STRING);

        let top = Environment::extend(Arc::new(middle)).with_binding("c".parse().unwrap(), BOOLEAN);

        // Should find field at each level
        assert_eq!(top.lookup(&"c".parse().unwrap()), Ok(BOOLEAN)); // local
        assert_eq!(top.lookup(&"b".parse().unwrap()), Ok(STRING)); // middle
        assert_eq!(top.lookup(&"a".parse().unwrap()), Ok(INT)); // base
    }

    #[test]
    fn test_flatten_simple_environment() {
        // Test flattening a simple non-extended environment
        let env = Environment::default()
            .with_binding("x".parse().unwrap(), INT)
            .with_binding("y".parse().unwrap(), STRING);

        let flattened = env.flatten();

        // Should have both fields
        assert_eq!(flattened.lookup(&"x".parse().unwrap()).unwrap(), &INT);
        assert_eq!(flattened.lookup(&"y".parse().unwrap()).unwrap(), &STRING);
    }

    #[test]
    fn test_flatten_with_chain() {
        use std::sync::Arc;

        // Create a three-level chain
        let base = Environment::default().with_binding("x".parse().unwrap(), INT);

        let level1 = Environment::extend(Arc::new(base)).with_binding("y".parse().unwrap(), STRING);

        let level2 =
            Environment::extend(Arc::new(level1)).with_binding("z".parse().unwrap(), BOOLEAN);

        let flattened = level2.flatten();

        // Should have all fields from all levels
        assert_eq!(flattened.lookup(&"x".parse().unwrap()).unwrap(), &INT);
        assert_eq!(flattened.lookup(&"y".parse().unwrap()).unwrap(), &STRING);
        assert_eq!(flattened.lookup(&"z".parse().unwrap()).unwrap(), &BOOLEAN);
    }

    #[test]
    fn test_flatten_with_shadowing() {
        use std::sync::Arc;

        // Base has 'value' as INT
        let base = Environment::default()
            .with_binding("value".parse().unwrap(), INT)
            .with_binding("other".parse().unwrap(), BOOLEAN);

        // Extended shadows 'value' with STRING
        let extended =
            Environment::extend(Arc::new(base)).with_binding("value".parse().unwrap(), STRING);

        let flattened = extended.flatten();

        // Local field should override base field
        assert_eq!(
            flattened.lookup(&"value".parse().unwrap()).unwrap(),
            &STRING
        );
        // Base field should still be present
        assert_eq!(
            flattened.lookup(&"other".parse().unwrap()).unwrap(),
            &BOOLEAN
        );
    }

    #[test]
    fn test_flatten_with_struct_merging() {
        use hamelin_lib::types::struct_type::Struct;
        use std::sync::Arc;

        // Base has a struct with field 'a'
        let base_struct = Struct::default().with_str("a", INT);
        let base =
            Environment::default().with_binding("record".parse().unwrap(), base_struct.into());

        // Extended has the same struct field but adds field 'b'
        let extended_struct = Struct::default().with_str("b", STRING);
        let extended = Environment::extend(Arc::new(base))
            .with_binding("record".parse().unwrap(), extended_struct.into());

        let flattened = extended.flatten();

        // The local struct should override (not merge)
        let result_type = flattened.lookup(&"record".parse().unwrap()).unwrap();
        if let Type::Struct(s) = result_type {
            // Should only have 'b', not 'a' (override semantics)
            assert!(s.lookup(&"b".parse().unwrap()).is_some());
            assert!(s.lookup(&"a".parse().unwrap()).is_none());
        } else {
            panic!("Expected struct type");
        }
    }

    #[test]
    fn test_nest_into_simple() {
        use hamelin_lib::sql::expression::identifier::SimpleIdentifier;

        // Create environment with fields
        let env = Environment::default()
            .with_binding("x".parse().unwrap(), INT)
            .with_binding("y".parse().unwrap(), STRING);

        // Nest into "foo"
        let nested = env.nest_into(SimpleIdentifier::new("foo"));

        // Should have top-level access
        assert_eq!(nested.lookup(&"x".parse().unwrap()), Ok(INT));
        assert_eq!(nested.lookup(&"y".parse().unwrap()), Ok(STRING));

        // Should have nested access through "foo"
        assert_eq!(nested.lookup(&"foo.x".parse().unwrap()), Ok(INT));
        assert_eq!(nested.lookup(&"foo.y".parse().unwrap()), Ok(STRING));

        // The "foo" field itself should be a struct
        let foo_type = nested.lookup(&"foo".parse().unwrap()).unwrap();
        assert!(matches!(foo_type, Type::Struct(_)));
    }

    #[test]
    fn test_nest_into_with_chain() {
        use hamelin_lib::sql::expression::identifier::SimpleIdentifier;
        use std::sync::Arc;

        // Create a two-level chain
        let base = Environment::default().with_binding("a".parse().unwrap(), INT);

        let extended =
            Environment::extend(Arc::new(base)).with_binding("b".parse().unwrap(), STRING);

        // Nest the entire chain into "alias"
        let nested = extended.nest_into(SimpleIdentifier::new("alias"));

        // Should have top-level access to all fields
        assert_eq!(nested.lookup(&"a".parse().unwrap()), Ok(INT));
        assert_eq!(nested.lookup(&"b".parse().unwrap()), Ok(STRING));

        // Should have nested access through "alias"
        assert_eq!(nested.lookup(&"alias.a".parse().unwrap()), Ok(INT));
        assert_eq!(nested.lookup(&"alias.b".parse().unwrap()), Ok(STRING));
    }

    #[test]
    fn test_nest_into_preserves_all_fields() {
        use hamelin_lib::sql::expression::identifier::SimpleIdentifier;
        use std::sync::Arc;

        // Create a complex environment
        let base = Environment::default()
            .with_binding("id".parse().unwrap(), INT)
            .with_binding("name".parse().unwrap(), STRING);

        let extended = Environment::extend(Arc::new(base))
            .with_binding("timestamp".parse().unwrap(), TIMESTAMP)
            .with_binding("active".parse().unwrap(), BOOLEAN);

        let nested = extended.nest_into(SimpleIdentifier::new("t"));

        // All fields should be accessible both ways
        let fields = vec![
            ("id", INT),
            ("name", STRING),
            ("timestamp", TIMESTAMP),
            ("active", BOOLEAN),
        ];

        for (field_name, expected_type) in fields {
            // Top-level access
            assert_eq!(
                nested.lookup(&field_name.parse().unwrap()),
                Ok(expected_type.clone())
            );
            // Nested access
            let nested_path = format!("t.{}", field_name);
            assert_eq!(
                nested.lookup(&nested_path.parse().unwrap()),
                Ok(expected_type)
            );
        }
    }

    #[test]
    fn test_nest_into_no_base_chain() {
        use hamelin_lib::sql::expression::identifier::SimpleIdentifier;

        // The nested environment should NOT have a base chain
        // (it's flattened into a single-level environment)
        let base_env = Environment::default().with_binding("x".parse().unwrap(), INT);

        let extended = Environment::extend(std::sync::Arc::new(base_env))
            .with_binding("y".parse().unwrap(), STRING);

        let nested = extended.nest_into(SimpleIdentifier::new("foo"));

        // The nested environment should have no base (it's flattened)
        assert!(nested.base.is_none());

        // But should still be able to look up all fields
        assert_eq!(nested.lookup(&"x".parse().unwrap()), Ok(INT));
        assert_eq!(nested.lookup(&"y".parse().unwrap()), Ok(STRING));
        assert_eq!(nested.lookup(&"foo.x".parse().unwrap()), Ok(INT));
        assert_eq!(nested.lookup(&"foo.y".parse().unwrap()), Ok(STRING));
    }
}