lix 0.12.2

use std::collections::{BTreeMap, BTreeSet};

use datafusion::arrow::datatypes::{Field, Schema};
use datafusion::common::metadata::FieldMetadata;
use datafusion::common::tree_node::{Transformed, TreeNode};
use datafusion::common::{DFSchema, DataFusionError, ScalarValue};
use datafusion::logical_expr::expr::{Between, InList};
use datafusion::logical_expr::{BinaryExpr, Expr, Like, Operator};

use crate::LixError;

use super::error::lix_error_to_datafusion_error;
use super::result_metadata::{LIX_VALUE_TYPE_JSONB, LIX_VALUE_TYPE_METADATA_KEY, field_is_json};

pub(crate) fn validate_json_predicate_filters(
    schema: &Schema,
    filters: &[Expr],
) -> Result<(), DataFusionError> {
    for filter in filters {
        validate_json_predicate_expr_with_arrow_schema(schema, filter)
            .map_err(lix_error_to_datafusion_error)?;
    }
    Ok(())
}

pub(crate) fn canonicalize_json_identity_text_filters(
    schema: &Schema,
    filters: &[Expr],
) -> Result<Vec<Expr>, DataFusionError> {
    filters
        .iter()
        .cloned()
        .map(|filter| canonicalize_json_identity_text_filter(schema, filter))
        .collect()
}

pub(crate) fn validate_json_predicate_expr_with_dfschema(
    schema: &DFSchema,
    expr: &Expr,
) -> Result<(), LixError> {
    validate_expr(expr, &|column| {
        schema
            .field_with_name(column.relation.as_ref(), &column.name)
            .ok()
            .map(AsRef::as_ref)
    })
}

pub(crate) fn json_predicate_placeholder_indexes_with_dfschema(
    schema: &DFSchema,
    expr: &Expr,
) -> BTreeSet<usize> {
    let mut indexes = BTreeSet::new();
    collect_json_predicate_placeholder_indexes(expr, &mut indexes, &|column| {
        schema
            .field_with_name(column.relation.as_ref(), &column.name)
            .ok()
            .map(AsRef::as_ref)
    });
    indexes
}

fn canonicalize_json_identity_text_filter(
    schema: &Schema,
    expr: Expr,
) -> Result<Expr, DataFusionError> {
    expr.transform(|expr| canonicalize_json_identity_text_expr(schema, expr))
        .map(|transformed| transformed.data)
}

fn canonicalize_json_identity_text_expr(
    schema: &Schema,
    expr: Expr,
) -> Result<Transformed<Expr>, DataFusionError> {
    match expr {
        Expr::BinaryExpr(binary) if is_comparison_operator(binary.op) => {
            canonicalize_json_identity_text_binary(schema, binary)
        }
        Expr::InList(in_list) => canonicalize_json_identity_text_in_list(schema, in_list),
        _ => Ok(Transformed::no(expr)),
    }
}

fn canonicalize_json_identity_text_binary(
    schema: &Schema,
    binary: BinaryExpr,
) -> Result<Transformed<Expr>, DataFusionError> {
    let BinaryExpr { left, op, right } = binary;
    let left_identity_json = is_identity_json_expr_for_arrow_schema(schema, &left);
    let right_identity_json = is_identity_json_expr_for_arrow_schema(schema, &right);
    let left = if right_identity_json {
        Box::new(canonicalize_json_text_literal(*left)?)
    } else {
        left
    };
    let right = if left_identity_json {
        Box::new(canonicalize_json_text_literal(*right)?)
    } else {
        right
    };
    Ok(Transformed::yes(Expr::BinaryExpr(BinaryExpr::new(
        left, op, right,
    ))))
}

fn canonicalize_json_identity_text_in_list(
    schema: &Schema,
    in_list: InList,
) -> Result<Transformed<Expr>, DataFusionError> {
    let expr_identity_json = is_identity_json_expr_for_arrow_schema(schema, &in_list.expr);
    let list_has_identity_json = in_list
        .list
        .iter()
        .any(|item| is_identity_json_expr_for_arrow_schema(schema, item));

    let expr = if list_has_identity_json {
        Box::new(canonicalize_json_text_literal(*in_list.expr)?)
    } else {
        in_list.expr
    };
    let list = if expr_identity_json {
        in_list
            .list
            .into_iter()
            .map(canonicalize_json_text_literal)
            .collect::<Result<Vec<_>, _>>()?
    } else {
        in_list.list
    };

    Ok(Transformed::yes(Expr::InList(InList::new(
        expr,
        list,
        in_list.negated,
    ))))
}

