use std::{ops::RangeInclusive, rc::Rc};

use crate::{err::TranslationError, tree::builder::*};
use derive_more::derive::From;

use crate::{
    sql::expression::identifier::Identifier as SqlIdentifier,
    tree::{
        ast::{
            clause::{
                Assignment, AssignmentClause, FromClause, Selection, TableAlias, TableReference,
            },
            expression::{Expression, ExpressionKind},
            identifier::{
                CompoundIdentifier, Identifier, ParsedIdentifier, ParsedSimpleIdentifier,
            },
            node::Spannable,
        },
        typed_ast::{context::ExpressionTranslationContext, environment::TypeEnvironment},
    },
    types::{matcher::ExactMatcher, struct_type::Struct, BOOLEAN},
};

use super::{context::StatementTranslationContext, expression::TypedExpression};

#[derive(Debug, Clone)]
pub struct Projections {
    pub assignments: Vec<TypedAssignment>,
    /// Parent identifiers that were desugared (need to be tombstoned in output environment)
    pub tombstoned_parents: Vec<Identifier>,
}

#[derive(Debug, Clone)]
pub struct TypedAssignment {
    pub identifier: ParsedIdentifier,
    pub expression: Rc<TypedExpression>,
}

impl Spannable for TypedAssignment {
    fn span(&self) -> Option<RangeInclusive<usize>> {
        // Try to span from start of identifier to end of expression
        // Fall back to just expression span if identifier has no span (e.g., inferred from struct literal)
        match (self.identifier.span(), self.expression.ast.span()) {
            (Some(id_span), Some(expr_span)) => Some(*id_span.start()..=*expr_span.end()),
            (None, Some(expr_span)) => Some(expr_span),
            (Some(id_span), None) => Some(id_span),
            (None, None) => None,
        }
    }
}

impl Projections {
    pub fn from_assignment_clauses_with_context(
        clauses: &Vec<Rc<AssignmentClause>>,
        ctx: &mut ExpressionTranslationContext,
        extends_current_environment: bool,
    ) -> Self {
        // Resolve identifiers from clauses (bare expressions are desugared at parse time)
        let parsed_pairs: Vec<(ParsedIdentifier, Rc<Expression>)> = clauses
            .iter()
            .filter_map(|clause| match clause.as_ref() {
                AssignmentClause::Assignment(assignment) => {
                    Some((assignment.identifier.clone(), assignment.expression.clone()))
                }
                AssignmentClause::Error(_) => None,
            })
            .collect();

        // Expand struct expressions (handles errors internally via .valid_ref())
        let (expanded, tombstoned_parents) = if extends_current_environment {
            Self::expand_struct_expressions(parsed_pairs, ctx.bindings.clone(), ctx)
        } else {
            (parsed_pairs, vec![])
        };

        // Flatten struct literals and type-check
        let assignments = expanded
            .into_iter()
            .flat_map(|(parsed_id, expr)| {
                // Only flatten if valid, otherwise keep as-is
                parsed_id
                    .valid_ref()
                    .map(|identifier| {
                        Self::flatten_struct_assignments(identifier, expr.as_ref())
                            .into_iter()
                            .map(|(id, e)| (ParsedIdentifier::Valid(id), e))
                            .collect::<Vec<_>>()
                    })
                    .unwrap_or_else(|_| vec![(parsed_id, expr)])
            })
            .map(|(parsed_id, expression)| TypedAssignment {
                identifier: parsed_id,
                expression: Rc::new(TypedExpression::from_ast_with_context(expression, ctx)),
            })
            .collect();

        Self {
            assignments,
            tombstoned_parents,
        }
    }

