//! Translation implementations for array operators and functions

use crate::utils::{
    array_negative_index_to_positive, convert_nested_transforms, direct_function_translation,
    hamelin_array_index_to_sql_with_negatives, to_sql_concat,
};
use crate::TranslationRegistry;
use hamelin_lib::func::defs::{
    ArrayAll, ArrayAny, ArrayAvg, ArrayConcat, ArrayDistinct, ArrayJoin2, ArrayJoin3, ArrayMax,
    ArrayMin, ArrayOrMapLen, ArraySum, ArrayVariantGet, ArrayVariantToJson, Filter, FilterNull,
    Flatten, Sequence2, Sequence3, Slice, Split, Transform, Zip,
};
use hamelin_lib::sql::expression::apply::{
    BinaryOperatorApply, FunctionCallApply, Lambda, UnaryOperatorApply,
};
use hamelin_lib::sql::expression::identifier::SimpleIdentifier;
use hamelin_lib::sql::expression::literal::{
    BooleanLiteral, ColumnReference, IntegerLiteral, NullLiteral, ScientificLiteral, StringLiteral,
};
use hamelin_lib::sql::expression::operator::Operator;
use hamelin_lib::sql::expression::{Case, Cast, Leaf, SQLExpression};
use hamelin_lib::sql::projection_builder::ProjectionBuilder;
use hamelin_lib::sql::types::{SQLAnonRowType, SQLArrayType, SQLBaseType, SQLDecimalType, SQLType};
use hamelin_lib::types::array::Array;
use hamelin_lib::types::{Type, UNKNOWN};