fn canonicalize_json_text_literal(expr: Expr) -> Result<Expr, DataFusionError> {
    let Expr::Literal(literal, metadata) = expr else {
        return Ok(expr);
    };
    let canonical = match &literal {
        ScalarValue::Utf8(Some(value))
        | ScalarValue::Utf8View(Some(value))
        | ScalarValue::LargeUtf8(Some(value)) => Some(
            crate::sql2::udfs::common::canonical_jsonb_text(value).map_err(|error| {
                lix_error_to_datafusion_error(LixError::new(
                    LixError::CODE_TYPE_MISMATCH,
                    format!("JSON comparison value is not valid JSON: {error}"),
                ))
            })?,
        ),
        _ => None,
    };
    Ok(canonical.map_or_else(
        || Expr::Literal(literal, metadata),
        |value| Expr::Literal(ScalarValue::Utf8(Some(value)), Some(json_field_metadata())),
    ))
}

fn json_field_metadata() -> FieldMetadata {
    FieldMetadata::new(BTreeMap::from([(
        LIX_VALUE_TYPE_METADATA_KEY.to_string(),
        LIX_VALUE_TYPE_JSONB.to_string(),
    )]))
}

fn is_identity_json_expr_for_arrow_schema(schema: &Schema, expr: &Expr) -> bool {
    is_identity_json_expr(expr)
        && is_json_expr(expr, &|column| {
            schema
                .fields()
                .iter()
                .find(|field| field.name() == &column.name)
                .map(AsRef::as_ref)
        })
}

fn validate_json_predicate_expr_with_arrow_schema(
    schema: &Schema,
    expr: &Expr,
) -> Result<(), LixError> {
    validate_expr(expr, &|column| {
        schema
            .fields()
            .iter()
            .find(|field| field.name() == &column.name)
            .map(AsRef::as_ref)
    })
}

fn validate_expr<'a>(
    expr: &'a Expr,
    lookup_field: &impl Fn(&datafusion::common::Column) -> Option<&'a Field>,
) -> Result<(), LixError> {
    match expr {
        Expr::BinaryExpr(binary) => validate_binary_expr(binary, lookup_field),
        Expr::InList(in_list) => validate_in_list(in_list, lookup_field),
        Expr::Between(between) => validate_between(between, lookup_field),
        Expr::Like(like) | Expr::SimilarTo(like) => validate_like(like, lookup_field),
        Expr::Alias(alias) => validate_expr(&alias.expr, lookup_field),
        Expr::Not(inner)
        | Expr::IsNotNull(inner)
        | Expr::IsNull(inner)
        | Expr::IsTrue(inner)
        | Expr::IsFalse(inner)
        | Expr::IsUnknown(inner)
        | Expr::IsNotTrue(inner)
        | Expr::IsNotFalse(inner)
        | Expr::IsNotUnknown(inner)
        | Expr::Negative(inner) => validate_expr(inner, lookup_field),
        Expr::Cast(cast) => validate_expr(&cast.expr, lookup_field),
        Expr::TryCast(cast) => validate_expr(&cast.expr, lookup_field),
        Expr::ScalarFunction(function) => {
            for arg in &function.args {
                validate_expr(arg, lookup_field)?;
            }
            Ok(())
        }
        Expr::Case(case) => {
            if let Some(expr) = &case.expr {
                validate_expr(expr, lookup_field)?;
            }
            for (when, then) in &case.when_then_expr {
                validate_expr(when, lookup_field)?;
                validate_expr(then, lookup_field)?;
            }
            if let Some(expr) = &case.else_expr {
                validate_expr(expr, lookup_field)?;
            }
            Ok(())
        }
        Expr::AggregateFunction(function) => {
            for arg in &function.params.args {
                validate_expr(arg, lookup_field)?;
            }
            Ok(())
        }
        Expr::WindowFunction(function) => {
            for arg in &function.params.args {
                validate_expr(arg, lookup_field)?;
            }
            Ok(())
        }
        _ => Ok(()),
    }
}

