use std::{iter::once, rc::Rc};

use crate::tree::{
    ast::{identifier::ParsedIdentifier, node::Span, pipeline::Pipeline, query::Query},
    typed_ast::{
        command::SideEffect, context::StatementTranslationContext, environment::TypeEnvironment,
        pipeline::TypedPipeline,
    },
};

/// A type-checked statement
///
/// All statements have WITH clauses (CTEs) and a main pipeline.
/// The side effect describes what mutation the statement performs.
#[derive(Debug, Clone)]
pub struct TypedStatement {
    /// Reference to the original AST query
    pub ast: Rc<Query>,

    /// Type-checked WITH clauses (CTEs)
    ///
    /// CTEs should be pure (no side effects). If a CTE has side effects,
    /// an error is emitted but the CTE is still registered for error recovery
    /// (to avoid cascading "unknown table" errors on references to it).
    pub with_clauses: Vec<TypedWithClause>,

    /// The main pipeline (unified type)
    pub pipeline: Rc<TypedPipeline>,

    /// Side effect of this statement (extracted from final command)
    pub side_effect: SideEffect,
}

/// A type-checked WITH clause (CTE)
///
/// CTEs should be pure (no side effects). If a CTE has side effects,
/// an error is emitted during type-checking but the CTE is still included
/// for error recovery purposes.
#[derive(Debug, Clone)]
pub struct TypedWithClause {
    /// The CTE name
    pub name: ParsedIdentifier,

    /// The typed pipeline for this CTE
    pub pipeline: Rc<TypedPipeline>,
}

// ============================================================================
// TypedStatement Construction
// ============================================================================

impl TypedStatement {
    /// Type-check a query
    ///
    /// Type-checks WITH clauses first (ensuring no side effects), then
    /// type-checks the main pipeline and extracts its side effect.
    pub fn from_ast_with_context(query: Rc<Query>, ctx: &mut StatementTranslationContext) -> Self {
        use crate::tree::ast::query::QueryKind;

        match &query.kind {
            QueryKind::Valid(valid_query) => {
                // Type-check WITH clauses first and register their schemas
                let mut typed_with_clauses = Vec::new();

                for with_clause in &valid_query.with_clauses {
                    let typed_pipeline =
                        TypedPipeline::from_ast_with_context(with_clause.pipeline.clone(), ctx);

                    // Check that CTE has no side effects
                    if !typed_pipeline.side_effect.is_none() {
                        // CTE has side effects - emit error but continue
                        ctx.error("WITH clause cannot have side effects (APPEND not allowed)")
                            .at(with_clause.pipeline.as_ref())
                            .emit();
                    }

                    // Register the CTE schema (only for valid names)
                    if let Ok(valid_name) = with_clause.name.clone().valid() {
                        ctx.register_cte(valid_name, typed_pipeline.environment());
                    }

                    typed_with_clauses.push(TypedWithClause {
                        name: with_clause.name.clone(),
                        pipeline: Rc::new(typed_pipeline),
                    });
                }

                // Type-check the main pipeline (can now reference all CTEs)
                let typed_pipeline =
                    TypedPipeline::from_ast_with_context(valid_query.main_pipeline.clone(), ctx);

                // Side effect comes from the pipeline
                let side_effect = typed_pipeline.side_effect.clone();

                Self {
                    ast: query,
                    with_clauses: typed_with_clauses,
                    pipeline: Rc::new(typed_pipeline),
                    side_effect,
                }
            }
            QueryKind::Error(_) => {
                // Query has parse errors - construct error pipeline
                // Use from_ast_with_context which will produce an Err for empty pipeline
                let empty_pipeline_ast = Rc::new(Pipeline {
                    span: Span::NONE,
                    commands: vec![],
                });
                let typed_pipeline = TypedPipeline::from_ast_with_context(empty_pipeline_ast, ctx);

                Self {
                    ast: query.clone(),
                    with_clauses: vec![],
                    pipeline: Rc::new(typed_pipeline),
                    side_effect: SideEffect::None,
                }
            }
        }
    }

    /// Get the type environment (final schema)
    ///
    /// Returns the output schema. For APPEND statements, this is the schema
    /// of the data being appended (before the APPEND command).
    pub fn environment(&self) -> Rc<TypeEnvironment> {
        self.pipeline.environment()
    }

    /// Get the resolved schema as a Struct (flattens the TypeEnvironment)
    pub fn schema(&self) -> crate::types::struct_type::Struct {
        self.environment().flatten()
    }

    /// Check if this is a pure query (no side effects)
    pub fn is_query(&self) -> bool {
        self.side_effect.is_none()
    }

    /// Check if this is a DML statement (has data modification side effect)
    pub fn is_dml(&self) -> bool {
        self.side_effect.is_dml()
    }

    /// Iterate over all the pipelines in this query
    pub fn iter(&self) -> impl Iterator<Item = &Rc<TypedPipeline>> {
        self.with_clauses
            .iter()
            .map(|wc| &wc.pipeline)
            .chain(once(&self.pipeline))
    }
}
