//! Tests for transform function and lambda evaluation

use crate::eval::environment::Environment;
use crate::eval::evaluator::eval;
use crate::value::Value;
use hamelin_lib::tree::ast::identifier::SimpleIdentifier;
use hamelin_lib::tree::builder::{array, call, field_ref, lambda1, multiply, ExpressionBuilder};
use hamelin_lib::tree::options::ExpressionTypeCheckOptions;
use hamelin_lib::tree::typed_ast::environment::TypeEnvironment;
use hamelin_lib::type_check_expression;
use hamelin_lib::types::INT;
use std::sync::Arc;

use super::test_helpers::TestContext;

#[test]
fn test_transform_multiply_by_constant() {
    // transform([1, 2, 3], x -> x * 2) = [2, 4, 6]
    let ctx = TestContext::default();

    let expr = call("transform")
        .arg(array().element(1).element(2).element(3))
        .arg(lambda1("x").body(multiply(field_ref("x"), 2)));

    let result = ctx.eval_expr(&expr);
    assert_eq!(
        result,
        Value::Array(vec![Value::Int(2), Value::Int(4), Value::Int(6)])
    );
}

#[test]
fn test_transform_add_constant() {
    // transform([10, 20, 30], x -> x + 5) = [15, 25, 35]
    use hamelin_lib::tree::builder::add;

    let ctx = TestContext::default();

    let expr = call("transform")
        .arg(array().element(10).element(20).element(30))
        .arg(lambda1("x").body(add(field_ref("x"), 5)));

    let result = ctx.eval_expr(&expr);
    assert_eq!(
        result,
        Value::Array(vec![Value::Int(15), Value::Int(25), Value::Int(35)])
    );
}

#[test]
fn test_transform_with_captured_variable() {
    // SET multiplier = 10
    // transform([1, 2, 3], x -> x * multiplier) = [10, 20, 30]
    let mut env = Environment::new();
    env.bind(SimpleIdentifier::new("multiplier"), Value::Int(10));

    let mut trans_env = TypeEnvironment::default();
    trans_env.bind_str("multiplier", INT);

    let expr = type_check_expression(
        call("transform")
            .arg(array().element(1).element(2).element(3))
            .arg(lambda1("x").body(multiply(field_ref("x"), field_ref("multiplier"))))
            .build(),
        ExpressionTypeCheckOptions::builder()
            .bindings(Arc::new(trans_env))
            .build(),
    )
    .output;

    let result = eval(&expr, &env).unwrap();
    assert_eq!(
        result,
        Value::Array(vec![Value::Int(10), Value::Int(20), Value::Int(30)])
    );
}

#[test]
fn test_transform_empty_array() {
    // transform([], x -> x * 2) = []
    let ctx = TestContext::default();

    let expr = call("transform")
        .arg(array())
        .arg(lambda1("x").body(multiply(field_ref("x"), 2)));

    let result = ctx.eval_expr(&expr);
    assert_eq!(result, Value::Array(vec![]));
}

#[test]
fn test_transform_null_array() {
    // When the array variable is null, transform returns null
    use hamelin_lib::types::array::Array;

    let mut env = Environment::new();
    env.bind(SimpleIdentifier::new("arr"), Value::Null);

    let mut trans_env = TypeEnvironment::default();
    trans_env.bind_str("arr", Array::new(INT).into());

    let expr = type_check_expression(
        call("transform")
            .arg(field_ref("arr"))
            .arg(lambda1("x").body(multiply(field_ref("x"), 2)))
            .build(),
        ExpressionTypeCheckOptions::builder()
            .bindings(Arc::new(trans_env))
            .build(),
    )
    .output;

    let result = eval(&expr, &env).unwrap();
    assert_eq!(result, Value::Null);
}

#[test]
fn test_transform_with_null_elements() {
    // transform([1, null, 3], x -> x * 2) = [2, null, 6]
    // Null propagates through the lambda
    use hamelin_lib::tree::builder::null;

    let ctx = TestContext::default();

    let arr = array().element(1).element(null()).element(3);
    let expr = call("transform")
        .arg(arr)
        .arg(lambda1("x").body(multiply(field_ref("x"), 2)));

    let result = ctx.eval_expr(&expr);
    assert_eq!(
        result,
        Value::Array(vec![Value::Int(2), Value::Null, Value::Int(6)])
    );
}

#[test]
fn test_transform_nested() {
    // Nested transform: transform(transform([1, 2], x -> x * 2), y -> y + 1)
    // = transform([2, 4], y -> y + 1) = [3, 5]
    use hamelin_lib::tree::builder::add;

    let ctx = TestContext::default();

    let inner = call("transform")
        .arg(array().element(1).element(2))
        .arg(lambda1("x").body(multiply(field_ref("x"), 2)));

    let expr = call("transform")
        .arg(inner)
        .arg(lambda1("y").body(add(field_ref("y"), 1)));

    let result = ctx.eval_expr(&expr);
    assert_eq!(result, Value::Array(vec![Value::Int(3), Value::Int(5)]));
}

#[test]
fn test_broadcast_multiply_via_coercion() {
    // [1, 2, 3] * 10 should be rewritten to transform and evaluate to [10, 20, 30]
    let ctx = TestContext::default();

    let expr = multiply(array().element(1).element(2).element(3), 10);

    let result = ctx.eval_expr(&expr);
    assert_eq!(
        result,
        Value::Array(vec![Value::Int(10), Value::Int(20), Value::Int(30)])
    );
}

#[test]
fn test_broadcast_add_via_coercion() {
    // [1, 2, 3] + 10 should be rewritten to transform and evaluate to [11, 12, 13]
    use hamelin_lib::tree::builder::add;

    let ctx = TestContext::default();

    let expr = add(array().element(1).element(2).element(3), 10);

    let result = ctx.eval_expr(&expr);
    assert_eq!(
        result,
        Value::Array(vec![Value::Int(11), Value::Int(12), Value::Int(13)])
    );
}

#[test]
fn test_broadcast_comparison() {
    // [1, 5, 10] > 3 should broadcast to [false, true, true]
    use hamelin_lib::tree::builder::gt;

    let ctx = TestContext::default();

    let expr = gt(array().element(1).element(5).element(10), 3);

    let result = ctx.eval_expr(&expr);
    assert_eq!(
        result,
        Value::Array(vec![
            Value::Boolean(false),
            Value::Boolean(true),
            Value::Boolean(true)
        ])
    );
}

#[test]
fn test_broadcast_null_array() {
    // When the array is null, broadcast returns null (like transform)
    use hamelin_lib::types::array::Array;

    let mut env = Environment::new();
    env.bind(SimpleIdentifier::new("arr"), Value::Null);

    let mut trans_env = TypeEnvironment::default();
    trans_env.bind_str("arr", Array::new(INT).into());

    let expr = type_check_expression(
        multiply(field_ref("arr"), 10).build(),
        ExpressionTypeCheckOptions::builder()
            .bindings(Arc::new(trans_env))
            .build(),
    )
    .output;

    let result = eval(&expr, &env).unwrap();
    assert_eq!(result, Value::Null);
}
