use std::{iter::once, sync::Arc};

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

/// A type-checked statement
///
/// All statements have DEF pipelines (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: Arc<Query>,

    /// Tabular `DEF` bindings in order (user and compiler-generated), each with a typed pipeline.
    pub pipeline_defs: Vec<TypedPipelineDef>,

    /// Scalar `DEF` bindings visible for this statement
    pub scalar_defs: Vec<TypedScalarDef>,

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

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

/// A type-checked tabular `DEF` (`DefBody::Pipeline`).
#[derive(Debug, Clone)]
pub struct TypedPipelineDef {
    pub name: ParsedIdentifier,
    pub pipeline: Arc<TypedPipeline>,
}

/// A type-checked scalar `DEF` (`DefBody::Expression`).
#[derive(Debug, Clone)]
pub struct TypedScalarDef {
    pub name: ParsedIdentifier,
    pub expression: Arc<TypedExpression>,
}

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

impl TypedStatement {
    /// Type-check a parsed query, including `DEF` statements in source order.
    pub fn from_ast_with_context(query: Arc<Query>, ctx: &mut StatementTranslationContext) -> Self {
        use crate::tree::ast::query::QueryKind;

        match &query.kind {
            QueryKind::Valid(valid_query) => {
                let (pipeline_defs, scalar_defs) = typecheck_defs(&valid_query.defs, ctx);

                let typed_pipeline =
                    TypedPipeline::from_ast_with_context(valid_query.main_pipeline.clone(), ctx);

                let side_effect = typed_pipeline.side_effect.clone();

                Self {
                    ast: query,
                    pipeline_defs,
                    scalar_defs,
                    pipeline: Arc::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 = Arc::new(Pipeline {
                    span: Span::NONE,
                    commands: vec![],
                });
                let typed_pipeline = TypedPipeline::from_ast_with_context(empty_pipeline_ast, ctx);

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

    /// Get the type environment (final schema)
    ///
    /// Returns the output schema of the last command in the pipeline.
    /// For APPEND statements (which must be in final position), this is
    /// the target table's schema, not the pipeline's pre-APPEND schema.
    pub fn environment(&self) -> Arc<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().as_struct().clone()
    }

    /// Get the environment in scope at a codepoint offset.
    ///
    /// This is used by editor features (autocomplete) to resolve bindings
    /// from the command currently being edited.
    pub fn environment_at_offset(&self, offset: usize) -> Option<Arc<TypeEnvironment>> {
        let mut last_spanned = None;
        let mut last_any = None;

        for cmd in self
            .iter()
            .flat_map(|pipeline| pipeline.commands().into_iter().flatten())
        {
            match cmd.ast.span.to_range() {
                Some(span) if span.start() <= &offset => {
                    last_spanned = Some(cmd);
                    last_any = Some(cmd);
                }
                Some(_) => {}
                None => {
                    last_any = Some(cmd);
                }
            }
        }

        last_spanned
            .or(last_any)
            .map(|cmd| cmd.input_schema.clone())
    }

    /// 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 = &Arc<TypedPipeline>> {
        self.pipeline_defs
            .iter()
            .map(|d| &d.pipeline)
            .chain(once(&self.pipeline))
    }
}

fn typecheck_defs(
    defs: &[DefStatement],
    ctx: &mut StatementTranslationContext,
) -> (Vec<TypedPipelineDef>, Vec<TypedScalarDef>) {
    let mut pipeline_defs = Vec::new();
    let mut scalar_defs: Vec<TypedScalarDef> = Vec::new();

    for def in defs {
        match &def.body {
            DefBody::Pipeline(pipeline) => {
                let typed_pipeline = TypedPipeline::from_ast_with_context(pipeline.clone(), ctx);
                if !typed_pipeline.side_effect.is_none() {
                    ctx.error("DEF cannot have side effects (APPEND not allowed)")
                        .at(pipeline.as_ref())
                        .emit();
                }
                if let Ok(valid_name) = def.name.clone().valid() {
                    ctx.register_cte(valid_name, typed_pipeline.environment());
                }
                pipeline_defs.push(TypedPipelineDef {
                    name: def.name.clone(),
                    pipeline: Arc::new(typed_pipeline),
                });
            }
            DefBody::Expression(expr) => {
                // Empty bindings: prior scalar DEFs are resolved only via
                // `register_scalar_def` / `lookup_scalar_def_type` so field refs get
                // `FieldReferenceResolution::ScalarDef`, not `Pipeline`.
                let empty_bindings = Arc::new(TypeEnvironment::default());
                let mut expr_ctx = ctx.default_expression_context(&empty_bindings);
                let typed_expr =
                    TypedExpression::from_ast_with_context(expr.clone(), &mut expr_ctx);
                if let Ok(valid_name) = def.name.clone().valid() {
                    ctx.register_scalar_def(valid_name, typed_expr.resolved_type.clone());
                }
                let expression = Arc::new(typed_expr);
                scalar_defs.push(TypedScalarDef {
                    name: def.name.clone(),
                    expression: Arc::clone(&expression),
                });
            }
        }
    }

    (pipeline_defs, scalar_defs)
}