    pub fn from_assignments_with_context(
        assignments: &Vec<Rc<Assignment>>,
        ctx: &mut ExpressionTranslationContext,
    ) -> Self {
        // Collect ParsedIdentifier pairs
        let parsed_pairs: Vec<(ParsedIdentifier, Rc<Expression>)> = assignments
            .iter()
            .map(|assignment| (assignment.identifier.clone(), assignment.expression.clone()))
            .collect();

        // Expand struct expressions (handles errors internally)
        let (expanded, tombstoned_parents) =
            Self::expand_struct_expressions(parsed_pairs, ctx.bindings.clone(), ctx);

        // Flatten struct literals and type-check
        let typed_assignments = expanded
            .into_iter()
            .flat_map(|(parsed_id, expr)| {
                // Only flatten if valid, otherwise keep as-is
                parsed_id
                    .valid_ref()
                    .map(|identifier| {
                        Self::flatten_struct_assignments(identifier, expr.as_ref())
                            .into_iter()
                            .map(|(id, e)| (ParsedIdentifier::Valid(id), e))
                            .collect::<Vec<_>>()
                    })
                    .unwrap_or_else(|_| vec![(parsed_id, expr)])
            })
            .map(|(parsed_id, expression)| TypedAssignment {
                identifier: parsed_id,
                expression: Rc::new(TypedExpression::from_ast_with_context(expression, ctx)),
            })
            .collect();

        Self {
            assignments: typed_assignments,
            tombstoned_parents,
        }
    }

    /// Helper: flatten struct literal assignments into compound identifier assignments
    ///
    /// For example, `x = {y: {z: 1, a: 2}}` becomes:
    /// - `x.y.z = 1`
    /// - `x.y.a = 2`
    fn flatten_struct_assignments(
        base_identifier: &Identifier,
        expression: &Expression,
    ) -> Vec<(Identifier, Rc<Expression>)> {
        match &expression.kind {
            ExpressionKind::StructLiteral(struct_lit) => {
                let mut result = Vec::new();
                for (maybe_field_name, field_expr) in struct_lit.fields.iter() {
                    let field_name = match maybe_field_name.clone().valid() {
                        Ok(f) => f,
                        Err(_) => continue,
                    };

                    // Build compound identifier: base.field
                    let compound_id = match base_identifier {
                        Identifier::Simple(simple) => {
                            CompoundIdentifier::new(simple.clone(), field_name.clone(), vec![])
                                .into()
                        }
                        Identifier::Compound(compound) => {
                            let mut parts = compound.parts.clone();
                            parts.push(field_name.clone());
                            CompoundIdentifier { parts }.into()
                        }
                    };

                    // Recursively flatten if the field value is also a struct literal
                    let nested = Self::flatten_struct_assignments(&compound_id, field_expr);
                    result.extend(nested);
                }
                result
            }
            _ => {
                // Not a struct literal - return as-is
                vec![(base_identifier.clone(), expression.clone().into())]
            }
        }
    }

    /// Shared helper: expand struct expressions by desugaring parent struct bindings
    ///
    /// For each assignment, checks if a parent prefix has a struct binding in the lookup environment.
    /// If yes, expands the struct into all its fields (with the target field first).
    /// If no, keeps the assignment as-is.
    ///
    /// Returns (expanded_assignments, tombstoned_parents)
    fn expand_struct_expressions(
        assignments: Vec<(ParsedIdentifier, Rc<Expression>)>,
        lookup_env: Rc<TypeEnvironment>,
        ctx: &mut ExpressionTranslationContext,
    ) -> (Vec<(ParsedIdentifier, Rc<Expression>)>, Vec<Identifier>) {
        let mut expanded = Vec::new();
        let mut tombstoned_parents = Vec::new();

        for (parsed_id, expression) in assignments {
            // Extract valid identifier for struct operations, or pass through errors unchanged
            let identifier = match parsed_id.valid_ref() {
                Ok(id) => id,
                Err(_) => {
                    // Error identifier - preserve as-is without processing
                    expanded.push((parsed_id, expression));
                    continue;
                }
            };

            // Check if any parent prefix has a direct struct binding in lookup environment
            match lookup_env.lookup_struct(identifier) {
                Ok(Some((parent, parent_struct))) => {
                    // Parent has struct binding - must desugar

                    // Add the target field FIRST (modified field comes to front)
                    expanded.push((parsed_id.clone(), expression.clone()));

                    // Iterate through ALL flattened fields of the parent struct
                    // Add non-target fields (preserving their relative order)
                    for (field_id, _field_type) in parent_struct.iter_flattened() {
                        // Convert SQL identifier to AST identifier and prepend parent
                        let field_id: Identifier = field_id.into();
                        let full_id = parent.clone() + field_id;

                        if &full_id == identifier {
                            // Target field - already added at the front, skip
                            continue;
                        } else if full_id.has_prefix(identifier) {
                            // This field is a child of the target - will be replaced
                            continue;
                        } else {
                            // Preserve existing field - create column reference expression
                            let col_ref_expr = Self::create_column_reference_expression(&full_id);
                            expanded
                                .push((ParsedIdentifier::Valid(full_id), Rc::new(col_ref_expr)));
                        }
                    }

                    // Track parent for tombstoning
                    if !tombstoned_parents.contains(&parent) {
                        tombstoned_parents.push(parent);
                    }
                }
                Ok(None) => {
                    // All prefixes are clean - no desugaring needed
                    expanded.push((parsed_id, expression));
                }
                Err(env_error) => {
                    ctx.error(format!("{}", env_error)).at(identifier).emit();
                    // Don't add to expanded - skip this assignment
                }
            }
        }

        (expanded, tombstoned_parents)
    }