/// Register all array operator and function translations.
pub fn register(registry: &mut TranslationRegistry) {
    // ArrayConcat: array + array -> SQL CONCAT with struct widening if needed
    registry.register::<ArrayConcat>(|name, mut bindings| {
        let left_typ = bindings.get_by_index(0)?.typ.clone();
        let right_typ = bindings.get_by_index(1)?.typ.clone();

        // If types are equal, no transformation needed
        if left_typ == right_typ {
            return to_sql_concat(name, bindings);
        }

        // Compute the merged type (this is also done during type checking)
        let merge_type = left_typ.clone().merge(right_typ.clone())?;

        // Helper to transform an operand to the target merged type
        let transform_to_target =
            |sql: hamelin_lib::sql::expression::SQLExpression,
             from_type: &Type,
             target_type: &Type|
             -> anyhow::Result<hamelin_lib::sql::expression::SQLExpression> {
                // If types match, no transformation needed
                if from_type == target_type {
                    return Ok(sql);
                }

                // Check if from_type is array(unknown) - needs simple cast to target
                if let Type::Array(Array { element_type }) = from_type {
                    if **element_type == UNKNOWN {
                        // Cast empty array to the target type
                        return Ok(Cast::new(sql, target_type.clone().to_sql()?).into());
                    }
                }

                // Both are struct arrays with different fields - need struct widening
                ProjectionBuilder::expand_repeated_struct(sql, from_type, target_type)
                    .map_err(|e| e.into())
            };

        let left = bindings.take()?;
        let right = bindings.take()?;

        let left_transformed = transform_to_target(left.sql, &left_typ, &merge_type)?;
        let right_transformed = transform_to_target(right.sql, &right_typ, &merge_type)?;

        Ok(BinaryOperatorApply::new(Operator::Concat, left_transformed, right_transformed).into())
    });

    // filter_null(array) -> filter(array, _e -> _e IS NOT NULL)
    registry.register::<FilterNull>(|_, mut bindings| {
        let x = bindings.take()?;
        Ok(FunctionCallApply::with_two(
            "filter",
            x.sql,
            Lambda::new(
                vec![SimpleIdentifier::new("_e")],
                BinaryOperatorApply::new(
                    Operator::IsNot,
                    ColumnReference::new(SimpleIdentifier::new("_e").into()).into(),
                    NullLiteral::default().into(),
                )
                .into(),
            )
            .into(),
        )
        .into())
    });

    // len(array|map) -> cardinality(x)
    registry.register::<ArrayOrMapLen>(|_, mut bindings| {
        let x = bindings.take()?;
        Ok(FunctionCallApply::with_one("cardinality", x.sql).into())
    });

    // array_distinct(array) - pass through
    registry.register::<ArrayDistinct>(direct_function_translation);

    // slice(array, start, end) - convert indices from Hamelin to SQL space
    registry.register::<Slice>(|_, mut bindings| {
        let array = bindings.take()?;
        let start = bindings.take()?;
        let end = bindings.take()?;

        // Convert negative indices to positive in Hamelin space
        let positive_end = array_negative_index_to_positive(end.sql.clone(), array.sql.clone());
        let positive_start = array_negative_index_to_positive(start.sql.clone(), array.sql.clone());

        // Calculate length as positive_end - positive_start (both in Hamelin zero-based space)
        let length = BinaryOperatorApply::new(Operator::Minus, positive_end, positive_start).into();

        // Convert start to SQL space: add 1 only for originally positive indices
        let sql_start = hamelin_array_index_to_sql_with_negatives(start.sql, array.sql.clone());

        Ok(FunctionCallApply::with_three("slice", array.sql, sql_start, length).into())
    });

    // split(string, delimiter) - pass through
    registry.register::<Split>(direct_function_translation);

    // array_join(array, delimiter) - pass through
    registry.register::<ArrayJoin2>(direct_function_translation);

    // array_join(array, delimiter, null_replacement) - pass through
    registry.register::<ArrayJoin3>(direct_function_translation);

    // flatten(array<array<T>>) - pass through
    registry.register::<Flatten>(direct_function_translation);

    // transform(array, lambda) - pass through
    registry.register::<Transform>(direct_function_translation);

    // filter(array, lambda) - pass through
    registry.register::<Filter>(direct_function_translation);

    // array_variant_get(arr, path) -> transform(arr, x -> json_extract(x, '$["path"]'))
    //
    // Internal operator emitted by the DataFusion transform-lowering pipeline
    // as a fast path. For Trino, which natively supports `transform` + JSON,
    // we translate back to the equivalent lambda form so the query stays
    // valid and runs efficiently.
    registry.register::<ArrayVariantGet>(|_, mut bindings| {
        let arr = bindings.take()?;
        let path = bindings.take()?;
        let path_string = match &path.sql {
            SQLExpression::Leaf(Leaf::StringLiteral(s)) => s.value.clone(),
            _ => anyhow::bail!("array_variant_get: path must be a string literal"),
        };
        let json_path = format!("$[\"{}\"]", path_string);

        let lambda_param = SimpleIdentifier::new("x");
        let param_ref: SQLExpression = ColumnReference::new(lambda_param.clone().into()).into();
        let body: SQLExpression = FunctionCallApply::with_two(
            "json_extract",
            param_ref,
            StringLiteral::new(&json_path).into(),
        )
        .into();
        let lambda: SQLExpression = Lambda::new(vec![lambda_param], body).into();
        Ok(FunctionCallApply::with_two("transform", arr.sql, lambda).into())
    });

    // array_variant_to_json(arr) -> transform(arr, x -> json_format(x))
    registry.register::<ArrayVariantToJson>(|_, mut bindings| {
        let arr = bindings.take()?;
        let lambda_param = SimpleIdentifier::new("x");
        let param_ref: SQLExpression = ColumnReference::new(lambda_param.clone().into()).into();
        let body: SQLExpression = FunctionCallApply::with_one("json_format", param_ref).into();
        let lambda: SQLExpression = Lambda::new(vec![lambda_param], body).into();
        Ok(FunctionCallApply::with_two("transform", arr.sql, lambda).into())
    });

    // any(array<boolean>) - contains(x, true)
    // Three-valued logic: true if any true, null if any null (and no true), false otherwise
    registry.register::<ArrayAny>(|_, mut bindings| {
        let x = bindings.take()?;

        match x.sql {
            hamelin_lib::sql::expression::SQLExpression::FunctionCallApply(fca)
                if fca.function_name == "transform" =>
            {
                Ok(convert_nested_transforms(fca.into(), "any_match"))
            }
            _ => Ok(FunctionCallApply::with_two(
                "contains",
                x.sql,
                BooleanLiteral::new(true).into(),
            )
            .into()),
        }
    });

    // all(array<boolean>) - NOT contains(x, false)
    // Three-valued logic: false if any false, null if any null (and no false), true otherwise
    registry.register::<ArrayAll>(|_, mut bindings| {
        let x = bindings.take()?;
        match x.sql {
            hamelin_lib::sql::expression::SQLExpression::FunctionCallApply(fca)
                if fca.function_name == "transform" =>
            {
                Ok(convert_nested_transforms(fca.into(), "all_match"))
            }
            _ => Ok(UnaryOperatorApply::new(
                Operator::Not,
                FunctionCallApply::with_two("contains", x.sql, BooleanLiteral::new(false).into())
                    .into(),
            )
            .into()),
        }
    });

    // max(array) - uses reduce to ignore nulls (Trino array_max returns null if any element is null)
    registry.register::<ArrayMax>(|_, mut bindings| {
        let x = bindings.take()?;
        let s = SimpleIdentifier::new("s");
        let n = SimpleIdentifier::new("n");

        let fca = FunctionCallApply::with_four(
            "reduce",
            x.sql,
            NullLiteral::default().into(),
            Lambda::new(
                vec![s.clone(), n.clone()],
                Case::new(
                    &[(
                        BinaryOperatorApply::new(
                            Operator::Is,
                            ColumnReference::new(s.clone().into()).into(),
                            NullLiteral::default().into(),
                        )
                        .into(),
                        ColumnReference::new(n.clone().into()).into(),
                    )],
                    Some(
                        Case::new(
                            &[(
                                BinaryOperatorApply::new(
                                    Operator::Gt,
                                    ColumnReference::new(n.clone().into()).into(),
                                    ColumnReference::new(s.clone().into()).into(),
                                )
                                .into(),
                                ColumnReference::new(n.clone().into()).into(),
                            )],
                            Some(ColumnReference::new(s.clone().into()).into()),
                        )
                        .into(),
                    ),
                )
                .into(),
            )
            .into(),
            Lambda::new(
                vec![s.clone()],
                ColumnReference::new(s.clone().into()).into(),
            )
            .into(),
        );

        Ok(fca.into())
    });

    // min(array) - uses reduce to ignore nulls (Trino array_min returns null if any element is null)
    registry.register::<ArrayMin>(|_, mut bindings| {
        let x = bindings.take()?;
        let s = SimpleIdentifier::new("s");
        let n = SimpleIdentifier::new("n");

        let fca = FunctionCallApply::with_four(
            "reduce",
            x.sql,
            NullLiteral::default().into(),
            Lambda::new(
                vec![s.clone(), n.clone()],
                Case::new(
                    &[(
                        BinaryOperatorApply::new(
                            Operator::Is,
                            ColumnReference::new(s.clone().into()).into(),
                            NullLiteral::default().into(),
                        )
                        .into(),
                        ColumnReference::new(n.clone().into()).into(),
                    )],
                    Some(
                        Case::new(
                            &[(
                                BinaryOperatorApply::new(
                                    Operator::Lt,
                                    ColumnReference::new(n.clone().into()).into(),
                                    ColumnReference::new(s.clone().into()).into(),
                                )
                                .into(),
                                ColumnReference::new(n.clone().into()).into(),
                            )],
                            Some(ColumnReference::new(s.clone().into()).into()),
                        )
                        .into(),
                    ),
                )
                .into(),
            )
            .into(),
            Lambda::new(
                vec![s.clone()],
                ColumnReference::new(s.clone().into()).into(),
            )
            .into(),
        );

        Ok(fca.into())
    });

    // sum(array<numeric>) - filter nulls then reduce to sum elements
    registry.register::<ArraySum>(|_, mut bindings| {
        let x = bindings.take()?;
        let s = SimpleIdentifier::new("s");
        let n = SimpleIdentifier::new("n");
        let (zero, accum_body) = reduce_parts(&x.typ, &s, &n);
        let filtered: SQLExpression = FunctionCallApply::with_two(
            "filter",
            x.sql,
            Lambda::new(
                vec![SimpleIdentifier::new("_e")],
                BinaryOperatorApply::new(
                    Operator::IsNot,
                    ColumnReference::new(SimpleIdentifier::new("_e").into()).into(),
                    NullLiteral::default().into(),
                )
                .into(),
            )
            .into(),
        )
        .into();
        let fca = FunctionCallApply::with_four(
            "reduce",
            filtered,
            zero,
            Lambda::new(vec![s.clone(), n], accum_body).into(),
            Lambda::new(
                vec![s.clone()],
                ColumnReference::new(s.clone().into()).into(),
            )
            .into(),
        );

        Ok(fca.into())
    });

    // avg(array<numeric>) - filter nulls, then sum / cardinality
    registry.register::<ArrayAvg>(|_, mut bindings| {
        let x = bindings.take()?;
        let s = SimpleIdentifier::new("s");
        let n = SimpleIdentifier::new("n");
        let (zero, accum_body) = reduce_parts(&x.typ, &s, &n);
        let filtered: SQLExpression = FunctionCallApply::with_two(
            "filter",
            x.sql,
            Lambda::new(
                vec![SimpleIdentifier::new("_e")],
                BinaryOperatorApply::new(
                    Operator::IsNot,
                    ColumnReference::new(SimpleIdentifier::new("_e").into()).into(),
                    NullLiteral::default().into(),
                )
                .into(),
            )
            .into(),
        )
        .into();

        // Calculate sum using reduce
        let sum_expr: SQLExpression = FunctionCallApply::with_four(
            "reduce",
            filtered.clone(),
            zero,
            Lambda::new(vec![s.clone(), n], accum_body).into(),
            Lambda::new(
                vec![s.clone()],
                ColumnReference::new(s.clone().into()).into(),
            )
            .into(),
        )
        .into();

        // For integer arrays, cast sum to double so we get float division
        let element_type = match &x.typ {
            Type::Array(a) => &*a.element_type,
            _ => anyhow::bail!("ArrayAvg expected array type, got {:?}", x.typ),
        };
        let dividend = if *element_type == Type::Int {
            Cast::new(sum_expr, SQLBaseType::Double.into()).into()
        } else {
            sum_expr
        };

        let card: SQLExpression = FunctionCallApply::with_one("cardinality", filtered).into();

        // Divide sum by filtered array length
        let avg_expr: SQLExpression =
            BinaryOperatorApply::new(Operator::Slash, dividend, card.clone()).into();

        // Guard against empty/all-null arrays: 0/0 gives NaN in Trino, we want NULL
        let guarded = Case::new(
            &[(
                BinaryOperatorApply::new(Operator::Eq, card, IntegerLiteral::new("0").into())
                    .into(),
                NullLiteral::default().into(),
            )],
            Some(avg_expr),
        );

        Ok(guarded.into())
    });

    // sequence(start, stop) -> sequence(start, stop)
    // Trino/Presto has native sequence function with same semantics (inclusive)
    registry.register::<Sequence2>(|_, mut bindings| {
        let start = bindings.take()?;
        let stop = bindings.take()?;
        Ok(FunctionCallApply::with_two("sequence", start.sql, stop.sql).into())
    });

    // sequence(start, stop, step) -> sequence(start, stop, step)
    registry.register::<Sequence3>(|_, mut bindings| {
        let start = bindings.take()?;
        let stop = bindings.take()?;
        let step = bindings.take()?;
        Ok(FunctionCallApply::with_three("sequence", start.sql, stop.sql, step.sql).into())
    });

    // zip(left, right) -> CAST(zip(left, right) AS ARRAY(ROW(T1, T2)))
    // Trino's native zip returns anonymous rows; the CAST matches Hamelin's tuple element types.
    registry.register::<Zip>(|_, mut bindings| {
        let left = bindings.take()?;
        let right = bindings.take()?;

        let left_elem = match &left.typ {
            Type::Array(a) => (*a.element_type).clone(),
            _ => anyhow::bail!("zip requires array arguments"),
        };
        let right_elem = match &right.typ {
            Type::Array(a) => (*a.element_type).clone(),
            _ => anyhow::bail!("zip requires array arguments"),
        };

        let row_type = SQLAnonRowType::new(vec![left_elem.to_sql()?, right_elem.to_sql()?]);
        let array_type = SQLArrayType::new(row_type.into());

        let zip_call = FunctionCallApply::with_two("zip", left.sql, right.sql);
        Ok(Cast::new(zip_call.into(), array_type.into()).into())
    });
}

