// Copyright 2026 RisingWave Labs
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use chrono::Datelike;
use iceberg::expr::{Predicate as IcebergPredicate, Reference};
use iceberg::spec::Datum as IcebergDatum;
use risingwave_common::catalog::Field;
use risingwave_common::types::ScalarImpl;

use crate::expr::{Expr, ExprImpl, ExprType, Literal};
use crate::utils::Condition;

pub struct ExtractIcebergPredicateResult {
    pub iceberg_predicate: IcebergPredicate,
    pub extracted_condition: Condition,
    pub remaining_condition: Condition,
}

/// NOTE(kwannoel): We do predicate pushdown to the iceberg-sdk here.
/// zone-map is used to evaluate predicates on iceberg tables.
/// Without zone-map, iceberg-sdk will still apply the predicate on its own.
/// See: <https://github.com/apache/iceberg-rust/blob/5c1a9e68da346819072a15327080a498ad91c488/crates/iceberg/src/arrow/reader.rs#L229-L235>.
pub fn extract_iceberg_predicate(
    predicate: Condition,
    fields: &[Field],
) -> ExtractIcebergPredicateResult {
    if predicate.always_true() {
        return ExtractIcebergPredicateResult {
            iceberg_predicate: IcebergPredicate::AlwaysTrue,
            extracted_condition: Condition {
                conjunctions: vec![],
            },
            remaining_condition: Condition {
                conjunctions: vec![],
            },
        };
    }

    let mut iceberg_predicates = Vec::new();
    let mut extracted_conjunctions = Vec::new();
    let mut remaining_conjunctions = Vec::new();

    for conjunction in predicate.conjunctions {
        match rw_expr_to_iceberg_predicate(&conjunction, fields) {
            Some(iceberg_predicate) => {
                iceberg_predicates.push(iceberg_predicate);
                extracted_conjunctions.push(conjunction);
            }
            None => remaining_conjunctions.push(conjunction),
        }
    }

    let iceberg_predicate = iceberg_predicates
        .into_iter()
        .reduce(IcebergPredicate::and)
        .unwrap_or(IcebergPredicate::AlwaysTrue);

    ExtractIcebergPredicateResult {
        iceberg_predicate,
        extracted_condition: Condition {
            conjunctions: extracted_conjunctions,
        },
        remaining_condition: Condition {
            conjunctions: remaining_conjunctions,
        },
    }
}

fn rw_literal_to_iceberg_datum(literal: &Literal) -> Option<IcebergDatum> {
    let Some(scalar) = literal.get_data() else {
        return None;
    };
    match scalar {
        ScalarImpl::Bool(b) => Some(IcebergDatum::bool(*b)),
        ScalarImpl::Int32(i) => Some(IcebergDatum::int(*i)),
        ScalarImpl::Int64(i) => Some(IcebergDatum::long(*i)),
        ScalarImpl::Float32(f) => Some(IcebergDatum::float(*f)),
        ScalarImpl::Float64(f) => Some(IcebergDatum::double(*f)),
        ScalarImpl::Decimal(_) => {
            // TODO(iceberg): iceberg-rust doesn't support decimal predicate pushdown yet.
            None
        }
        ScalarImpl::Date(d) => {
            let Ok(datum) = IcebergDatum::date_from_ymd(d.0.year(), d.0.month(), d.0.day()) else {
                return None;
            };
            Some(datum)
        }
        ScalarImpl::Timestamp(t) => Some(IcebergDatum::timestamp_micros(
            t.0.and_utc().timestamp_micros(),
        )),
        ScalarImpl::Timestamptz(t) => Some(IcebergDatum::timestamptz_micros(t.timestamp_micros())),
        ScalarImpl::Utf8(s) => Some(IcebergDatum::string(s)),
        ScalarImpl::Bytea(b) => Some(IcebergDatum::binary(b.clone())),
        _ => None,
    }
}