    /// Helper: create an expression that references a column by identifier
    ///
    /// This is used when desugaring struct assignments to preserve existing fields.
    /// For simple identifiers (x), creates a ColumnReference.
    /// For compound identifiers (x.y.z), creates a ColumnReference wrapped in FieldLookup chain.
    pub fn create_column_reference_expression(id: &Identifier) -> Expression {
        match id {
            Identifier::Simple(simple) => {
                // Simple identifier: just a column reference
                column_ref(simple.as_str()).build()
            }
            Identifier::Compound(compound) => {
                // Compound identifier: column reference + field lookups
                // Start with a column reference to the first part
                let builder = column_ref(compound.parts[0].as_str());

                // Chain field lookups for each subsequent part
                let first_field = &compound.parts[1];
                let mut current = field(builder, first_field.as_str());

                for field_part in &compound.parts[2..] {
                    current = field(current, field_part.as_str());
                }

                current.build()
            }
        }
    }

    pub fn types(&self) -> Struct {
        let mut new_fields = Struct::default();
        for typed_assignment in &self.assignments {
            // Convert AST identifier to SQL identifier (only for valid identifiers)
            if let Ok(identifier) = typed_assignment.identifier.clone().valid() {
                let sql_ident = SqlIdentifier::from(identifier);
                new_fields.insert(
                    sql_ident,
                    (*typed_assignment.expression.resolved_type).clone(),
                );
            }
            // Skip error identifiers - they don't contribute to the type environment
        }

        new_fields
    }

    pub fn lookup(&self, id: &Identifier) -> Option<&TypedAssignment> {
        self.assignments
            .iter()
            .find(|assignment| assignment.identifier.valid_ref().ok() == Some(id))
    }
}

// ============================================================================
// ResolvedEnvironment - Catalog Resolution Result
// ============================================================================

/// Result of resolving a table reference from the catalog
///
/// Separates resolution concerns from aliasing concerns.
/// Resolution errors (table not found) are stored here.
#[derive(Debug, Clone, From)]
pub enum ResolvedEnvironment {
    /// Table was found in catalog
    Resolved(Rc<TypeEnvironment>),

    /// Table not found in catalog
    Error(Rc<TranslationError>),
}

impl ResolvedEnvironment {
    /// Get the environment, or empty if error
    pub fn environment(&self) -> Rc<TypeEnvironment> {
        match self {
            ResolvedEnvironment::Resolved(env) => env.clone(),
            ResolvedEnvironment::Error(_) => Rc::new(TypeEnvironment::default()),
        }
    }
}

// ============================================================================
// TypedFromClause - Type-checked FROM clause
// ============================================================================

/// A type-checked FROM clause
///
/// Three-variant enum that models the AST structure:
/// - Reference: Plain table reference (FROM events)
/// - Alias: Aliased table reference (FROM x = events)
/// - Error: Clause-level structural errors (e.g., schema merge failure)
#[derive(Debug, Clone, From)]
pub enum TypedFromClause {
    Reference(TypedTableReference),
    Alias(TypedTableAlias),
    Error(Rc<TranslationError>),
}

