#include "duckdb/optimizer/unnest_rewriter.hpp"

#include "duckdb/common/pair.hpp"

#include "duckdb/planner/expression/bound_columnref_expression.hpp"

#include "duckdb/planner/expression/bound_unnest_expression.hpp"

#include "duckdb/planner/operator/logical_comparison_join.hpp"

#include "duckdb/planner/operator/logical_delim_get.hpp"

#include "duckdb/planner/operator/logical_projection.hpp"

#include "duckdb/planner/operator/logical_unnest.hpp"

#include "duckdb/planner/operator/logical_window.hpp"

namespace duckdb {

void UnnestRewriterPlanUpdater::VisitOperator(LogicalOperator &op) {

  VisitOperatorChildren(op);

  VisitOperatorExpressions(op);

}

void UnnestRewriterPlanUpdater::VisitExpression(unique_ptr<Expression> *expression) {

  auto &expr = *expression;

  if (expr->GetExpressionClass() == ExpressionClass::BOUND_COLUMN_REF) {

    auto &bound_column_ref = expr->Cast<BoundColumnRefExpression>();

    for (idx_t i = 0; i < replace_bindings.size(); i++) {

      if (bound_column_ref.binding == replace_bindings[i].old_binding) {

        bound_column_ref.binding = replace_bindings[i].new_binding;

        break;

      }

    }

  }

  VisitExpressionChildren(**expression);

}

unique_ptr<LogicalOperator> UnnestRewriter::Optimize(unique_ptr<LogicalOperator> op) {

  UnnestRewriterPlanUpdater updater;

  vector<reference<unique_ptr<LogicalOperator>>> candidates;

  FindCandidates(op, candidates);

  // rewrite the plan and update the bindings

  for (auto &candidate : candidates) {

    // rearrange the logical operators

    if (RewriteCandidate(candidate)) {

      updater.overwritten_tbl_idx = overwritten_tbl_idx;

      // update the bindings of the BOUND_UNNEST expression

      UpdateBoundUnnestBindings(updater, candidate);

      // update the sequence of LOGICAL_PROJECTION(s)

      UpdateRHSBindings(op, candidate, updater);

      // reset

      delim_columns.clear();

      lhs_bindings.clear();

    }

  }

  return op;

}

void UnnestRewriter::FindCandidates(unique_ptr<LogicalOperator> &op,

                                    vector<reference<unique_ptr<LogicalOperator>>> &candidates) {

  // search children before adding, so that we add candidates bottom-up

  for (auto &child : op->children) {

    FindCandidates(child, candidates);

  }

  // search for operator that has a LOGICAL_DELIM_JOIN as its child

  if (op->children.size() != 1) {

    return;

  }

  if (op->children[0]->type != LogicalOperatorType::LOGICAL_DELIM_JOIN) {

    return;

  }

  // found a delim join

  auto &delim_join = op->children[0]->Cast<LogicalComparisonJoin>();

  // only support INNER delim joins

  if (delim_join.join_type != JoinType::INNER) {

    return;

  }

  // INNER delim join must have exactly one condition

  if (delim_join.conditions.size() != 1) {

    return;

  }

  // LHS child is a window

  idx_t delim_idx = delim_join.delim_flipped ? 1 : 0;

  idx_t other_idx = 1 - delim_idx;

  if (delim_join.children[delim_idx]->type != LogicalOperatorType::LOGICAL_WINDOW) {

    return;

  }

  // RHS child must be projection(s) followed by an UNNEST

  auto curr_op = &delim_join.children[other_idx];

  while (curr_op->get()->type == LogicalOperatorType::LOGICAL_PROJECTION) {

    if (curr_op->get()->children.size() != 1) {

      break;

    }

    curr_op = &curr_op->get()->children[0];

  }

  if (curr_op->get()->type == LogicalOperatorType::LOGICAL_UNNEST &&

      curr_op->get()->children[0]->type == LogicalOperatorType::LOGICAL_DELIM_GET) {

    candidates.push_back(op);

  }

}

bool UnnestRewriter::RewriteCandidate(unique_ptr<LogicalOperator> &candidate) {

  auto &topmost_op = *candidate;

  if (topmost_op.type != LogicalOperatorType::LOGICAL_PROJECTION &&

      topmost_op.type != LogicalOperatorType::LOGICAL_WINDOW &&

      topmost_op.type != LogicalOperatorType::LOGICAL_FILTER &&

      topmost_op.type != LogicalOperatorType::LOGICAL_AGGREGATE_AND_GROUP_BY &&

      topmost_op.type != LogicalOperatorType::LOGICAL_UNNEST) {

    return false;

  }

  // get the LOGICAL_DELIM_JOIN, which is a child of the candidate

  D_ASSERT(topmost_op.children.size() == 1);

  auto &delim_join = topmost_op.children[0]->Cast<LogicalComparisonJoin>();

  D_ASSERT(delim_join.type == LogicalOperatorType::LOGICAL_DELIM_JOIN);

  GetDelimColumns(delim_join);

  // LHS of the LOGICAL_DELIM_JOIN is a LOGICAL_WINDOW that contains a LOGICAL_PROJECTION/LOGICAL_CROSS_JOIN

  // this lhs_proj later becomes the child of the UNNEST

  idx_t delim_idx = delim_join.delim_flipped ? 1 : 0;

  idx_t other_idx = 1 - delim_idx;

  auto &window = *delim_join.children[delim_idx];

  auto &lhs_op = window.children[0];

  GetLHSExpressions(*lhs_op);

  // find the LOGICAL_UNNEST

  // and get the path down to the LOGICAL_UNNEST

  vector<unique_ptr<LogicalOperator> *> path_to_unnest;

  auto curr_op = &delim_join.children[other_idx];

  while (curr_op->get()->type == LogicalOperatorType::LOGICAL_PROJECTION) {

    path_to_unnest.push_back(curr_op);

    curr_op = &curr_op->get()->children[0];

  }

  // store the table index of the child of the LOGICAL_UNNEST

  // then update the plan by making the lhs_proj the child of the LOGICAL_UNNEST

  D_ASSERT(curr_op->get()->type == LogicalOperatorType::LOGICAL_UNNEST);

  auto &unnest = curr_op->get()->Cast<LogicalUnnest>();

  D_ASSERT(unnest.children[0]->type == LogicalOperatorType::LOGICAL_DELIM_GET);

  overwritten_tbl_idx = unnest.children[0]->Cast<LogicalDelimGet>().table_index;

  D_ASSERT(!unnest.children.empty());

  auto &delim_get = unnest.children[0]->Cast<LogicalDelimGet>();

  D_ASSERT(delim_get.chunk_types.size() > 1);

  distinct_unnest_count = delim_get.chunk_types.size();

  unnest.children[0] = std::move(lhs_op);

  // replace the LOGICAL_DELIM_JOIN with its RHS child operator

  topmost_op.children[0] = std::move(*path_to_unnest.front());

  return true;

}

void UnnestRewriter::UpdateRHSBindings(unique_ptr<LogicalOperator> &plan, unique_ptr<LogicalOperator> &candidate,

                                       UnnestRewriterPlanUpdater &updater) {

  auto &topmost_op = *candidate;

  idx_t shift = lhs_bindings.size();

  vector<unique_ptr<LogicalOperator> *> path_to_unnest;

  auto curr_op = &topmost_op.children[0];

  while (curr_op->get()->type == LogicalOperatorType::LOGICAL_PROJECTION) {

    path_to_unnest.push_back(curr_op);

    D_ASSERT(curr_op->get()->type == LogicalOperatorType::LOGICAL_PROJECTION);

    auto &proj = curr_op->get()->Cast<LogicalProjection>();

    // pop the unnest columns and the delim index

    D_ASSERT(proj.expressions.size() > distinct_unnest_count);

    for (idx_t i = 0; i < distinct_unnest_count; i++) {

      proj.expressions.pop_back();

    }

    // store all shifted current bindings

    idx_t tbl_idx = proj.table_index;

    for (idx_t i = 0; i < proj.expressions.size(); i++) {

      ReplaceBinding replace_binding(ColumnBinding(tbl_idx, i), ColumnBinding(tbl_idx, i + shift));

      updater.replace_bindings.push_back(replace_binding);

    }

    curr_op = &curr_op->get()->children[0];

  }

  // update all bindings by shifting them

  updater.VisitOperator(*plan);

  updater.replace_bindings.clear();

  // update all bindings coming from the LHS to RHS bindings

  D_ASSERT(topmost_op.children[0]->type == LogicalOperatorType::LOGICAL_PROJECTION);

  auto &top_proj = topmost_op.children[0]->Cast<LogicalProjection>();

  for (idx_t i = 0; i < lhs_bindings.size(); i++) {

    ReplaceBinding replace_binding(lhs_bindings[i].binding, ColumnBinding(top_proj.table_index, i));

    updater.replace_bindings.push_back(replace_binding);

  }

  // temporarily remove the BOUND_UNNESTs and the child of the LOGICAL_UNNEST from the plan

  D_ASSERT(curr_op->get()->type == LogicalOperatorType::LOGICAL_UNNEST);

  auto &unnest = curr_op->get()->Cast<LogicalUnnest>();

  vector<unique_ptr<Expression>> temp_bound_unnests;

  for (auto &temp_bound_unnest : unnest.expressions) {

    temp_bound_unnests.push_back(std::move(temp_bound_unnest));

  }

  D_ASSERT(unnest.children.size() == 1);

  auto temp_unnest_child = std::move(unnest.children[0]);

  unnest.expressions.clear();

  unnest.children.clear();

  // update the bindings of the plan

  updater.VisitOperator(*plan);

  updater.replace_bindings.clear();

  // add the children again

  for (auto &temp_bound_unnest : temp_bound_unnests) {

    unnest.expressions.push_back(std::move(temp_bound_unnest));

  }

  unnest.children.push_back(std::move(temp_unnest_child));

  // add the LHS expressions to each LOGICAL_PROJECTION

  for (idx_t i = path_to_unnest.size(); i > 0; i--) {

    D_ASSERT(path_to_unnest[i - 1]->get()->type == LogicalOperatorType::LOGICAL_PROJECTION);

    auto &proj = path_to_unnest[i - 1]->get()->Cast<LogicalProjection>();

    // temporarily store the existing expressions

    vector<unique_ptr<Expression>> existing_expressions;

    for (idx_t expr_idx = 0; expr_idx < proj.expressions.size(); expr_idx++) {

      existing_expressions.push_back(std::move(proj.expressions[expr_idx]));

    }

    proj.expressions.clear();

    // add the new expressions

    for (idx_t expr_idx = 0; expr_idx < lhs_bindings.size(); expr_idx++) {

      auto new_expr = make_uniq<BoundColumnRefExpression>(

          lhs_bindings[expr_idx].alias, lhs_bindings[expr_idx].type, lhs_bindings[expr_idx].binding);

      proj.expressions.push_back(std::move(new_expr));

      // update the table index

      lhs_bindings[expr_idx].binding.table_index = proj.table_index;

      lhs_bindings[expr_idx].binding.column_index = expr_idx;

    }

    // add the existing expressions again

    for (idx_t expr_idx = 0; expr_idx < existing_expressions.size(); expr_idx++) {

      proj.expressions.push_back(std::move(existing_expressions[expr_idx]));

    }

  }

}

void UnnestRewriter::UpdateBoundUnnestBindings(UnnestRewriterPlanUpdater &updater,

                                               unique_ptr<LogicalOperator> &candidate) {

  auto &topmost_op = *candidate;

  // traverse LOGICAL_PROJECTION(s)

  auto curr_op = &topmost_op.children[0];

  while (curr_op->get()->type == LogicalOperatorType::LOGICAL_PROJECTION) {

    curr_op = &curr_op->get()->children[0];

  }

  // found the LOGICAL_UNNEST

  D_ASSERT(curr_op->get()->type == LogicalOperatorType::LOGICAL_UNNEST);

  auto &unnest = curr_op->get()->Cast<LogicalUnnest>();

  D_ASSERT(unnest.children.size() == 1);

  auto unnest_cols = unnest.children[0]->GetColumnBindings();

  for (idx_t i = 0; i < delim_columns.size(); i++) {

    auto delim_binding = delim_columns[i];

    auto unnest_it = unnest_cols.begin();

    while (unnest_it != unnest_cols.end()) {

      auto unnest_binding = *unnest_it;

      if (delim_binding.table_index == unnest_binding.table_index) {

        unnest_binding.table_index = overwritten_tbl_idx;

        unnest_binding.column_index = i;

        updater.replace_bindings.emplace_back(unnest_binding, delim_binding);

        unnest_cols.erase(unnest_it);

        break;

      }

      unnest_it++;

    }

  }

  // update bindings

  for (auto &unnest_expr : unnest.expressions) {

    updater.VisitExpression(&unnest_expr);

  }

  updater.replace_bindings.clear();

}

void UnnestRewriter::GetDelimColumns(LogicalOperator &op) {

  D_ASSERT(op.type == LogicalOperatorType::LOGICAL_DELIM_JOIN);

  auto &delim_join = op.Cast<LogicalComparisonJoin>();

  for (idx_t i = 0; i < delim_join.duplicate_eliminated_columns.size(); i++) {

    auto &expr = *delim_join.duplicate_eliminated_columns[i];

    D_ASSERT(expr.GetExpressionType() == ExpressionType::BOUND_COLUMN_REF);

    auto &bound_colref_expr = expr.Cast<BoundColumnRefExpression>();

    delim_columns.push_back(bound_colref_expr.binding);

  }

}

void UnnestRewriter::GetLHSExpressions(LogicalOperator &op) {

  op.ResolveOperatorTypes();

  auto col_bindings = op.GetColumnBindings();

  D_ASSERT(op.types.size() == col_bindings.size());

  bool set_alias = false;

  // we can easily extract the alias for LOGICAL_PROJECTION(s)

  if (op.type == LogicalOperatorType::LOGICAL_PROJECTION) {

    auto &proj = op.Cast<LogicalProjection>();

    if (proj.expressions.size() == op.types.size()) {

      set_alias = true;

    }

  }

  for (idx_t i = 0; i < op.types.size(); i++) {

    lhs_bindings.emplace_back(col_bindings[i], op.types[i]);

    if (set_alias) {

      auto &proj = op.Cast<LogicalProjection>();

      lhs_bindings.back().alias = proj.expressions[i]->GetAlias();

    }

  }

}

} // namespace duckdb