/// Build the zero literal and accumulator lambda body for Trino's `reduce`.
///
/// Trino's `reduce` requires the initial value type to exactly match the accumulator.
/// For decimal types, `s + n` causes precision escalation that confuses the
/// SignatureBinder, so we wrap the addition in a CAST back to the element type.
fn reduce_parts(
    array_type: &Type,
    s: &SimpleIdentifier,
    n: &SimpleIdentifier,
) -> (SQLExpression, SQLExpression) {
    let sum_expr: SQLExpression = BinaryOperatorApply::new(
        Operator::Plus,
        ColumnReference::new(s.clone().into()).into(),
        ColumnReference::new(n.clone().into()).into(),
    )
    .into();

    let element_type = match array_type {
        Type::Array(a) => &*a.element_type,
        _ => return (IntegerLiteral::new("0").into(), sum_expr),
    };
    match element_type {
        Type::Double => (ScientificLiteral::new("0E0").into(), sum_expr),
        Type::Decimal(d) => {
            let sql_type: SQLType = SQLDecimalType::new(d.precision as u64, d.scale as u64).into();
            let zero = Cast::new(IntegerLiteral::new("0").into(), sql_type.clone()).into();
            let cast_sum = Cast::new(sum_expr, sql_type).into();
            (zero, cast_sum)
        }
        _ => (IntegerLiteral::new("0").into(), sum_expr),
    }
}
