use std::{collections::HashMap, sync::Arc};

use crate::{
    err::{DiagnosticSink, TranslationError, TranslationErrors},
    func::{def::FunctionTranslationContext, registry::FunctionRegistry},
    interner::Interner,
    provider::EnvironmentProvider,
    tree::{
        ast::{
            context::ErrorBuilder,
            expression::Expression,
            identifier::{CompoundIdentifier, Identifier, SimpleIdentifier},
        },
        typed_ast::environment::TypeEnvironment,
    },
    types::Type,
};

/// Context needed for type-checking pipelines and queries
///
/// This context contains pipeline-wide information like the function registry
/// and environment provider for catalog lookups.
#[derive(Debug, Clone)]
pub struct StatementTranslationContext {
    /// Function registry for resolving function calls
    pub registry: Arc<FunctionRegistry>,

    /// Environment provider for catalog/schema lookups (used by FROM commands)
    pub provider: Arc<dyn EnvironmentProvider>,

    /// Type interner for deduplicating types in the TypedAST
    pub interner: Arc<Interner>,

    /// The timestamp field identifier used by WITHIN (typically "timestamp")
    pub timestamp_field: Identifier,

    /// The message field identifier used by PARSE (typically "event.original")
    pub message_field: Identifier,

    /// Optional time range expression for injecting WITHIN after FROM
    pub time_range: Option<Arc<Expression>>,

    /// Whether to lower filter() calls to flatten(transform(...)) pattern.
    /// Needed for backends that don't have a native array filter function.
    pub lower_filter: bool,

    /// Whether to lower transform() calls to EXPLODE + AGG pattern.
    /// Needed for backends like DataFusion that don't support lambdas.
    pub lower_transform: bool,

    /// Skip statement-level normalization passes (lower_match, nest_from_aliases, etc.)
    pub skip_statement_passes: bool,

    /// Skip all pipeline-level normalization passes (implies skip_projection_fusion)
    pub skip_pipeline_passes: bool,

    /// Skip just the fuse_projections pass
    pub skip_projection_fusion: bool,

    /// CTE registry - maps CTE names to their resolved schemas
    /// CTEs are checked first before falling back to the catalog provider
    cte_schemas: HashMap<Identifier, Arc<TypeEnvironment>>,

    /// Dashboard `${param}` expression types (from typecheck options).
    template_param_types: Arc<HashMap<SimpleIdentifier, Type>>,

    /// String enum domains for parameters used in table paths (for catalog expansion).
    template_param_enums: Arc<HashMap<SimpleIdentifier, Arc<[String]>>>,

    /// Scalar `DEF` types (name → type) for expression resolution.
    scalar_def_types: HashMap<Identifier, Arc<Type>>,

    /// Default space for unqualified dataset references in `FROM` / `JOIN` / etc.
    default_space: Option<SimpleIdentifier>,

    /// Diagnostic sink - accumulates errors during TypedAST construction
    /// Stores Arc<TranslationError> for efficient sharing without cloning
    errors: Vec<Arc<TranslationError>>,
}

impl StatementTranslationContext {
    /// Create a new pipeline translation context with default timestamp field
    pub fn new(registry: Arc<FunctionRegistry>, provider: Arc<dyn EnvironmentProvider>) -> Self {
        Self {
            registry,
            provider,
            interner: Arc::new(Interner::new()),
            timestamp_field: SimpleIdentifier::new("timestamp").into(),
            message_field: CompoundIdentifier::new(
                SimpleIdentifier::new("event"),
                SimpleIdentifier::new("original"),
                vec![],
            )
            .into(),
            time_range: None,
            lower_filter: false,
            lower_transform: false,
            skip_statement_passes: false,
            skip_pipeline_passes: false,
            skip_projection_fusion: false,
            cte_schemas: HashMap::new(),
            template_param_types: Arc::new(HashMap::new()),
            template_param_enums: Arc::new(HashMap::new()),
            scalar_def_types: HashMap::new(),
            default_space: None,
            errors: Vec::new(),
        }
    }

