use std::sync::Arc;

use crate::{
    err::TranslationError,
    tree::{
        ast::{
            command::{Command, CommandClass, CommandKind, WithinCommand},
            node::Span,
            pipeline::Pipeline,
        },
        typed_ast::{
            command::{SideEffect, TypedCommand, TypedCommandKind},
            context::StatementTranslationContext,
            environment::TypeEnvironment,
        },
    },
    types::Type,
};

/// A valid typed pipeline with commands and final schema
#[derive(Debug, Clone)]
pub struct ValidPipeline {
    /// The typed commands in execution order
    pub commands: Vec<Arc<TypedCommand>>,
    /// The final schema after all commands
    pub final_schema: Arc<TypeEnvironment>,
}

/// A type-checked pipeline
///
/// TypedPipeline represents a sequence of commands that transform data.
/// It tracks the flow of schemas through the pipeline and has a final schema.
///
/// If the final command has a side effect (like APPEND), it's recorded here.
///
/// The `kind` field is `Err` for pipeline-level errors (e.g., empty pipeline).
/// Command-level errors are stored in individual `TypedCommand` nodes as
/// `TypedCommandKind::Error`, allowing error recovery and continued type-checking.
///
/// Example:
/// ```hamelin
/// FROM events
/// | WHERE severity > 3
/// | SET severity_label = if(severity > 5, "critical", "normal")
/// | SELECT timestamp, severity, severity_label
/// ```
#[derive(Debug, Clone)]
pub struct TypedPipeline {
    /// Reference to the original AST pipeline
    pub ast: Arc<Pipeline>,

    /// The pipeline contents - `Err` for empty pipelines, `Ok` otherwise.
    pub kind: Result<ValidPipeline, Arc<TranslationError>>,

    /// Side effect of this pipeline (from final command, if any)
    pub side_effect: SideEffect,
}

// ============================================================================
// TypedPipeline Construction
// ============================================================================

impl TypedPipeline {
    /// Type-check a pipeline
    ///
    /// Type-checks each command in sequence, threading schemas through.
    /// If the final command has a side effect, it's recorded in the pipeline.
    ///
    /// Side-effecting commands (DML/DDL like APPEND) must be the final command.
    /// If they appear earlier, an error is emitted.
    ///
    /// Returns an error pipeline if no commands were produced (empty input or
    /// all commands failed to parse).
    pub fn from_ast_with_context(
        pipeline: Arc<Pipeline>,
        ctx: &mut StatementTranslationContext,
    ) -> Self {
        let mut current_schema = Arc::new(TypeEnvironment::default());
        let mut typed_commands = Vec::new();
        let last_idx = pipeline.commands.len().saturating_sub(1);
        let time_range = ctx.time_range.clone();

        // Type-check each command in sequence
        for (i, ast_command) in pipeline.commands.iter().enumerate() {
            // Check if this is a source command (FROM/UNION/ROWS) in non-first position
            if ast_command.kind.is_source_command() && i > 0 {
                let command_name = ast_command.kind.command_name();
                let error = ctx
                    .error(format!(
                        "{} must be the first command in a pipeline",
                        command_name
                    ))
                    .at(ast_command.as_ref())
                    .emit();
                let typed_command = TypedCommand {
                    ast: ast_command.clone(),
                    input_schema: current_schema.clone(),
                    output_schema: current_schema.clone(),
                    kind: error.into(),
                };
                typed_commands.push(Arc::new(typed_command));
                continue;
            }

            // Check if this is a misplaced non-query command (DML/DDL in non-final position)
            let command_class = ast_command.kind.command_class();
            let is_non_query = command_class != CommandClass::Query;
            if is_non_query && i != last_idx {
                // Create error command for misplaced non-query command
                let command_name = ast_command.kind.command_name();
                let error = ctx
                    .error(format!(
                        "{} command must be the final command in a pipeline",
                        command_name
                    ))
                    .at(ast_command.as_ref())
                    .emit();
                let typed_command = TypedCommand {
                    ast: ast_command.clone(),
                    input_schema: current_schema.clone(),
                    output_schema: current_schema.clone(), // Preserve schema for error recovery
                    kind: error.into(),
                };
                typed_commands.push(Arc::new(typed_command));
                continue;
            }

            let typed_command = TypedCommand::from_ast_with_context(
                ast_command.clone(),
                current_schema.clone(),
                ctx,
            );

            // Update the current schema for the next command
            current_schema = typed_command.output_schema.clone();
            typed_commands.push(Arc::new(typed_command));

            let should_insert = time_range.is_some()
                && typed_commands.last().is_some_and(|c| {
                    matches!(
                        &c.kind,
                        TypedCommandKind::From(_)
                            | TypedCommandKind::Union(_)
                            | TypedCommandKind::Match(_)
                    )
                })
                && schema_has_timestamp(&current_schema, ctx)
                && !next_is_within_ast(&pipeline.commands, i);

            if let Some(time_range) = time_range.clone() {
                if should_insert {
                    let within_command = Command {
                        span: Span::NONE,
                        kind: WithinCommand {
                            duration: time_range.clone(),
                        }
                        .into(),
                    };
                    let typed_within = TypedCommand::from_ast_with_context(
                        Arc::new(within_command),
                        current_schema.clone(),
                        ctx,
                    );
                    current_schema = typed_within.output_schema.clone();
                    typed_commands.push(Arc::new(typed_within));
                }
            }
        }

        // Empty pipelines are invalid (no commands, or all commands failed to parse)
        if typed_commands.is_empty() {
            let error = ctx
                .error("Pipeline must have at least one command")
                .at(pipeline.as_ref())
                .emit();
            return TypedPipeline {
                ast: pipeline,
                kind: Err(error),
                side_effect: SideEffect::None,
            };
        }

        // Extract side effect from final command
        let side_effect = typed_commands
            .last()
            .and_then(|cmd| cmd.kind.side_effect())
            .unwrap_or(SideEffect::None);

        TypedPipeline {
            ast: pipeline,
            kind: Ok(ValidPipeline {
                commands: typed_commands,
                final_schema: current_schema,
            }),
            side_effect,
        }
    }

    /// Get the type environment (final schema)
    pub fn environment(&self) -> Arc<TypeEnvironment> {
        match &self.kind {
            Ok(valid) => valid.final_schema.clone(),
            Err(_) => Arc::new(TypeEnvironment::default()),
        }
    }

    /// 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 commands if this is a valid pipeline
    pub fn commands(&self) -> Option<&Vec<Arc<TypedCommand>>> {
        self.kind.as_ref().ok().map(|v| &v.commands)
    }

    pub fn valid_ref(&self) -> Result<&ValidPipeline, Arc<TranslationError>> {
        match &self.kind {
            Ok(valid) => Ok(valid),
            Err(err) => Err(err.clone()),
        }
    }
}

fn schema_has_timestamp(env: &TypeEnvironment, ctx: &StatementTranslationContext) -> bool {
    env.lookup(&ctx.timestamp_field)
        .is_some_and(|typ| matches!(typ.as_ref(), Type::Timestamp))
}

fn next_is_within_ast(commands: &[Arc<Command>], idx: usize) -> bool {
    commands
        .get(idx.saturating_add(1))
        .is_some_and(|cmd| matches!(cmd.kind, CommandKind::Within(_)))
}
