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

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

/// 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>,

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

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

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

    /// Diagnostic sink - accumulates errors during TypedAST construction
    /// Stores Rc<TranslationError> for efficient sharing without cloning
    errors: Vec<Rc<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,
            timestamp_field: SimpleIdentifier::new("timestamp").into(),
            message_field: SimpleIdentifier::new("_raw").into(),
            cte_schemas: HashMap::new(),
            errors: Vec::new(),
        }
    }

    /// 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
    }

    /// 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: &Rc<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: &Rc<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: Rc<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<Rc<TypeEnvironment>> {
        self.cte_schemas.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: Rc<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<Rc<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: Rc<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: Rc<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: Rc<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: Rc<TranslationError>) {
        // Delegate to the public add_error method
        self.add_error(error);
    }
}
