// Copyright 2025 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 itertools::Itertools;
use risingwave_common::catalog::Field;
use risingwave_common::gap_fill::FillStrategy;
use risingwave_common::types::{DataType, Interval};
use risingwave_sqlparser::ast::{Expr as AstExpr, FunctionArg, FunctionArgExpr, TableAlias};

use super::{Binder, Relation};
use crate::binder::BoundFillStrategy;
use crate::error::{ErrorCode, Result as RwResult};
use crate::expr::{Expr, ExprImpl, InputRef};

#[derive(Debug, Clone)]
pub struct BoundGapFill {
    pub input: Relation,
    pub time_col: InputRef,
    pub interval: ExprImpl,
    pub fill_strategies: Vec<BoundFillStrategy>,
    pub partition_by_cols: Vec<InputRef>,
}

impl Binder {
    pub(super) fn bind_gap_fill(
        &mut self,
        alias: Option<&TableAlias>,
        args: &[FunctionArg],
    ) -> RwResult<BoundGapFill> {
        if args.len() < 3 {
            return Err(ErrorCode::BindError(
                "GAP_FILL requires at least 3 arguments: input, time_col, interval".to_owned(),
            )
            .into());
        }

        let mut args_iter = args.iter();

        self.push_context();

        let (input, table_name) = self.bind_relation_by_function_arg(
            args_iter.next(),
            "The 1st arg of GAP_FILL should be a table name",
        )?;

        let time_col = self.bind_column_by_function_args(
            args_iter.next(),
            "The 2nd arg of GAP_FILL should be a column name",
        )?;

        if !matches!(
            time_col.data_type,
            DataType::Timestamp | DataType::Timestamptz
        ) {
            return Err(ErrorCode::BindError(
                "The 2nd arg of GAP_FILL should be a column of type timestamp or timestamptz"
                    .to_owned(),
            )
            .into());
        }

        let interval_arg = args_iter.next().unwrap();
        let interval_exprs = self.bind_function_arg(interval_arg)?;
        let interval = Itertools::exactly_one(interval_exprs.into_iter()).map_err(|_| {
            ErrorCode::BindError("The 3rd arg of GAP_FILL should be a single expression".to_owned())
        })?;
        if interval.return_type() != DataType::Interval {
            return Err(ErrorCode::BindError(
                "The 3rd arg of GAP_FILL should be an interval".to_owned(),
            )
            .into());
        }

        // Reject non-positive intervals. Fold constant expressions
        // (e.g. `INTERVAL '2' MINUTE - INTERVAL '3' MINUTE`) so the check is not limited to bare
        // literals.
        if let Some(folded) = interval.try_fold_const()
            && let Some(risingwave_common::types::ScalarImpl::Interval(interval_value)) = folded?
            && interval_value <= Interval::from_month_day_usec(0, 0, 0)
        {
            return Err(
                ErrorCode::BindError("The gap fill interval must be positive".to_owned()).into(),
            );
        }

        let mut fill_strategies = vec![];
        let mut partition_by_cols = vec![];
        let mut seen_partition_by_cols = std::collections::HashSet::new();
        for arg in args_iter {
            if let FunctionArg::Unnamed(FunctionArgExpr::Expr(AstExpr::Function(func))) = arg {
                let name = func.name.0[0].real_value().to_ascii_lowercase();

                // Handle PARTITION_BY(col1, col2, ...) as a special directive
                if name == "partition_by" {
                    if func.arg_list.args.is_empty() {
                        return Err(ErrorCode::BindError(
                            "PARTITION_BY requires at least one column argument".to_owned(),
                        )
                        .into());
                    }
                    for partition_arg in &func.arg_list.args {
                        let arg_exprs = self.bind_function_arg(partition_arg)?;
                        let arg_expr =
                            Itertools::exactly_one(arg_exprs.into_iter()).map_err(|_| {
                                ErrorCode::BindError(
                                    "PARTITION_BY argument should be a single column reference"
                                        .to_owned(),
                                )
                            })?;
                        if let ExprImpl::InputRef(input_ref) = arg_expr {
                            if input_ref.index() == time_col.index() {
                                return Err(ErrorCode::BindError(
                                    "PARTITION_BY cannot include the time column".to_owned(),
                                )
                                .into());
                            }
                            // Canonicalize duplicate partition columns eagerly so downstream
                            // planning always sees a unique partition key list.
                            if !seen_partition_by_cols.insert(input_ref.index()) {
                                continue;
                            }
                            partition_by_cols.push(*input_ref);
                        } else {
                            return Err(ErrorCode::BindError(
                                "PARTITION_BY argument must be a column reference".to_owned(),
                            )
                            .into());
                        }
                    }
                    continue;
                }

                let strategy = match name.as_str() {
                    "interpolate" => FillStrategy::Interpolate,
                    "locf" => FillStrategy::Locf,
                    "keepnull" => FillStrategy::Null,
                    _ => {
                        return Err(ErrorCode::BindError(format!(
                            "Unsupported fill strategy: {}",
                            name
                        ))
                        .into());
                    }
                };

                if func.arg_list.args.len() != 1 {
                    return Err(ErrorCode::BindError(format!(
                        "Fill strategy function {} expects exactly one argument",
                        name
                    ))
                    .into());
                }

                let arg_exprs = self.bind_function_arg(&func.arg_list.args[0])?;
                let arg_expr = Itertools::exactly_one(arg_exprs.into_iter()).map_err(|_| {
                    ErrorCode::BindError(
                        "Fill strategy argument should be a single expression".to_owned(),
                    )
                })?;

                if let ExprImpl::InputRef(input_ref) = arg_expr {
                    if input_ref.index() == time_col.index() {
                        return Err(ErrorCode::BindError(
                            "Cannot apply a fill strategy to the time column".to_owned(),
                        )
                        .into());
                    }
                    // Check datatype for interpolate
                    if matches!(strategy, FillStrategy::Interpolate) {
                        let data_type = &input_ref.data_type;
                        if !data_type.is_numeric() || matches!(data_type, DataType::Serial) {
                            return Err(ErrorCode::BindError(format!(
                                "INTERPOLATE only supports numeric types, got {}",
                                data_type
                            ))
                            .into());
                        }
                    }
                    fill_strategies.push(BoundFillStrategy {
                        strategy,
                        target_col: *input_ref,
                    });
                } else {
                    return Err(ErrorCode::BindError(
                        "Fill strategy argument must be a column reference".to_owned(),
                    )
                    .into());
                }
            } else {
                return Err(ErrorCode::BindError(
                    "Fill strategy must be a function call like LOCF(col) or PARTITION_BY(col)"
                        .to_owned(),
                )
                .into());
            }
        }

        // A PARTITION_BY column defines the time-series identity and is always carried into
        // generated rows, so a fill strategy on it would be silently ignored. Reject it.
        for strategy in &fill_strategies {
            if partition_by_cols
                .iter()
                .any(|c| c.index() == strategy.target_col.index())
            {
                return Err(ErrorCode::BindError(
                    "Cannot apply a fill strategy to a PARTITION_BY column".to_owned(),
                )
                .into());
            }
        }

        let base_columns = std::mem::take(&mut self.context.columns);
        self.pop_context()?;

        let columns = base_columns
            .into_iter()
            .map(|c| (c.is_hidden, c.field))
            .collect::<Vec<(bool, Field)>>();

        // Parse into the relation argument of the `gap_fill` function.
        let (schema_name, table_name) =
            Self::resolve_schema_qualified_name(&self.db_name, &table_name)?;
        self.bind_table_to_context(columns, table_name, schema_name, alias)?;

        Ok(BoundGapFill {
            input,
            time_col: *time_col,
            interval,
            fill_strategies,
            partition_by_cols,
        })
    }
}
