use std::sync::Arc;

use thiserror::Error;

use crate::{
    tree::ast::identifier::{Identifier, SimpleIdentifier},
    types::{struct_type::Struct, Type},
};

/// A type environment representing the schema available at a point in a pipeline.
///
/// This is a newtype around `Struct` — each command produces a fully materialized
/// environment with no lazy merging or parent chains. Field ordering is explicit
/// and determined by each command's semantics.
#[derive(Debug, Clone, Default)]
pub struct TypeEnvironment {
    schema: Struct,
}

/// Errors that can occur when working with TypeEnvironment
#[derive(Debug, Error)]
pub enum TypeEnvironmentError {
    #[error("Identifier not found: '{ident}'")]
    IdentifierNotFound { ident: String },
}

/// Error produced when a field reference cannot be resolved in the environment.
///
/// Display finds close matches via case-insensitive substring or edit distance,
/// otherwise shows the full environment schema.
#[derive(Debug)]
pub struct UnboundFieldReference {
    pub field: String,
    pub schema: Struct,
}

impl UnboundFieldReference {
    /// Recursively collect matching field paths from a struct.
    ///
    /// Scores: 2.0 for exact case-insensitive match, 1.0 for substring,
    /// 0..1 for fuzzy (Jaro-Winkler + Damerau-Levenshtein).
    fn collect_matches(
        schema: &Struct,
        needle: &str,
        prefix: String,
        results: &mut Vec<(String, f64)>,
    ) {
        for (key, typ) in schema.iter() {
            let path = if prefix.is_empty() {
                key.name().to_string()
            } else {
                format!("{}.{}", prefix, key.name())
            };
            let key_lower = key.name().to_lowercase();

            if key_lower == needle {
                results.push((path.clone(), 2.0));
            } else if key_lower.contains(needle)
                || (key_lower.len() > 4 && needle.contains(&key_lower))
            {
                results.push((path.clone(), 1.0));
            } else {
                let jw = strsim::jaro_winkler(needle, &key_lower);
                let ndl = strsim::normalized_damerau_levenshtein(needle, &key_lower);
                let score = (jw * 0.7) + (ndl * 0.3);
                if score > 0.7 {
                    results.push((path.clone(), score));
                }
            }

            if let Type::Struct(nested) = typ {
                Self::collect_matches(nested, needle, path, results);
            }
        }
    }
}

impl std::fmt::Display for UnboundFieldReference {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let needle = self.field.to_lowercase();
        let mut results = Vec::new();
        Self::collect_matches(&self.schema, &needle, String::new(), &mut results);
        results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));

        if !results.is_empty() {
            writeln!(f, "the following fields are similar:")?;
            for (path, _) in results {
                writeln!(f, "- {}", path)?;
            }
        } else {
            writeln!(f, "environment at this location:")?;
            self.schema.fmt_indented(f, 0, 5)?;
            writeln!(f)?;
        }
        Ok(())
    }
}

impl std::error::Error for UnboundFieldReference {}

impl TypeEnvironment {
    /// Create a new environment from a Struct schema
    pub fn new(schema: Struct) -> Self {
        Self { schema }
    }

