// 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 itertools::Itertools;
use risingwave_pb::plan_common::JoinType;

use super::prelude::{PlanRef, *};
use crate::expr::{CorrelatedId, CorrelatedInputRef, Expr, ExprImpl, ExprRewriter, InputRef};
use crate::optimizer::plan_node::generic::GenericPlanRef;
use crate::optimizer::plan_node::{
    LogicalApply, LogicalProject, LogicalProjectSet, LogicalValues, PlanTreeNodeUnary,
};

/// Convert a correlated `LogicalApply` with a converted table function RHS into a `LogicalProjectSet`
/// over LHS, so that we can avoid the Domain/Distinct introduced by general apply translation.
///
/// This targets the pattern produced by [`TableFunctionToProjectSetRule`]:
///
/// ```text
/// LogicalApply (Inner, on true)
///  /                       \\
/// LHS                 LogicalProject
///                          |
///                    LogicalProjectSet
///                          |
///                     LogicalValues
/// ```
///
/// After:
///
/// ```text
/// LogicalProject
///      |
/// LogicalProjectSet
///      |
///     LHS
/// ```
///
/// The resulting `ProjectSet` is stateless in streaming, and its hidden `projected_row_id` can be
/// used to derive ordinality.
///
/// [`TableFunctionToProjectSetRule`]: crate::optimizer::rule::TableFunctionToProjectSetRule
pub struct ApplyTableFunctionToProjectSetRule {}

impl Rule<Logical> for ApplyTableFunctionToProjectSetRule {
    fn apply(&self, plan: PlanRef) -> Option<PlanRef> {
        let apply: &LogicalApply = plan.as_logical_apply()?;
        let (left, right, on, join_type, correlated_id, correlated_indices, max_one_row) =
            apply.clone().decompose();

        if max_one_row || join_type != JoinType::Inner || !on.always_true() {
            return None;
        }

        // Only handle RHS generated by `TableFunctionToProjectSetRule`.
        let right_project = right.as_logical_project()?;
        let right_project_input = right_project.input();
        let right_project_set: &LogicalProjectSet = right_project_input.as_logical_project_set()?;
        let right_project_set_input = right_project_set.input();
        let right_values: &LogicalValues = right_project_set_input.as_logical_values()?;
        if !right_values.is_empty_scalar() {
            return None;
        }

        let left_len = left.schema().len();

        // Build a new ProjectSet over LHS:
        // - keep all LHS columns, so they are duplicated per unnested row
        // - append RHS ProjectSet expressions, rewriting CorrelatedInputRef -> InputRef on LHS
        let mut select_list: Vec<ExprImpl> = left
            .schema()
            .data_types()
            .into_iter()
            .enumerate()
            .map(|(idx, ty)| InputRef::new(idx, ty).into())
            .collect();

        let mut corr_rewriter =
            CorrelatedInputRefToInputRefRewriter::new(correlated_id, correlated_indices, left_len);
        let rhs_select_list = right_project_set
            .select_list()
            .iter()
            .cloned()
            .map(|e| corr_rewriter.rewrite_expr(e))
            .collect_vec();
        if !corr_rewriter.touched() {
            // If RHS doesn't reference LHS, this is likely not a lateral table function apply.
            return None;
        }
        select_list.extend(rhs_select_list);

        let new_project_set: PlanRef = LogicalProjectSet::new(left.clone(), select_list).into();

        // Reconstruct the RHS `LogicalProject` on top of the new ProjectSet output.
        // New ProjectSet output layout:
        //   0: projected_row_id
        //   1..=left_len: LHS columns
        //   left_len+1..: RHS ProjectSet outputs (old idx >= 1)
        //
        // We output: [LHS columns..., RHS project exprs...]
        let mut out_exprs: Vec<ExprImpl> = left
            .schema()
            .data_types()
            .into_iter()
            .enumerate()
            .map(|(idx, ty)| InputRef::new(idx + 1, ty).into())
            .collect();

        let mut shift_rewriter = ShiftRhsInputRefRewriter::new(left_len);
        let rhs_project_exprs = right_project
            .exprs()
            .iter()
            .cloned()
            .map(|e| shift_rewriter.rewrite_expr(e))
            .collect_vec();
        out_exprs.extend(rhs_project_exprs);

        Some(LogicalProject::new(new_project_set, out_exprs).into())
    }
}

impl ApplyTableFunctionToProjectSetRule {
    pub fn create() -> BoxedRule {
        Box::new(ApplyTableFunctionToProjectSetRule {})
    }
}

struct CorrelatedInputRefToInputRefRewriter {
    correlated_id: CorrelatedId,
    correlated_indices: Vec<usize>,
    left_len: usize,
    touched: bool,
}

impl CorrelatedInputRefToInputRefRewriter {
    fn new(correlated_id: CorrelatedId, correlated_indices: Vec<usize>, left_len: usize) -> Self {
        Self {
            correlated_id,
            correlated_indices,
            left_len,
            touched: false,
        }
    }

    fn touched(&self) -> bool {
        self.touched
    }
}

impl ExprRewriter for CorrelatedInputRefToInputRefRewriter {
    fn rewrite_correlated_input_ref(
        &mut self,
        correlated_input_ref: CorrelatedInputRef,
    ) -> ExprImpl {
        if correlated_input_ref.correlated_id() != self.correlated_id {
            return correlated_input_ref.into();
        }
        let idx = correlated_input_ref.index();
        if idx >= self.left_len || !self.correlated_indices.contains(&idx) {
            // Be conservative: only rewrite the correlated indices recorded on the apply.
            return correlated_input_ref.into();
        }
        self.touched = true;
        InputRef::new(idx, correlated_input_ref.return_type()).into()
    }
}

struct ShiftRhsInputRefRewriter {
    left_len: usize,
}

impl ShiftRhsInputRefRewriter {
    fn new(left_len: usize) -> Self {
        Self { left_len }
    }
}

impl ExprRewriter for ShiftRhsInputRefRewriter {
    fn rewrite_input_ref(&mut self, input_ref: InputRef) -> ExprImpl {
        if input_ref.index() == 0 {
            // projected_row_id stays at 0
            input_ref.into()
        } else {
            InputRef::new(input_ref.index() + self.left_len, input_ref.return_type()).into()
        }
    }
}