fn collect_json_predicate_placeholder_indexes<'a>(
    expr: &'a Expr,
    indexes: &mut BTreeSet<usize>,
    lookup_field: &impl Fn(&datafusion::common::Column) -> Option<&'a Field>,
) {
    match expr {
        Expr::BinaryExpr(binary) if is_comparison_operator(binary.op) => {
            collect_json_predicate_placeholder_indexes(&binary.left, indexes, lookup_field);
            collect_json_predicate_placeholder_indexes(&binary.right, indexes, lookup_field);
            if is_json_expr(&binary.left, lookup_field) {
                collect_placeholder_indexes(&binary.right, indexes);
            }
            if is_json_expr(&binary.right, lookup_field) {
                collect_placeholder_indexes(&binary.left, indexes);
            }
        }
        Expr::BinaryExpr(binary) => {
            collect_json_predicate_placeholder_indexes(&binary.left, indexes, lookup_field);
            collect_json_predicate_placeholder_indexes(&binary.right, indexes, lookup_field);
        }
        Expr::InList(in_list) => {
            collect_json_predicate_placeholder_indexes(&in_list.expr, indexes, lookup_field);
            for item in &in_list.list {
                collect_json_predicate_placeholder_indexes(item, indexes, lookup_field);
            }
            if is_json_expr(&in_list.expr, lookup_field) {
                for item in &in_list.list {
                    collect_placeholder_indexes(item, indexes);
                }
            }
            for item in &in_list.list {
                if is_json_expr(item, lookup_field) {
                    collect_placeholder_indexes(&in_list.expr, indexes);
                }
            }
        }
        Expr::Between(between) => {
            collect_json_predicate_placeholder_indexes(&between.expr, indexes, lookup_field);
            collect_json_predicate_placeholder_indexes(&between.low, indexes, lookup_field);
            collect_json_predicate_placeholder_indexes(&between.high, indexes, lookup_field);
            if is_json_expr(&between.expr, lookup_field) {
                collect_placeholder_indexes(&between.low, indexes);
                collect_placeholder_indexes(&between.high, indexes);
            }
        }
        Expr::Alias(alias) => {
            collect_json_predicate_placeholder_indexes(&alias.expr, indexes, lookup_field);
        }
        Expr::Not(inner)
        | Expr::IsNotNull(inner)
        | Expr::IsNull(inner)
        | Expr::IsTrue(inner)
        | Expr::IsFalse(inner)
        | Expr::IsUnknown(inner)
        | Expr::IsNotTrue(inner)
        | Expr::IsNotFalse(inner)
        | Expr::IsNotUnknown(inner)
        | Expr::Negative(inner) => {
            collect_json_predicate_placeholder_indexes(inner, indexes, lookup_field);
        }
        Expr::Cast(cast) => {
            collect_json_predicate_placeholder_indexes(&cast.expr, indexes, lookup_field);
        }
        Expr::TryCast(cast) => {
            collect_json_predicate_placeholder_indexes(&cast.expr, indexes, lookup_field);
        }
        Expr::ScalarFunction(function) => {
            for arg in &function.args {
                collect_json_predicate_placeholder_indexes(arg, indexes, lookup_field);
            }
        }
        Expr::Case(case) => {
            if let Some(expr) = &case.expr {
                collect_json_predicate_placeholder_indexes(expr, indexes, lookup_field);
            }
            for (when, then) in &case.when_then_expr {
                collect_json_predicate_placeholder_indexes(when, indexes, lookup_field);
                collect_json_predicate_placeholder_indexes(then, indexes, lookup_field);
            }
            if let Some(expr) = &case.else_expr {
                collect_json_predicate_placeholder_indexes(expr, indexes, lookup_field);
            }
        }
        _ => {}
    }
}

fn collect_placeholder_indexes(expr: &Expr, indexes: &mut BTreeSet<usize>) {
    if let Expr::Placeholder(placeholder) = expr {
        if let Some(index) = placeholder
            .id
            .strip_prefix('$')
            .and_then(|value| value.parse::<usize>().ok())
        {
            indexes.insert(index);
        }
    }
}

fn validate_binary_expr<'a>(
    binary: &'a BinaryExpr,
    lookup_field: &impl Fn(&datafusion::common::Column) -> Option<&'a Field>,
) -> Result<(), LixError> {
    validate_expr(&binary.left, lookup_field)?;
    validate_expr(&binary.right, lookup_field)?;

    if !is_comparison_operator(binary.op) {
        return Ok(());
    }

    validate_comparison_operands(&binary.left, &binary.right, lookup_field)
}

fn validate_in_list<'a>(
    in_list: &'a InList,
    lookup_field: &impl Fn(&datafusion::common::Column) -> Option<&'a Field>,
) -> Result<(), LixError> {
    validate_expr(&in_list.expr, lookup_field)?;
    for item in &in_list.list {
        validate_expr(item, lookup_field)?;
    }

    if is_json_expr(&in_list.expr, lookup_field) {
        for item in &in_list.list {
            require_json_comparison_operand(item, lookup_field)?;
        }
    }

    for item in &in_list.list {
        if is_json_expr(item, lookup_field) {
            require_json_comparison_operand(&in_list.expr, lookup_field)?;
        }
    }

    Ok(())
}