fn rw_expr_to_iceberg_predicate(expr: &ExprImpl, fields: &[Field]) -> Option<IcebergPredicate> {
    match expr {
        ExprImpl::Literal(l) => match l.get_data() {
            Some(ScalarImpl::Bool(b)) => {
                if *b {
                    Some(IcebergPredicate::AlwaysTrue)
                } else {
                    Some(IcebergPredicate::AlwaysFalse)
                }
            }
            _ => None,
        },
        ExprImpl::FunctionCall(f) => {
            let args = f.inputs();
            match f.func_type() {
                ExprType::Not => {
                    let arg = rw_expr_to_iceberg_predicate(&args[0], fields)?;
                    Some(IcebergPredicate::negate(arg))
                }
                ExprType::And => {
                    let arg0 = rw_expr_to_iceberg_predicate(&args[0], fields)?;
                    let arg1 = rw_expr_to_iceberg_predicate(&args[1], fields)?;
                    Some(IcebergPredicate::and(arg0, arg1))
                }
                ExprType::Or => {
                    let arg0 = rw_expr_to_iceberg_predicate(&args[0], fields)?;
                    let arg1 = rw_expr_to_iceberg_predicate(&args[1], fields)?;
                    Some(IcebergPredicate::or(arg0, arg1))
                }
                ExprType::Equal if args[0].return_type() == args[1].return_type() => {
                    match [&args[0], &args[1]] {
                        [ExprImpl::InputRef(lhs), ExprImpl::Literal(rhs)]
                        | [ExprImpl::Literal(rhs), ExprImpl::InputRef(lhs)] => {
                            let column_name = &fields[lhs.index].name;
                            let reference = Reference::new(column_name);
                            let datum = rw_literal_to_iceberg_datum(rhs)?;
                            Some(reference.equal_to(datum))
                        }
                        _ => None,
                    }
                }
                ExprType::NotEqual if args[0].return_type() == args[1].return_type() => {
                    match [&args[0], &args[1]] {
                        [ExprImpl::InputRef(lhs), ExprImpl::Literal(rhs)]
                        | [ExprImpl::Literal(rhs), ExprImpl::InputRef(lhs)] => {
                            let column_name = &fields[lhs.index].name;
                            let reference = Reference::new(column_name);
                            let datum = rw_literal_to_iceberg_datum(rhs)?;
                            Some(reference.not_equal_to(datum))
                        }
                        _ => None,
                    }
                }
                ExprType::GreaterThan if args[0].return_type() == args[1].return_type() => {
                    match [&args[0], &args[1]] {
                        [ExprImpl::InputRef(lhs), ExprImpl::Literal(rhs)] => {
                            let column_name = &fields[lhs.index].name;
                            let reference = Reference::new(column_name);
                            let datum = rw_literal_to_iceberg_datum(rhs)?;
                            Some(reference.greater_than(datum))
                        }
                        [ExprImpl::Literal(rhs), ExprImpl::InputRef(lhs)] => {
                            let column_name = &fields[lhs.index].name;
                            let reference = Reference::new(column_name);
                            let datum = rw_literal_to_iceberg_datum(rhs)?;
                            Some(reference.less_than_or_equal_to(datum))
                        }
                        _ => None,
                    }
                }
                ExprType::GreaterThanOrEqual if args[0].return_type() == args[1].return_type() => {
                    match [&args[0], &args[1]] {
                        [ExprImpl::InputRef(lhs), ExprImpl::Literal(rhs)] => {
                            let column_name = &fields[lhs.index].name;
                            let reference = Reference::new(column_name);
                            let datum = rw_literal_to_iceberg_datum(rhs)?;
                            Some(reference.greater_than_or_equal_to(datum))
                        }
                        [ExprImpl::Literal(rhs), ExprImpl::InputRef(lhs)] => {
                            let column_name = &fields[lhs.index].name;
                            let reference = Reference::new(column_name);
                            let datum = rw_literal_to_iceberg_datum(rhs)?;
                            Some(reference.less_than(datum))
                        }
                        _ => None,
                    }
                }
                ExprType::LessThan if args[0].return_type() == args[1].return_type() => {
                    match [&args[0], &args[1]] {
                        [ExprImpl::InputRef(lhs), ExprImpl::Literal(rhs)] => {
                            let column_name = &fields[lhs.index].name;
                            let reference = Reference::new(column_name);
                            let datum = rw_literal_to_iceberg_datum(rhs)?;
                            Some(reference.less_than(datum))
                        }
                        [ExprImpl::Literal(rhs), ExprImpl::InputRef(lhs)] => {
                            let column_name = &fields[lhs.index].name;
                            let reference = Reference::new(column_name);
                            let datum = rw_literal_to_iceberg_datum(rhs)?;
                            Some(reference.greater_than_or_equal_to(datum))
                        }
                        _ => None,
                    }
                }
                ExprType::LessThanOrEqual if args[0].return_type() == args[1].return_type() => {
                    match [&args[0], &args[1]] {
                        [ExprImpl::InputRef(lhs), ExprImpl::Literal(rhs)] => {
                            let column_name = &fields[lhs.index].name;
                            let reference = Reference::new(column_name);
                            let datum = rw_literal_to_iceberg_datum(rhs)?;
                            Some(reference.less_than_or_equal_to(datum))
                        }
                        [ExprImpl::Literal(rhs), ExprImpl::InputRef(lhs)] => {
                            let column_name = &fields[lhs.index].name;
                            let reference = Reference::new(column_name);
                            let datum = rw_literal_to_iceberg_datum(rhs)?;
                            Some(reference.greater_than(datum))
                        }
                        _ => None,
                    }
                }
                ExprType::IsNull => match &args[0] {
                    ExprImpl::InputRef(lhs) => {
                        let column_name = &fields[lhs.index].name;
                        let reference = Reference::new(column_name);
                        Some(reference.is_null())
                    }
                    _ => None,
                },
                ExprType::IsNotNull => match &args[0] {
                    ExprImpl::InputRef(lhs) => {
                        let column_name = &fields[lhs.index].name;
                        let reference = Reference::new(column_name);
                        Some(reference.is_not_null())
                    }
                    _ => None,
                },
                ExprType::In => match &args[0] {
                    ExprImpl::InputRef(lhs) => {
                        let column_name = &fields[lhs.index].name;
                        let reference = Reference::new(column_name);
                        let mut datums = Vec::with_capacity(args.len() - 1);
                        for arg in &args[1..] {
                            if args[0].return_type() != arg.return_type() {
                                return None;
                            }
                            if let ExprImpl::Literal(l) = arg {
                                if let Some(datum) = rw_literal_to_iceberg_datum(l) {
                                    datums.push(datum);
                                } else {
                                    return None;
                                }
                            } else {
                                return None;
                            }
                        }
                        Some(reference.is_in(datums))
                    }
                    _ => None,
                },
                _ => None,
            }
        }
        _ => None,
    }
}