    /// Build a SET/WINDOW/EXPLODE output environment from assignments and a parent env.
    ///
    /// Assignment keys appear first (in declaration order), with compound
    /// identifiers expanding into the parent's existing struct (preserving
    /// sibling fields). Simple identifiers shadow/replace entirely.
    /// Untouched parent fields are appended at the end.
    pub fn extend(
        assignments: &[(Identifier, Arc<Type>)],
        parent: &TypeEnvironment,
    ) -> TypeEnvironment {
        let capacity = assignments.len() + parent.schema.len();
        let mut result = TypeEnvironment::new(Struct::with_capacity(capacity));
        let mut touched_keys: Vec<SimpleIdentifier> = Vec::new();

        for (id, typ) in assignments {
            let top_key = id.first();
            if !touched_keys.contains(&top_key) {
                touched_keys.push(top_key.clone());
                // Seed with parent's version so compound ids preserve siblings
                if let Some(parent_type) = parent.schema.lookup_arc(&top_key) {
                    result
                        .schema
                        .insert_arc(top_key.clone().into(), parent_type.clone());
                }
            }
            // Apply the assignment (overwrites for simple, expands for compound)
            result.bind_arc(id.clone(), typ.clone());
        }

        // Append untouched parent fields (Arc clone, not deep Type clone)
        for (key, typ) in parent.schema.iter_arc() {
            if !touched_keys.contains(key) {
                result.schema.insert_arc(key.clone().into(), typ.clone());
            }
        }

        result.schema.shrink_to_fit();
        result
    }

    /// Look up the type of an identifier.
    ///
    /// For simple identifiers, looks up the top-level field.
    /// For compound identifiers (x.y.z), navigates nested structs.
    /// Also reconstructs struct types for intermediate lookups
    /// (e.g., looking up "x" when "x.y" and "x.z" exist returns Struct{y, z}).
    pub fn lookup(&self, id: &Identifier) -> Option<Arc<Type>> {
        let segments = id.segments();

        let mut current_struct = &self.schema;
        // Navigate through all segments except the last
        for segment in &segments[..segments.len() - 1] {
            match current_struct.lookup(segment) {
                Some(Type::Struct(nested)) => {
                    current_struct = nested;
                }
                _ => return None,
            }
        }

        // Look up the final segment
        let last = &segments[segments.len() - 1];
        current_struct.lookup_arc(last).cloned()
    }

    /// Append an identifier binding at the end.
    ///
    /// For simple identifiers, appends at the top level.
    /// For compound identifiers (x.a), appends within the parent struct.
    /// Used by SELECT, AGG, and struct literal construction where fields
    /// appear in declaration order.
    pub fn bind(&mut self, id: Identifier, typ: Type) {
        self.schema.insert(id, typ);
    }

    /// Append a pre-wrapped `Arc<Type>` binding.
    ///
    /// For simple identifiers, avoids re-wrapping in `Arc::new()`.
    /// For compound identifiers, falls back to unwrapping the Arc.
    pub fn bind_arc(&mut self, id: Identifier, typ: Arc<Type>) {
        self.schema.insert_arc(id, typ);
    }

    /// Append a simple string-keyed type binding.
    pub fn bind_str(&mut self, name: impl Into<String>, typ: Type) {
        self.bind(SimpleIdentifier::new(name).into(), typ);
    }

    /// Append a type binding (consuming self, builder-style)
    pub fn with(self, id: Identifier, typ: Type) -> Self {
        let mut env = self;
        env.bind(id, typ);
        env
    }

    /// Append a simple string-keyed type binding (consuming self, builder-style)
    pub fn with_str(self, name: impl Into<String>, typ: Type) -> Self {
        let mut env = self;
        env.bind_str(name, typ);
        env
    }

    /// Drop a field from the environment.
    ///
    /// For simple identifiers, removes the top-level field.
    /// For compound identifiers (x.a), removes the nested field from its parent struct.
    pub fn drop(&mut self, id: &Identifier) -> Result<(), TypeEnvironmentError> {
        self.schema =
            self.schema
                .drop(id.clone())
                .map_err(|_| TypeEnvironmentError::IdentifierNotFound {
                    ident: id.to_string(),
                })?;
        Ok(())
    }

    /// Drop a field, consuming self
    pub fn without(self, id: Identifier) -> Result<Self, TypeEnvironmentError> {
        let new_schema =
            self.schema
                .drop(id.clone())
                .map_err(|_| TypeEnvironmentError::IdentifierNotFound {
                    ident: id.to_string(),
                })?;
        Ok(Self::new(new_schema))
    }

