use std::marker::PhantomData;
use std::rc::Rc;

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

use super::{ExpressionBuilder, IntoExpressionBuilder};

// Typestate markers for LambdaBuilder
#[derive(Debug)]
pub struct NoBody;
#[derive(Debug)]
pub struct WithBody;

/// Builder for lambda expressions
///
/// Uses typestate pattern to ensure body is provided before building.
#[derive(Debug)]
pub struct LambdaBuilder<State = NoBody> {
    parameters: Vec<LambdaParameter>,
    body: Option<Box<dyn ExpressionBuilder>>,
    _state: PhantomData<State>,
}

impl LambdaBuilder<NoBody> {
    pub fn new() -> Self {
        Self {
            parameters: Vec::new(),
            body: None,
            _state: PhantomData,
        }
    }

    /// 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: 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) -> LambdaBuilder<WithBody> {
        LambdaBuilder {
            parameters: self.parameters,
            body: Some(body.into_expression_builder()),
            _state: PhantomData,
        }
    }
}

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

impl LambdaBuilder<WithBody> {
    /// 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: 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 = Some(body.into_expression_builder());
        self
    }
}

impl ExpressionBuilder for LambdaBuilder<WithBody> {
    fn build(&self) -> Expression {
        let body = self
            .body
            .as_ref()
            .map(|b| Rc::new(b.build()))
            .expect("body is guaranteed by typestate");

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

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

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

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

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