    /// Set the default space for unqualified dataset references.
    pub fn with_default_space(mut self, default_space: SimpleIdentifier) -> Self {
        self.default_space = Some(default_space);
        self
    }

    pub fn default_space(&self) -> Option<&SimpleIdentifier> {
        self.default_space.as_ref()
    }

    /// Set template parameter types for `${name}` in expressions.
    pub fn with_template_param_types(
        mut self,
        template_param_types: Arc<HashMap<SimpleIdentifier, Type>>,
    ) -> Self {
        self.template_param_types = template_param_types;
        self
    }

    /// Set string enum domains for template parameters used in table paths.
    pub fn with_template_param_enums(
        mut self,
        template_param_enums: Arc<HashMap<SimpleIdentifier, Arc<[String]>>>,
    ) -> Self {
        self.template_param_enums = template_param_enums;
        self
    }

    /// Look up the declared type for a template parameter, if any.
    pub fn lookup_template_param_type(&self, name: &SimpleIdentifier) -> Option<Type> {
        self.template_param_types.get(name).cloned()
    }

    /// Look up the declared identifier fragment domain for a template parameter, if any.
    pub fn lookup_template_param_enum(&self, name: &SimpleIdentifier) -> Option<Arc<[String]>> {
        self.template_param_enums.get(name).cloned()
    }

    /// Set a custom interner for type deduplication
    pub fn with_interner(mut self, interner: Arc<Interner>) -> Self {
        self.interner = interner;
        self
    }

    /// Set a time range expression to inject WITHIN after FROM commands.
    pub fn with_time_range(mut self, time_range: impl Into<Arc<Expression>>) -> Self {
        self.time_range = Some(time_range.into());
        self
    }

    /// Set a custom timestamp field identifier
    pub fn with_timestamp_field(mut self, timestamp_field: Identifier) -> Self {
        self.timestamp_field = timestamp_field;
        self
    }

    /// Set a custom message field identifier
    pub fn with_message_field(mut self, message_field: Identifier) -> Self {
        self.message_field = message_field;
        self
    }

    /// Enable filter() lowering for backends that don't have a native array filter function.
    pub fn with_lower_filter(mut self) -> Self {
        self.lower_filter = true;
        self
    }

    /// Enable transform() lowering for backends that don't support lambdas.
    pub fn with_lower_transform(mut self) -> Self {
        self.lower_transform = true;
        self
    }

    /// Skip statement-level normalization passes.
    pub fn with_skip_statement_passes(mut self) -> Self {
        self.skip_statement_passes = true;
        self
    }

    /// Skip all pipeline-level normalization passes.
    pub fn with_skip_pipeline_passes(mut self) -> Self {
        self.skip_pipeline_passes = true;
        self
    }

    /// Skip just the fuse_projections pass.
    pub fn with_skip_projection_fusion(mut self) -> Self {
        self.skip_projection_fusion = true;
        self
    }

