//! Lowering pass for BroadcastApply → Transform(array, lambda)
//!
//! This module provides the lowering transformation that rewrites
//! `BroadcastApply` nodes into explicit `transform()` calls.
//!
//! BroadcastApply represents a function call where one argument is an array
//! and the operation should be applied element-wise. During lowering, this
//! becomes: `transform(array_arg, param -> function(param, other_args...))`
//!
//! Example:
//!   Input:  BroadcastApply { function: Add, broadcast_position: 0, args: [array(int), int] }
//!   Output: transform(array_arg, __e -> __e + scalar_arg)

use std::rc::Rc;
use std::sync::Arc;

use crate::func::def::ParameterBinding;
use crate::func::defs::transform::Transform;
use crate::tree::ast::identifier::SimpleIdentifier;
use crate::tree::typed_ast::environment::TypeEnvironment;
use crate::tree::typed_ast::expression::{
    TypedApply, TypedBroadcastApply, TypedExpression, TypedExpressionKind, TypedLambda,
    TypedLambdaParameter,
};
use crate::types::array::Array;
use crate::types::Type;

/// Generate a unique parameter name that doesn't conflict with bindings in scope.
///
/// Starts with `__broadcast_element` and appends a numeric suffix if needed.
fn unique_broadcast_param(bindings: &TypeEnvironment) -> SimpleIdentifier {
    let base = "__broadcast_element";
    if bindings.lookup(&SimpleIdentifier::new(base).into()).is_none() {
        return SimpleIdentifier::new(base);
    }

    let mut counter = 0;
    loop {
        let candidate = format!("{}{}", base, counter);
        if bindings
            .lookup(&SimpleIdentifier::new(&candidate).into())
            .is_none()
        {
            return SimpleIdentifier::new(candidate);
        }
        counter += 1;
    }
}

/// Lower a BroadcastApply into an equivalent Transform apply.
///
/// This creates:
/// - A TypedLambda that applies the original function with the broadcast
///   position replaced by the lambda parameter
/// - A TypedApply of Transform with the array and lambda as arguments
///
/// The resulting TypedExpression has the same resolved_type as the input
/// (which should be Array(element_result_type)).
pub fn lower_broadcast_apply(
    broadcast: &TypedBroadcastApply,
    original_ast: &Rc<crate::tree::ast::expression::Expression>,
    bindings: &TypeEnvironment,
) -> TypedExpression {
    // Generate a unique parameter name
    let param_name = unique_broadcast_param(bindings);

    // Extract the array argument and its element type
    let array_arg = broadcast
        .parameter_binding
        .get_positional(broadcast.broadcast_position)
        .expect("broadcast_position should be valid");

    let element_type = match array_arg.resolved_type.as_ref() {
        Type::Array(arr) => arr.element_type.as_ref().clone(),
        _ => panic!("broadcast argument should be an array"),
    };

    // Build the lambda parameter reference (what replaces the array in the function call)
    let param_ref = TypedExpression {
        ast: original_ast.clone(), // Use original AST for span preservation
        resolved_type: Rc::new(element_type.clone()),
        kind: TypedExpressionKind::ColumnReference(
            crate::tree::typed_ast::expression::TypedColumnReference {
                column: param_name.clone().into(),
            },
        ),
    };

    // Build new arguments with the array replaced by the parameter reference
    let new_positional: Vec<Rc<TypedExpression>> = broadcast
        .parameter_binding
        .positional()
        .iter()
        .enumerate()
        .map(|(i, arg)| {
            if i == broadcast.broadcast_position {
                Rc::new(param_ref.clone())
            } else {
                arg.clone()
            }
        })
        .collect();

    // Build the inner function application (the lambda body)
    let inner_apply = TypedApply {
        function_def: broadcast.function_def.clone(),
        parameter_binding: ParameterBinding::new(
            new_positional,
            broadcast.parameter_binding.named().clone(),
        ),
    };

    let lambda_body = TypedExpression {
        ast: original_ast.clone(),
        resolved_type: broadcast.element_result_type.clone(),
        kind: inner_apply.into(),
    };

    // Build the typed lambda
    let typed_lambda = TypedLambda {
        parameters: vec![TypedLambdaParameter {
            name: param_name,
            typ: element_type,
        }],
        body: Rc::new(lambda_body),
    };

    let lambda_type = crate::types::function::Function::new(
        typed_lambda
            .parameters
            .iter()
            .map(|p| p.typ.clone())
            .collect(),
        broadcast.element_result_type.as_ref().clone(),
    );

    let lambda_expr = TypedExpression {
        ast: original_ast.clone(),
        resolved_type: Rc::new(lambda_type.into()),
        kind: typed_lambda.into(),
    };

    // Build the transform() call
    let transform_binding = ParameterBinding::new(
        vec![array_arg.clone(), Rc::new(lambda_expr)],
        Default::default(),
    );

    let transform_apply = TypedApply {
        function_def: Arc::new(Transform),
        parameter_binding: transform_binding,
    };

    // The result type is Array(element_result_type)
    let result_type = Array::new(broadcast.element_result_type.as_ref().clone());

    TypedExpression {
        ast: original_ast.clone(),
        resolved_type: Rc::new(result_type.into()),
        kind: transform_apply.into(),
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_unique_broadcast_param_no_conflict() {
        let env = TypeEnvironment::default();
        let name = unique_broadcast_param(&env);
        assert_eq!(name.as_str(), "__broadcast_element");
    }

    #[test]
    fn test_unique_broadcast_param_with_conflict() {
        let mut env = TypeEnvironment::default();
        env.bind_base(
            SimpleIdentifier::new("__broadcast_element").into(),
            crate::types::INT,
        )
        .unwrap();

        let name = unique_broadcast_param(&env);
        assert_eq!(name.as_str(), "__broadcast_element0");
    }
}
