//! Transform function for applying lambdas to arrays.

use std::any::TypeId;
use std::sync::Arc;

use anyhow::bail;

use crate::func::def::{FunctionDef, ParameterBinding, ParameterBindingProvider, Parameters};
use crate::tree::ast::expression::{Expression, ExpressionKind};
use crate::tree::typed_ast::context::ExpressionTranslationContext;
use crate::tree::typed_ast::expression::TypedExpression;
use crate::types::array::Array;
use crate::types::matcher::AnyMatcher;
use crate::types::Type;

/// transform(array, lambda) - Apply a lambda to each element of an array.
///
/// The lambda receives each array element and returns a new value.
/// The result is an array of the lambda's return values.
///
/// Example: `transform([1, 2, 3], x -> x * 2)` returns `[2, 4, 6]`
#[derive(Default, Clone)]
pub struct Transform;

impl FunctionDef for Transform {
    fn name(&self) -> &'static str {
        "transform"
    }

    fn parameters(&self) -> Parameters {
        // AnyMatcher used for arity checking only - custom_bind validates types
        Parameters::new()
            .with("array", AnyMatcher)
            .with("lambda", AnyMatcher)
    }

    fn custom_bind(
        &self,
        ast_binding: &ParameterBinding<Arc<Expression>>,
        ctx: &mut ExpressionTranslationContext,
    ) -> anyhow::Result<Option<ParameterBinding<Arc<TypedExpression>>>> {
        // 1. Type the array arg first
        let array_ast = ast_binding.get_by_name("array")?.clone();
        let typed_array = Arc::new(TypedExpression::from_ast_with_context(array_ast, ctx));

        // 2. Validate it's an array - if not, return error
        let element_type = match typed_array.resolved_type.as_ref() {
            Type::Array(arr) => arr.element_type.clone(),
            other => bail!("Expected array type, got {}", other),
        };

        // 3. Get lambda AST
        let lambda_ast = ast_binding.get_by_name("lambda")?.clone();

        let ExpressionKind::Lambda(lambda) = &lambda_ast.kind else {
            bail!("Expected lambda expression");
        };

        // 4. Apply inferred types and type-check the lambda
        let lambda_with_hints = lambda.with_param_types(&[element_type])?;

        let lambda_expr = Expression {
            kind: lambda_with_hints.into(),
            span: lambda_ast.span.clone(),
        };
        let typed_lambda = Arc::new(TypedExpression::from_ast_with_context(
            Arc::new(lambda_expr),
            ctx,
        ));

        // 5. Build final binding
        Ok(Some(ParameterBinding::from_named([
            ("array".to_string(), typed_array),
            ("lambda".to_string(), typed_lambda),
        ])))
    }

    fn return_type(&self, bindings: &dyn ParameterBindingProvider) -> anyhow::Result<Type> {
        // Get the lambda's return type from the resolved binding
        let lambda_arg = bindings.get_by_name("lambda")?;

        let return_type = match lambda_arg.typ() {
            Type::Function(f) => f.return_type.as_ref().clone(),
            // If still unknown, the result is array<unknown>
            Type::Unknown => Type::Unknown,
            other => bail!("Expected function type for lambda, got {}", other),
        };

        Ok(Array::new(return_type).into())
    }

    fn type_id(&self) -> TypeId {
        TypeId::of::<Self>()
    }
}