fn validate_between<'a>(
    between: &'a Between,
    lookup_field: &impl Fn(&datafusion::common::Column) -> Option<&'a Field>,
) -> Result<(), LixError> {
    validate_expr(&between.expr, lookup_field)?;
    validate_expr(&between.low, lookup_field)?;
    validate_expr(&between.high, lookup_field)?;

    if is_json_expr(&between.expr, lookup_field) {
        require_json_comparison_operand(&between.low, lookup_field)?;
        require_json_comparison_operand(&between.high, lookup_field)?;
    }

    Ok(())
}

fn validate_like<'a>(
    like: &'a Like,
    lookup_field: &impl Fn(&datafusion::common::Column) -> Option<&'a Field>,
) -> Result<(), LixError> {
    validate_expr(&like.expr, lookup_field)?;
    validate_expr(&like.pattern, lookup_field)?;

    if is_json_expr(&like.expr, lookup_field) {
        return Err(json_predicate_type_error(&like.expr));
    }

    Ok(())
}

fn validate_comparison_operands<'a>(
    left: &'a Expr,
    right: &'a Expr,
    lookup_field: &impl Fn(&datafusion::common::Column) -> Option<&'a Field>,
) -> Result<(), LixError> {
    let left_is_json = is_json_expr(left, lookup_field);
    let right_is_json = is_json_expr(right, lookup_field);

    if left_is_json {
        require_json_comparison_operand(right, lookup_field)?;
    }
    if right_is_json {
        require_json_comparison_operand(left, lookup_field)?;
    }

    Ok(())
}

fn require_json_comparison_operand<'a>(
    expr: &'a Expr,
    lookup_field: &impl Fn(&datafusion::common::Column) -> Option<&'a Field>,
) -> Result<(), LixError> {
    if is_json_expr(expr, lookup_field)
        || is_null_literal(expr)
        || matches!(expr, Expr::Placeholder(_))
    {
        return Ok(());
    }

    Err(json_predicate_type_error(expr))
}

fn is_json_expr<'a>(
    expr: &'a Expr,
    lookup_field: &impl Fn(&datafusion::common::Column) -> Option<&'a Field>,
) -> bool {
    match expr {
        Expr::Column(column) => lookup_field(column).is_some_and(field_is_json),
        Expr::Literal(_, Some(metadata)) => metadata
            .inner()
            .get(LIX_VALUE_TYPE_METADATA_KEY)
            .is_some_and(|value| value == LIX_VALUE_TYPE_JSONB),
        Expr::ScalarFunction(function) => matches!(
            function.name(),
            "__lix_json_get" | "__lix_json_path_get" | "__lix_jsonb"
        ),
        Expr::Alias(alias) => is_json_expr(&alias.expr, lookup_field),
        Expr::Cast(cast) => is_json_expr(&cast.expr, lookup_field),
        Expr::TryCast(cast) => is_json_expr(&cast.expr, lookup_field),
        _ => false,
    }
}

fn is_identity_json_expr(expr: &Expr) -> bool {
    match expr {
        Expr::Column(column) => matches!(column.name.as_str(), "row_pk" | "lixcol_row_pk"),
        Expr::Alias(alias) => is_identity_json_expr(&alias.expr),
        Expr::Cast(cast) => is_identity_json_expr(&cast.expr),
        Expr::TryCast(cast) => is_identity_json_expr(&cast.expr),
        _ => false,
    }
}

fn is_null_literal(expr: &Expr) -> bool {
    matches!(expr, Expr::Literal(value, _) if matches!(value, ScalarValue::Null))
}

fn is_comparison_operator(op: Operator) -> bool {
    matches!(
        op,
        Operator::Eq
            | Operator::NotEq
            | Operator::Lt
            | Operator::LtEq
            | Operator::Gt
            | Operator::GtEq
            | Operator::IsDistinctFrom
            | Operator::IsNotDistinctFrom
    )
}

fn json_predicate_type_error(expr: &Expr) -> LixError {
    LixError::new(
        LixError::CODE_TYPE_MISMATCH,
        format!("JSON columns can only be compared with JSON expressions, got {expr}"),
    )
    .with_hint("Cast JSON text with ::jsonb, use PostgreSQL -> or ->> for JSON access, or use IS NULL for null checks.")
}