    /// Create an expression translation context from this pipeline context
    ///
    /// This combines the pipeline-wide context with expression-specific information:
    /// - The current environment (bindings available at this point)
    /// - Function translation context (for window functions, special positions, etc.)
    pub fn expression_context<'ctx>(
        &'ctx mut self,
        bindings: &Arc<TypeEnvironment>,
        fctx: FunctionTranslationContext,
    ) -> ExpressionTranslationContext<'ctx> {
        ExpressionTranslationContext {
            fctx,
            bindings: bindings.clone(),
            registry: self.registry.clone(),
            pipeline_ctx: self,
        }
    }

    /// Create a default expression translation context (with default function context)
    pub fn default_expression_context<'ctx>(
        &'ctx mut self,
        bindings: &Arc<TypeEnvironment>,
    ) -> ExpressionTranslationContext<'ctx> {
        self.expression_context(bindings, FunctionTranslationContext::default())
    }

    /// Register a CTE schema for lookup during type-checking
    ///
    /// This allows subsequent CTEs and the main pipeline to reference earlier CTEs
    /// by name in FROM clauses.
    pub fn register_cte(&mut self, name: Identifier, schema: Arc<TypeEnvironment>) {
        self.cte_schemas.insert(name, schema);
    }

    /// Look up a CTE schema by name
    ///
    /// Returns Some(schema) if the CTE exists, None otherwise.
    /// This is checked before falling back to the catalog provider.
    pub fn lookup_cte(&self, name: &Identifier) -> Option<Arc<TypeEnvironment>> {
        self.cte_schemas.get(name).cloned()
    }

    /// Register a scalar `DEF` (for subsequent expression lookups).
    pub fn register_scalar_def(&mut self, name: Identifier, typ: Arc<Type>) {
        self.scalar_def_types.insert(name, typ);
    }

    /// Look up a scalar `DEF` type by identifier.
    pub fn lookup_scalar_def_type(&self, name: &Identifier) -> Option<Arc<Type>> {
        self.scalar_def_types.get(name).cloned()
    }

    /// Create an error builder for constructing errors with this context
    pub fn error(&mut self, message: impl Into<String>) -> ErrorBuilder<'_, Self> {
        ErrorBuilder::new(message.into(), self)
    }

    /// Add an already-constructed error to the diagnostic sink
    /// Used when propagating errors from AST Error variants
    pub fn add_error(&mut self, error: Arc<TranslationError>) {
        self.errors.push(error);
    }

    /// Take ownership of accumulated errors (consumes self)
    /// Converts to TranslationErrors at the boundary
    pub fn take_errors(self) -> TranslationErrors {
        // Convert Vec<Arc<TranslationError>> to TranslationErrors
        // This is the boundary where we unwrap Rc's for the old API
        let unwrapped: Vec<TranslationError> =
            self.errors.into_iter().map(|rc| (*rc).clone()).collect();
        TranslationErrors(unwrapped)
    }
}

impl Default for StatementTranslationContext {
    fn default() -> Self {
        use crate::provider::NoOpProvider;
        Self::new(
            Arc::new(FunctionRegistry::default()),
            Arc::new(NoOpProvider {}),
        )
    }
}

/// Context needed for type-checking expressions
///
/// This context extends the pipeline context with expression-specific information:
/// - Current bindings (environment at this point in the pipeline)
/// - Function translation context (for special cases like window functions)
pub struct ExpressionTranslationContext<'ctx> {
    /// Function translation context (window reference, special positions, order by)
    pub fctx: FunctionTranslationContext,

    /// Current bindings available to this expression
    pub bindings: Arc<TypeEnvironment>,

    /// Function registry (inherited from pipeline context)
    pub registry: Arc<FunctionRegistry>,

    /// Reference to pipeline context for error emission
    pub pipeline_ctx: &'ctx mut StatementTranslationContext,
}

impl<'ctx> ExpressionTranslationContext<'ctx> {
    /// Create a new expression translation context
    pub fn new(
        bindings: Arc<TypeEnvironment>,
        registry: Arc<FunctionRegistry>,
        fctx: FunctionTranslationContext,
        pipeline_ctx: &'ctx mut StatementTranslationContext,
    ) -> Self {
        Self {
            fctx,
            bindings,
            registry,
            pipeline_ctx,
        }
    }

    /// Create a child context with new bindings and specials_allowed cleared.
    ///
    /// Used for type-checking lambda bodies where aggregates are not allowed.
    pub fn without_specials(
        &mut self,
        bindings: Arc<TypeEnvironment>,
    ) -> ExpressionTranslationContext<'_> {
        let mut fctx = self.fctx.clone();
        fctx.specials_allowed.clear();
        ExpressionTranslationContext {
            bindings,
            fctx,
            registry: self.registry.clone(),
            pipeline_ctx: self.pipeline_ctx,
        }
    }

    /// Create an error builder (delegates to pipeline context)
    pub fn error(
        &mut self,
        message: impl Into<String>,
    ) -> ErrorBuilder<'_, StatementTranslationContext> {
        self.pipeline_ctx.error(message)
    }
}

/// Implement DiagnosticSink for StatementTranslationContext
impl DiagnosticSink for StatementTranslationContext {
    fn add_error(&mut self, error: Arc<TranslationError>) {
        // Delegate to the public add_error method
        self.add_error(error);
    }
}