    /// Replace the type of an existing field while preserving field ordering.
    ///
    /// Used by EXPLODE to change a field from ARRAY<T> to T in-place.
    pub fn replace(
        &self,
        id: &Identifier,
        new_type: Type,
    ) -> Result<TypeEnvironment, TypeEnvironmentError> {
        // Check if the field exists first
        if self.lookup(id).is_none() {
            return Err(TypeEnvironmentError::IdentifierNotFound {
                ident: id.to_string(),
            });
        }
        Ok(Self::new(self.schema.replace(id.clone(), new_type)))
    }

    /// Get a reference to the underlying Struct schema
    pub fn as_struct(&self) -> &Struct {
        &self.schema
    }
}

impl From<Struct> for TypeEnvironment {
    fn from(s: Struct) -> Self {
        Self::new(s)
    }
}

impl std::fmt::Display for TypeEnvironment {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "TypeEnvironment {{")?;
        for (i, (k, v)) in self.schema.iter().enumerate() {
            if i > 0 {
                write!(f, ",")?;
            }
            write!(f, " {}: {}", k, v)?;
        }
        write!(f, " }}")
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::tree::builder::identifier::ident;
    use crate::types::array::Array;
    use crate::types::{BOOLEAN, INT, STRING};

    #[test]
    fn test_simple_bind_and_lookup() {
        let mut env = TypeEnvironment::default();
        env.bind(ident("x").into(), INT);
        env.bind(ident("y").into(), STRING);

        assert_eq!(env.lookup(&ident("x").into()).as_deref(), Some(&INT));
        assert_eq!(env.lookup(&ident("y").into()).as_deref(), Some(&STRING));
        assert_eq!(env.lookup(&ident("z").into()), None);
    }

    #[test]
    fn test_compound_bind_and_lookup() {
        let mut env = TypeEnvironment::default();
        env.bind(ident("x").dot("y").into(), INT);
        env.bind(ident("x").dot("z").into(), STRING);

        assert_eq!(
            env.lookup(&ident("x").dot("y").into()).as_deref(),
            Some(&INT)
        );
        assert_eq!(
            env.lookup(&ident("x").dot("z").into()).as_deref(),
            Some(&STRING)
        );
        // Intermediate lookup returns struct
        let x_type = env.lookup(&ident("x").into());
        assert!(matches!(x_type.as_deref(), Some(Type::Struct(_))));
    }

    #[test]
    fn test_from_struct() {
        let s = Struct::default().with_str("a", INT).with_str("b", STRING);
        let env = TypeEnvironment::from(s.clone());

        assert_eq!(env.lookup(&ident("a").into()).as_deref(), Some(&INT));
        assert_eq!(env.lookup(&ident("b").into()).as_deref(), Some(&STRING));
        assert_eq!(*env.as_struct(), s);
    }

    #[test]
    fn test_drop_simple() {
        let env = TypeEnvironment::from(
            Struct::default()
                .with_str("x", INT)
                .with_str("y", STRING)
                .with_str("z", BOOLEAN),
        );
        let env = env.without(ident("y").into()).unwrap();

        assert_eq!(env.lookup(&ident("x").into()).as_deref(), Some(&INT));
        assert_eq!(env.lookup(&ident("y").into()), None);
        assert_eq!(env.lookup(&ident("z").into()).as_deref(), Some(&BOOLEAN));
    }

    #[test]
    fn test_drop_nested() {
        let env = TypeEnvironment::from(
            Struct::default().with_str(
                "x",
                Struct::default()
                    .with_str("a", INT)
                    .with_str("b", STRING)
                    .into(),
            ),
        );
        let env = env.without(ident("x").dot("a").into()).unwrap();

        assert_eq!(env.lookup(&ident("x").dot("a").into()), None);
        assert_eq!(
            env.lookup(&ident("x").dot("b").into()).as_deref(),
            Some(&STRING)
        );
    }

    #[test]
    fn test_replace() {
        let env = TypeEnvironment::from(Struct::default().with_str("x", INT).with_str("y", STRING));
        let env = env.replace(&ident("x").into(), BOOLEAN).unwrap();

        assert_eq!(env.lookup(&ident("x").into()).as_deref(), Some(&BOOLEAN));
        assert_eq!(env.lookup(&ident("y").into()).as_deref(), Some(&STRING));
    }

    #[test]
    fn test_replace_compound() {
        let env = TypeEnvironment::from(
            Struct::default().with_str(
                "data",
                Struct::default()
                    .with_str("arr", Array::new(INT).into())
                    .with_str("other", STRING)
                    .into(),
            ),
        );
        let env = env.replace(&ident("data").dot("arr").into(), INT).unwrap();

        assert_eq!(
            env.lookup(&ident("data").dot("arr").into()).as_deref(),
            Some(&INT)
        );
        assert_eq!(
            env.lookup(&ident("data").dot("other").into()).as_deref(),
            Some(&STRING)
        );
    }

    #[test]
    fn test_replace_not_found() {
        let env = TypeEnvironment::from(Struct::default().with_str("x", INT));
        let result = env.replace(&ident("nonexistent").into(), STRING);
        assert!(result.is_err());
    }

    #[test]
    fn test_from_set_struct_literal_shadow() {
        // SET x = {a: 1, b: 2} when x.c already exists should replace x entirely
        let parent = TypeEnvironment::from(
            Struct::default().with_str("x", Struct::default().with_str("c", BOOLEAN).into()),
        );
        let assignments: Vec<(Identifier, Arc<Type>)> = vec![(
            ident("x").into(),
            Arc::new(
                Struct::default()
                    .with_str("a", INT)
                    .with_str("b", STRING)
                    .into(),
            ),
        )];
        let result = TypeEnvironment::extend(&assignments, &parent);

        // x.c is gone, x.a and x.b exist
        assert_eq!(result.lookup(&ident("x").dot("c").into()), None);
        assert_eq!(
            result.lookup(&ident("x").dot("a").into()).as_deref(),
            Some(&INT)
        );
        assert_eq!(
            result.lookup(&ident("x").dot("b").into()).as_deref(),
            Some(&STRING)
        );
    }

    #[test]
    fn test_from_set_compound_preserves_siblings() {
        // SET x.e = 1 when input has x: {b: INT, c: INT} should keep b and c
        let parent = TypeEnvironment::from(
            Struct::default()
                .with_str("a", INT)
                .with_str(
                    "x",
                    Struct::default()
                        .with_str("b", INT)
                        .with_str("c", INT)
                        .into(),
                )
                .with_str("d", STRING),
        );
        let assignments = vec![(ident("x").dot("e").into(), Arc::new(INT))];
        let result = TypeEnvironment::extend(&assignments, &parent);

        // x is first (from assignments), then a, d from parent appended
        let keys: Vec<_> = result
            .as_struct()
            .iter()
            .map(|(k, _)| k.name().to_string())
            .collect();
        assert_eq!(keys, vec!["x", "a", "d"]);

        // Within x: parent's b, c before new e
        let x_type = result.lookup(&ident("x").into()).unwrap();
        if let Type::Struct(x_struct) = x_type.as_ref() {
            let x_keys: Vec<_> = x_struct.iter().map(|(k, _)| k.name().to_string()).collect();
            assert_eq!(x_keys, vec!["b", "c", "e"]);
        } else {
            panic!("expected x to be a struct");
        }
    }

    #[test]
    fn test_flatten_roundtrip() {
        let s = Struct::default()
            .with_str(
                "x",
                Struct::default()
                    .with_str("y", INT)
                    .with_str("z", STRING)
                    .into(),
            )
            .with_str("a", BOOLEAN);
        let env = TypeEnvironment::from(s.clone());
        assert_eq!(*env.as_struct(), s);
    }
}