impl TypedFromClause {
    /// Type-check a FROM clause by resolving it from the catalog
    pub fn from_ast_with_context(
        clause: Rc<FromClause>,
        ctx: &mut StatementTranslationContext,
    ) -> Self {
        match clause.as_ref() {
            FromClause::TableReference(table_ref) => {
                match TypedTableReference::try_from_ast_with_context(table_ref.clone(), ctx) {
                    Ok(typed_ref) => typed_ref.into(),
                    Err(err) => err.into(),
                }
            }
            FromClause::TableAlias(table_alias) => {
                match TypedTableAlias::try_from_ast_with_context(table_alias.clone(), ctx) {
                    Ok(typed_alias) => typed_alias.into(),
                    Err(err) => err.into(),
                }
            }
        }
    }

    /// Get the resolved environment (or empty if error)
    pub fn environment(&self) -> Rc<TypeEnvironment> {
        match self {
            TypedFromClause::Reference(r) => r.resolved.environment(),
            TypedFromClause::Alias(a) => a.resolved.environment(),
            TypedFromClause::Error(_) => Rc::new(TypeEnvironment::default()),
        }
    }

    /// Get the alias if this is an Alias variant
    pub fn alias(&self) -> Option<&ParsedSimpleIdentifier> {
        match self {
            TypedFromClause::Alias(a) => Some(&a.alias),
            _ => None,
        }
    }

    /// Get the table reference
    pub fn reference(&self) -> Result<&TableReference, &Rc<TranslationError>> {
        match self {
            TypedFromClause::Reference(r) => Ok(&r.ast),
            TypedFromClause::Alias(a) => Ok(&a.ast.table),
            TypedFromClause::Error(translation_error) => Err(translation_error),
        }
    }
}

// ============================================================================
// TypedTableReference - Plain table reference without alias
// ============================================================================

/// A type-checked table reference without explicit alias
///
/// Models "FROM events" syntax - no alias, fields merged at top level.
#[derive(Debug, Clone)]
pub struct TypedTableReference {
    /// Reference to the AST table reference
    pub ast: Rc<TableReference>,

    /// The resolved environment (or error if table not found)
    pub resolved: ResolvedEnvironment,
}

impl TypedTableReference {
    /// Type-check a table reference by resolving it from the catalog
    /// Returns Err if the identifier itself is invalid
    pub fn try_from_ast_with_context(
        table_ref: Rc<TableReference>,
        ctx: &mut StatementTranslationContext,
    ) -> Result<Self, Rc<TranslationError>> {
        let identifier = table_ref.identifier.valid_ref()?;
        let resolved = Self::resolve_table_schema(&identifier, ctx, table_ref.as_ref());

        Ok(Self {
            ast: table_ref,
            resolved,
        })
    }

    /// Resolve a table schema from CTEs first, then fall back to catalog provider
    ///
    /// Returns the resolved environment or an error if not found.
    fn resolve_table_schema(
        ident: &Identifier,
        ctx: &mut StatementTranslationContext,
        error_node: &impl Spannable,
    ) -> ResolvedEnvironment {
        // First check if this is a CTE reference
        if let Some(cte_schema) = ctx.lookup_cte(ident) {
            // CTE found - use its schema
            ResolvedEnvironment::Resolved(cte_schema)
        } else {
            // Not a CTE - resolve from catalog provider
            // Convert AST identifier to SQL identifier only for catalog lookup
            let sql_ident = match SqlIdentifier::try_from(ident.clone()) {
                Ok(id) => id,
                Err(_) => {
                    let err = ctx
                        .error(format!("Invalid table identifier: {:?}", ident))
                        .at(error_node)
                        .emit();
                    return err.into();
                }
            };

            let sql_table_ref = crate::sql::query::TableReference {
                name: sql_ident,
                alias: None,
            };

            match ctx.provider.reflect_columns(sql_table_ref) {
                Ok(old_env) => {
                    // Convert old Environment to TypeEnvironment by flattening
                    ResolvedEnvironment::Resolved(Rc::new(old_env.into()))
                }
                Err(e) => {
                    let err = ctx
                        .error(format!("Dataset not found: {}", e))
                        .at(error_node)
                        .emit();
                    err.into()
                }
            }
        }
    }
}

