use std::sync::Arc;

use crate::tree::ast::expression::{Expression, Lambda, LambdaParameter};
use crate::tree::ast::node::Span;
use crate::types::Type;

use super::{ExpressionBuilder, IntoExpressionBuilder};

/// Builder for lambda expressions (before body is set)
#[derive(Debug)]
pub struct LambdaBuilder {
    parameters: Vec<LambdaParameter>,
}

impl LambdaBuilder {
    pub fn new() -> Self {
        Self {
            parameters: Vec::new(),
        }
    }

    /// Add a parameter without a type hint
    pub fn param(mut self, name: impl Into<String>) -> Self {
        self.parameters.push(LambdaParameter::new(name));
        self
    }

    /// Add a parameter with a type hint
    pub fn param_with_type(mut self, name: impl Into<String>, typ: Arc<Type>) -> Self {
        self.parameters.push(LambdaParameter::with_type(name, typ));
        self
    }

    /// Set the lambda body (transitions to WithBody state)
    pub fn body(self, body: impl IntoExpressionBuilder) -> LambdaBuilderWithBody {
        LambdaBuilderWithBody {
            parameters: self.parameters,
            body: body.into_expression_builder(),
        }
    }
}

impl Default for LambdaBuilder {
    fn default() -> Self {
        Self::new()
    }
}

/// Builder for lambda expressions (after body is set)
#[derive(Debug)]
pub struct LambdaBuilderWithBody {
    parameters: Vec<LambdaParameter>,
    body: Box<dyn ExpressionBuilder>,
}

impl LambdaBuilderWithBody {
    /// Add a parameter without a type hint
    pub fn param(mut self, name: impl Into<String>) -> Self {
        self.parameters.push(LambdaParameter::new(name));
        self
    }

    /// Add a parameter with a type hint
    pub fn param_with_type(mut self, name: impl Into<String>, typ: Arc<Type>) -> Self {
        self.parameters.push(LambdaParameter::with_type(name, typ));
        self
    }

    /// Replace the lambda body
    pub fn body(mut self, body: impl IntoExpressionBuilder) -> Self {
        self.body = body.into_expression_builder();
        self
    }
}

impl ExpressionBuilder for LambdaBuilderWithBody {
    fn build(&self) -> Expression {
        let body = Arc::new(self.body.build());

        Expression {
            span: Span::NONE,
            kind: Lambda {
                parameters: self.parameters.clone(),
                body,
            }
            .into(),
        }
    }
}

impl IntoExpressionBuilder for LambdaBuilderWithBody {
    fn into_expression_builder(self) -> Box<dyn ExpressionBuilder> {
        Box::new(self)
    }
}

/// Create a lambda builder
pub fn lambda() -> LambdaBuilder {
    LambdaBuilder::new()
}

/// Create a single-parameter lambda (convenience)
pub fn lambda1(param: impl Into<String>) -> LambdaBuilder {
    LambdaBuilder::new().param(param)
}

/// Create a two-parameter lambda (convenience)
pub fn lambda2(param1: impl Into<String>, param2: impl Into<String>) -> LambdaBuilder {
    LambdaBuilder::new().param(param1).param(param2)
}
