//! Translation implementations for map functions

use hamelin_lib::sql::expression::apply::{BinaryOperatorApply, FunctionCallApply, Lambda};
use hamelin_lib::sql::expression::identifier::SimpleIdentifier;
use hamelin_lib::sql::expression::literal::{ArrayLiteral, ColumnReference, StringLiteral};
use hamelin_lib::sql::expression::operator::Operator;
use hamelin_lib::sql::expression::{Case, Leaf, SQLExpression};

use crate::utils::direct_function_translation;
use crate::TranslationRegistry;
use hamelin_lib::func::defs::{
    MapEmpty, MapFromArrays, MapFromKeyValue, MapFromPairs, MapKeys, MapValues,
    MapVariantGetValues, MapVariantToJsonValues, TransformValues,
};

/// Register all map function translations.
pub fn register(registry: &mut TranslationRegistry) {
    // map(keys_array, values_array) ->
    //   CASE WHEN cardinality(keys) = cardinality(values) THEN map(keys, values) END
    //
    // Trino's map() throws on mismatched lengths; this guard returns NULL instead.
    registry.register::<MapFromArrays>(|_, mut bindings| {
        let keys = bindings.take()?;
        let values = bindings.take()?;
        let keys_len: SQLExpression =
            FunctionCallApply::with_one("cardinality", keys.sql.clone()).into();
        let values_len: SQLExpression =
            FunctionCallApply::with_one("cardinality", values.sql.clone()).into();
        let same_len: SQLExpression =
            BinaryOperatorApply::new(Operator::Eq, keys_len, values_len).into();
        let map_expr: SQLExpression =
            FunctionCallApply::with_two("map", keys.sql, values.sql).into();
        Ok(Case::new(&[(same_len, map_expr)], None).into())
    });

    // map(array_of_pairs) - SQL map_from_entries()
    registry.register::<MapFromPairs>(|_, mut bindings| {
        let elements = bindings.take()?;
        Ok(FunctionCallApply::with_one("map_from_entries", elements.sql).into())
    });

    // map() - empty map
    registry.register::<MapEmpty>(direct_function_translation);

    // map(pair1, pair2, ...) - SQL map_from_entries(ARRAY[...])
    // Duplicate key detection is left to the database at runtime
    registry.register::<MapFromKeyValue>(|_, mut bindings| {
        let mut pairs = vec![];
        while let Ok(pair) = bindings.take() {
            pairs.push(pair.sql);
        }
        Ok(FunctionCallApply::with_one("map_from_entries", ArrayLiteral::new(pairs).into()).into())
    });

    // map_keys(map) - pass through
    registry.register::<MapKeys>(direct_function_translation);

    // map_values(map) - pass through
    registry.register::<MapValues>(direct_function_translation);

    // transform_values(map, v -> body) -> transform_values(map, (k, v) -> body)
    //
    // Hamelin's `transform_values` lambda is 1-arg (value only). Trino's
    // `transform_values` requires a 2-arg lambda `(k, v)`. Wrap the 1-arg
    // lambda's body with a synthetic `k` parameter that the body ignores.
    registry.register::<TransformValues>(|_, mut bindings| {
        let map_expr = bindings.take()?;
        let lambda_expr = bindings.take()?;

        let wrapped_lambda: SQLExpression = match lambda_expr.sql {
            SQLExpression::Lambda(lambda) => {
                let mut args = vec![SimpleIdentifier::new("__k")];
                args.extend(lambda.arguments);
                Lambda::new(args, *lambda.body).into()
            }
            other => anyhow::bail!(
                "transform_values: expected lambda argument, got {:?}",
                other
            ),
        };

        Ok(FunctionCallApply::with_two("transform_values", map_expr.sql, wrapped_lambda).into())
    });

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

        let key_param = SimpleIdentifier::new("k");
        let value_param = SimpleIdentifier::new("v");
        let value_ref: SQLExpression = ColumnReference::new(value_param.clone().into()).into();
        let body: SQLExpression = FunctionCallApply::with_two(
            "json_extract",
            value_ref,
            StringLiteral::new(&json_path).into(),
        )
        .into();
        let lambda: SQLExpression = Lambda::new(vec![key_param, value_param], body).into();
        Ok(FunctionCallApply::with_two("transform_values", map_expr.sql, lambda).into())
    });

    // map_variant_to_json_values(m) ->
    //   transform_values(m, (k, v) -> json_format(v))
    registry.register::<MapVariantToJsonValues>(|_, mut bindings| {
        let map_expr = bindings.take()?;
        let key_param = SimpleIdentifier::new("k");
        let value_param = SimpleIdentifier::new("v");
        let value_ref: SQLExpression = ColumnReference::new(value_param.clone().into()).into();
        let body: SQLExpression = FunctionCallApply::with_one("json_format", value_ref).into();
        let lambda: SQLExpression = Lambda::new(vec![key_param, value_param], body).into();
        Ok(FunctionCallApply::with_two("transform_values", map_expr.sql, lambda).into())
    });
}