// ============================================================================
// TypedTableAlias - Aliased table reference
// ============================================================================

/// A type-checked table reference with explicit alias
///
/// Models "FROM x = events" syntax - alias is required and stored.
#[derive(Debug, Clone)]
pub struct TypedTableAlias {
    /// Reference to the AST table alias
    pub ast: Rc<TableAlias>,

    /// The alias identifier (simple identifier, matches grammar)
    pub alias: ParsedSimpleIdentifier,

    /// The resolved environment (or error if table not found)
    pub resolved: ResolvedEnvironment,
}

impl TypedTableAlias {
    /// Type-check a table alias by resolving it from the catalog
    /// Returns Err if the identifier itself is invalid
    fn try_from_ast_with_context(
        table_alias: Rc<TableAlias>,
        ctx: &mut StatementTranslationContext,
    ) -> Result<Self, Rc<TranslationError>> {
        // Extract valid identifiers for resolution, but keep parsed types in struct
        let table_identifier = table_alias.table.identifier.clone().valid()?;

        // Resolve using TypedTableReference's logic
        let resolved =
            TypedTableReference::resolve_table_schema(&table_identifier, ctx, table_alias.as_ref());

        Ok(Self {
            ast: table_alias.clone(),
            alias: table_alias.alias.clone(),
            resolved,
        })
    }
}

// ============================================================================
// TypedConstraint - WHEN clause constraints
// ============================================================================

/// A type-checked constraint from the WHEN clause
#[derive(Debug, Clone)]
pub struct TypedConstraint {
    /// Reference to the AST assignment
    pub ast: Rc<Assignment>,

    /// The pattern name (left side of assignment)
    pub identifier: ParsedIdentifier,

    /// The constraint expression (right side of assignment)
    pub expression: Rc<TypedExpression>,
}

impl TypedConstraint {
    /// Type-check a constraint from an Assignment
    /// Returns Err if the identifier is invalid
    pub fn try_from_ast_with_context(
        assignment: Rc<Assignment>,
        ctx: &mut ExpressionTranslationContext,
    ) -> Result<Self, Rc<TranslationError>> {
        // Validate that identifier is not a parse error (propagate error if it is)
        assignment.identifier.valid_ref()?;

        // Type-check the constraint expression (must be BOOLEAN)
        let expression = TypedExpression::from_ast_with_expected_type(
            assignment.expression.clone(),
            ctx,
            &ExactMatcher::of(BOOLEAN),
        );

        Ok(Self {
            ast: assignment.clone(),
            identifier: assignment.identifier.clone(),
            expression: Rc::new(expression),
        })
    }
}

// ============================================================================
// TypedSelection - Field selections (used in DISTINCT BY, DROP, etc.)
// ============================================================================

/// A type-checked selection (field identifier with schema validation)
#[derive(Debug, Clone)]
pub struct TypedSelection {
    /// Reference to the AST selection
    pub ast: Rc<Selection>,

    /// The resolved identifier (validated against schema) or error
    pub resolved: ResolvedSelection,
}

/// Result of resolving a selection against a schema
#[derive(Debug, Clone, From)]
pub enum ResolvedSelection {
    /// Successfully resolved - identifier exists in merged schema.
    Resolved(Identifier),

    /// Resolution failed - identifier not found or invalid
    Error(Rc<TranslationError>),
}

impl TypedSelection {
    /// Type-check a selection against an environment
    ///
    /// Validates that the identifier exists in the schema.
    pub fn from_ast_with_context(
        selection: Rc<Selection>,
        input_schema: &Rc<TypeEnvironment>,
        ctx: &mut StatementTranslationContext,
    ) -> Self {
        let resolved = match selection.identifier.clone().valid() {
            Ok(identifier) => {
                // Validate that the field exists in the input schema
                if input_schema.lookup(&identifier).is_some() {
                    ResolvedSelection::Resolved(identifier.clone())
                } else {
                    let err = ctx
                        .error(format!("Field '{:?}' not found", identifier))
                        .at(selection.as_ref())
                        .emit();
                    ResolvedSelection::Error(err)
                }
            }
            Err(err) => ResolvedSelection::Error(err),
        };

        Self {
            ast: selection,
            resolved,
        }
    }
}
