// Copyright 2023 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 std::collections::HashSet;
use std::fmt::Debug;
use std::sync::Arc;

use risingwave_common::array::{ArrayBuilder, ArrayRef, BoolArrayBuilder, DataChunk};
use risingwave_common::bail;
use risingwave_common::row::OwnedRow;
use risingwave_common::types::{DataType, Datum, Scalar, ToOwnedDatum};
use risingwave_common::util::iter_util::ZipEqFast;
use risingwave_expr::expr::{
    AsyncExpression, AsyncExpressionBoxExt, BoxedExpression, ExpressionInfo, SyncExpression,
    SyncExpressionBoxExt,
};
use risingwave_expr::{Result, build_function};

#[derive(Debug)]
pub struct InExpression<E> {
    left: E,
    set: HashSet<Datum>,
    return_type: DataType,
}

impl<E> InExpression<E> {
    pub fn new(left: E, data: impl Iterator<Item = Datum>, return_type: DataType) -> Self {
        Self {
            left,
            set: data.collect(),
            return_type,
        }
    }

    // Returns true if datum exists in set, null if datum is null or datum does not exist in set
    // but null does, and false if neither datum nor null exists in set.
    fn exists(&self, datum: &Datum) -> Option<bool> {
        if datum.is_none() {
            None
        } else if self.set.contains(datum) {
            Some(true)
        } else if self.set.contains(&None) {
            None
        } else {
            Some(false)
        }
    }
}

impl<E: ExpressionInfo> ExpressionInfo for InExpression<E> {
    fn return_type(&self) -> DataType {
        self.return_type.clone()
    }
}

macro_rules! eval_in {
    ($mode:ident, $this:expr, $input:expr) => {{
        let input_array = risingwave_expr::forward!($mode, $this.left, eval($input))?;
        let mut output_array = BoolArrayBuilder::new(input_array.len());
        for (data, vis) in input_array.iter().zip_eq_fast($input.visibility().iter()) {
            if vis {
                // TODO: avoid `to_owned_datum()`
                let ret = $this.exists(&data.to_owned_datum());
                output_array.append(ret);
            } else {
                output_array.append(None);
            }
        }
        Ok(Arc::new(output_array.finish().into()))
    }};
}

macro_rules! eval_row_in {
    ($mode:ident, $this:expr, $input:expr) => {{
        let data = risingwave_expr::forward!($mode, $this.left, eval_row($input))?;
        let ret = $this.exists(&data);
        Ok(ret.map(|b| b.to_scalar_value()))
    }};
}

impl<E: SyncExpression> SyncExpression for InExpression<E> {
    fn eval(&self, input: &DataChunk) -> Result<ArrayRef> {
        eval_in!(sync, self, input)
    }

    fn eval_row(&self, input: &OwnedRow) -> Result<Datum> {
        eval_row_in!(sync, self, input)
    }
}

impl<E: AsyncExpression> AsyncExpression for InExpression<E> {
    async fn eval(&self, input: &DataChunk) -> Result<ArrayRef> {
        eval_in!(async, self, input)
    }

    async fn eval_row(&self, input: &OwnedRow) -> Result<Datum> {
        eval_row_in!(async, self, input)
    }
}

#[build_function("in(any, ...) -> boolean")]
fn build(return_type: DataType, children: Vec<BoxedExpression>) -> Result<BoxedExpression> {
    let mut iter = children.into_iter();
    let left_expr = iter.next().unwrap();
    let mut data = Vec::with_capacity(iter.size_hint().0);
    let data_chunk = DataChunk::new_dummy(1);
    for child in iter {
        let BoxedExpression::Sync(child) = child else {
            bail!("IN set expression must be sync")
        };
        let array = child.eval(&data_chunk)?;
        let datum = array.value_at(0).to_owned_datum();
        data.push(datum);
    }
    Ok(match left_expr {
        BoxedExpression::Sync(left_expr) => {
            InExpression::new(left_expr, data.into_iter(), return_type).boxed()
        }
        left_expr @ BoxedExpression::Async(_) => {
            InExpression::new(left_expr, data.into_iter(), return_type).boxed()
        }
    })
}

#[cfg(test)]
mod tests {
    use risingwave_common::array::DataChunk;
    use risingwave_common::row::Row;
    use risingwave_common::test_prelude::DataChunkTestExt;
    use risingwave_common::types::ToOwnedDatum;
    use risingwave_common::util::iter_util::ZipEqDebug;
    use risingwave_expr::expr::build_from_pretty;

    #[tokio::test]
    async fn test_in_expr() {
        let expr = build_from_pretty("(in:boolean $0:varchar abc:varchar def:varchar)");
        let (input, expected) = DataChunk::from_pretty(
            "T   B
             abc t
             a   f
             def t
             abc t
             .   .",
        )
        .split_column_at(1);

        // test eval
        let output = expr.eval(&input).await.unwrap();
        assert_eq!(&output, expected.column_at(0));

        // test eval_row
        for (row, expected) in input.rows().zip_eq_debug(expected.rows()) {
            let result = expr.eval_row(&row.to_owned_row()).await.unwrap();
            assert_eq!(result, expected.datum_at(0).to_owned_datum());
        }
    }

    #[tokio::test]
    async fn test_in_expr_null() {
        let expr = build_from_pretty("(in:boolean $0:varchar abc:varchar null:varchar)");
        let (input, expected) = DataChunk::from_pretty(
            "T   B
             abc t
             a   .
             .   .",
        )
        .split_column_at(1);

        // test eval
        let output = expr.eval(&input).await.unwrap();
        assert_eq!(&output, expected.column_at(0));

        // test eval_row
        for (row, expected) in input.rows().zip_eq_debug(expected.rows()) {
            let result = expr.eval_row(&row.to_owned_row()).await.unwrap();
            assert_eq!(result, expected.datum_at(0).to_owned_datum());
        }
    }
}
