/* Copyright (c) 2021, 2025, Oracle and/or its affiliates.

   This program is free software; you can redistribute it and/or modify
   it under the terms of the GNU General Public License, version 2.0,
   as published by the Free Software Foundation.

   This program is designed to work with certain software (including
   but not limited to OpenSSL) that is licensed under separate terms,
   as designated in a particular file or component or in included license
   documentation.  The authors of MySQL hereby grant you an additional
   permission to link the program and your derivative works with the
   separately licensed software that they have either included with
   the program or referenced in the documentation.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU General Public License, version 2.0, for more details.

   You should have received a copy of the GNU General Public License
   along with this program; if not, write to the Free Software
   Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301  USA */

#ifndef UNITTEST_GUNIT_OPTIMIZER_TEST_H
#define UNITTEST_GUNIT_OPTIMIZER_TEST_H

#include <assert.h>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include <string.h>
#include <initializer_list>
#include <memory>
#include <string>
#include <unordered_map>
#include <vector>

#include "lex_string.h"
#include "mem_root_deque.h"
#include "my_alloc.h"
#include "my_sqlcommand.h"
#include "my_table_map.h"
#include "sql/field.h"
#include "sql/handler.h"
#include "sql/item.h"
#include "sql/item_subselect.h"
#include "sql/join_optimizer/bit_utils.h"
#include "sql/nested_join.h"
#include "sql/query_options.h"
#include "sql/sql_class.h"
#include "sql/sql_cmd.h"
#include "sql/sql_const.h"
#include "sql/sql_executor.h"
#include "sql/sql_lex.h"
#include "sql/sql_list.h"
#include "sql/sql_opt_exec_shared.h"
#include "sql/sql_optimizer.h"
#include "sql/sql_select.h"
#include "sql/table.h"
#include "sql/visible_fields.h"
#include "template_utils.h"
#include "unittest/gunit/fake_table.h"
#include "unittest/gunit/handler-t.h"
#include "unittest/gunit/mock_field_long.h"
#include "unittest/gunit/parsertest.h"
#include "unittest/gunit/test_utils.h"

inline Query_block *ParseAndResolve(
    const char *query, bool nullable, const Server_initializer &initializer,
    std::unordered_map<string, Fake_TABLE *> *fake_tables);
inline void ResolveQueryBlock(
    THD *thd, Query_block *query_block, bool nullable,
    std::unordered_map<string, Fake_TABLE *> *fake_tables);
inline void ResolveFieldToFakeTable(
    Item *item, const std::unordered_map<string, Fake_TABLE *> &fake_tables);
inline void ResolveAllFieldsToFakeTable(
    const mem_root_deque<Table_ref *> &join_list,
    const std::unordered_map<string, Fake_TABLE *> &fake_tables);
inline void SetJoinConditions(const mem_root_deque<Table_ref *> &join_list);

inline void DestroyFakeTables(
    const std::unordered_map<string, Fake_TABLE *> &fake_tables) {
  for (const auto &[name, table] : fake_tables) {
    destroy(table);
  }
}

namespace optimizer_test {
// Structure for storing table related information used by
// the old join optimizer.
struct Table {
  std::string m_table_name;
  /// Index of the table in the final plan generated by optimizer.
  int m_plan_idx;
  /// Tables including the current table that are looked at for this
  /// table's join. m_prefix_tables is needed for make_cond_for_table().
  table_map m_prefix_tables;

  Table(std::string table_name, int plan_idx, table_map prefix_tables)
      : m_table_name(table_name),
        m_plan_idx(plan_idx),
        m_prefix_tables(prefix_tables) {}
};

}  // namespace optimizer_test

// Base class for the optimizer unit tests (both old and new optimizer).
class OptimizerTestBase : public ::testing::Test {
 public:
  void SetUp() override {
    m_initializer.SetUp();
    m_thd = m_initializer.thd();
  }
  void TearDown() override {
    ClearFakeTables();
    m_initializer.TearDown();
  }

 protected:
  Query_block *ParseAndResolve(const char *query, bool nullable) {
    return ::ParseAndResolve(query, nullable, m_initializer, &m_fake_tables);
  }
  void ClearFakeTables() {
    DestroyFakeTables(m_fake_tables);
    m_fake_tables.clear();
  }
  handlerton *EnableSecondaryEngine(bool aggregation_is_unordered);
  void SetUpJoinTabs(Query_block *query_block, int num_tables,
                     std::vector<optimizer_test::Table> tables);
  void SetUpQEPTabs(Query_block *query_block, int num_tables,
                    std::vector<optimizer_test::Table> tables);

  Server_initializer m_initializer;
  THD *m_thd = nullptr;
  std::unordered_map<std::string, Fake_TABLE *> m_fake_tables;
};

// Template for parameterized optimizer tests.
template <class T>
class OptimizerTestWithParam : public OptimizerTestBase,
                               public ::testing::WithParamInterface<T> {};

inline Query_block *ParseAndResolve(
    const char *query, bool nullable, const Server_initializer &initializer,
    std::unordered_map<string, Fake_TABLE *> *fake_tables) {
  Query_block *query_block = ::parse(&initializer, query, 0);
  ResolveQueryBlock(initializer.thd(), query_block, nullable, fake_tables);
  return query_block;
}

inline void ResolveQueryBlock(
    THD *thd, Query_block *query_block, bool nullable,
    std::unordered_map<string, Fake_TABLE *> *fake_tables) {
  Query_block *const saved_current_query_block =
      thd->lex->current_query_block();
  thd->lex->set_current_query_block(query_block);

  // The hypergraph optimizer does not do const tables,
  // nor does it evaluate subqueries during optimization.
  query_block->add_active_options(OPTION_NO_CONST_TABLES |
                                  OPTION_NO_SUBQUERY_DURING_OPTIMIZATION);

  // Create fake TABLE objects for all tables mentioned in the query.
  int num_tables = 0;
  for (Table_ref *tl = query_block->get_table_list(); tl != nullptr;
       tl = tl->next_global) {
    // If we already have created a fake table with this name (for example to
    // get columns of specific types), use that one. Otherwise, create a new one
    // with four integer columns.
    Fake_TABLE *&fake_table = (*fake_tables)[tl->alias];
    if (fake_table == nullptr) {
      List<Field> fields;
      for (const char *field_name : {"x", "y", "z", "w"}) {
        fields.push_back(new (thd->mem_root) Mock_field_long(
            field_name, nullable, /*is_unsigned=*/false));
      }
      fake_table = new (thd->mem_root) Fake_TABLE(fields);
    }
    fake_table->alias = tl->alias;
    fake_table->pos_in_table_list = tl;
    fake_table->s->db = {tl->db, tl->db_length};
    fake_table->s->table_name = {tl->table_name, tl->table_name_length};
    tl->table = fake_table;
    tl->set_tableno(num_tables++);
    tl->set_updatable();
    tl->grant.privilege = ~(Access_bitmask)0;
  }

  // Find all Item_field objects, and resolve them to fields in the fake tables.
  ResolveAllFieldsToFakeTable(query_block->m_table_nest, *fake_tables);

  // Also in any conditions and subqueries within the WHERE condition.
  if (query_block->where_cond() != nullptr) {
    WalkItem(query_block->where_cond(), enum_walk::POSTFIX, [&](Item *item) {
      if (item->type() == Item::SUBSELECT_ITEM &&
          down_cast<Item_subselect *>(item)->substype() ==
              Item_subselect::IN_SUBS) {
        Item_in_subselect *item_subselect =
            down_cast<Item_in_subselect *>(item);
        ResolveFieldToFakeTable(item_subselect->left_expr, *fake_tables);
        Query_block *child_query_block =
            item_subselect->unit->first_query_block();
        ResolveAllFieldsToFakeTable(child_query_block->m_table_nest,
                                    *fake_tables);
        if (child_query_block->where_cond() != nullptr) {
          ResolveFieldToFakeTable(child_query_block->where_cond(),
                                  *fake_tables);
        }
        for (Item *field_item : child_query_block->fields) {
          ResolveFieldToFakeTable(field_item, *fake_tables);
        }
        return true;  // Don't go down into item_subselect->left_expr again.
      } else if (item->type() == Item::FIELD_ITEM) {
        ResolveFieldToFakeTable(item, *fake_tables);
      }
      return false;
    });
  }

  // And in the SELECT, GROUP BY and ORDER BY lists.
  for (Item *item : query_block->fields) {
    ResolveFieldToFakeTable(item, *fake_tables);
  }
  for (ORDER *cur_group = query_block->group_list.first; cur_group != nullptr;
       cur_group = cur_group->next) {
    ResolveFieldToFakeTable(*cur_group->item, *fake_tables);
  }
  for (ORDER *cur_group = query_block->order_list.first; cur_group != nullptr;
       cur_group = cur_group->next) {
    ResolveFieldToFakeTable(*cur_group->item, *fake_tables);
  }

  // Set up necessary context for UPDATE and single-table DELETE.
  if (thd->lex->sql_command == SQLCOM_DELETE ||
      thd->lex->sql_command == SQLCOM_UPDATE_MULTI) {
    assert(query_block->context.table_list == nullptr);
    assert(query_block->context.first_name_resolution_table == nullptr);
    query_block->context.table_list =
        query_block->context.first_name_resolution_table =
            query_block->get_table_list();
  }

  query_block->prepare(thd, nullptr);

  // Mark deleted and updated tables.
  switch (thd->lex->sql_command) {
    case SQLCOM_DELETE:
    case SQLCOM_DELETE_MULTI:
      for (Table_ref *tl = query_block->leaf_tables; tl != nullptr;
           tl = tl->next_leaf) {
        if (tl->updating) {
          tl->set_deleted();
        }
      }
      break;
    case SQLCOM_UPDATE:
    case SQLCOM_UPDATE_MULTI: {
      table_map update_tables = 0;
      for (Item *item : query_block->visible_fields()) {
        update_tables |= item->used_tables();
      }
      for (Table_ref *tl = query_block->leaf_tables; tl != nullptr;
           tl = tl->next_leaf) {
        if (Overlaps(tl->map(), update_tables)) {
          tl->set_updated();
        }
      }
    } break;
    default:
      break;
  }

  // Create a fake, tiny JOIN. (This would normally be done in optimization.)
  query_block->join = new (thd->mem_root) JOIN(thd, query_block);
  query_block->join->where_cond = query_block->where_cond();
  query_block->join->having_cond = query_block->having_cond();
  query_block->join->fields = &query_block->fields;
  query_block->join->alloc_func_list();
  SetJoinConditions(query_block->m_table_nest);
  count_field_types(query_block, &query_block->join->tmp_table_param,
                    query_block->fields, /*reset_with_sum_func=*/false,
                    /*save_sum_fields=*/false);

  if (query_block->select_limit != nullptr) {
    query_block->master_query_expression()->select_limit_cnt =
        query_block->select_limit->val_int();
  }

  if (query_block->offset_limit != nullptr) {
    query_block->master_query_expression()->offset_limit_cnt =
        query_block->offset_limit->val_int();
  }

  thd->lex->set_current_query_block(saved_current_query_block);
}

inline void ResolveFieldToFakeTable(
    Item *item_arg,
    const std::unordered_map<string, Fake_TABLE *> &fake_tables) {
  WalkItem(item_arg, enum_walk::POSTFIX, [&](Item *item) {
    if (item->type() == Item::FIELD_ITEM) {
      Item_field *item_field = down_cast<Item_field *>(item);
      Fake_TABLE *table = fake_tables.at(item_field->table_name);
      item_field->table_ref = table->pos_in_table_list;
      Field *field = nullptr;
      if (strcmp(item_field->field_name, "x") == 0) {
        field = table->field[0];
      } else if (strcmp(item_field->field_name, "y") == 0) {
        field = table->field[1];
      } else if (strcmp(item_field->field_name, "z") == 0) {
        field = table->field[2];
      } else if (strcmp(item_field->field_name, "w") == 0) {
        field = table->field[3];
      } else {
        assert(false);
      }
      item_field->field = field;
      item_field->set_nullable(field->is_nullable());
      item_field->set_data_type(field->type());
      item_field->collation.set(field->charset(), field->derivation(),
                                field->repertoire());
    }
    return false;
  });
}

inline void ResolveAllFieldsToFakeTable(
    const mem_root_deque<Table_ref *> &join_list,
    const std::unordered_map<string, Fake_TABLE *> &fake_tables) {
  for (Table_ref *tl : join_list) {
    if (tl->join_cond() != nullptr) {
      ResolveFieldToFakeTable(tl->join_cond(), fake_tables);
    }
    if (tl->nested_join != nullptr) {
      ResolveAllFieldsToFakeTable(tl->nested_join->m_tables, fake_tables);
    }
  }
}

inline void SetJoinConditions(const mem_root_deque<Table_ref *> &join_list) {
  for (Table_ref *tl : join_list) {
    tl->set_join_cond_optim(tl->join_cond());
    if (tl->nested_join != nullptr) {
      SetJoinConditions(tl->nested_join->m_tables);
    }
  }
}

inline handlerton *OptimizerTestBase::EnableSecondaryEngine(
    bool aggregation_is_unordered) {
  auto hton = new (m_thd->mem_root) Fake_handlerton;
  hton->flags = HTON_SUPPORTS_SECONDARY_ENGINE;
  if (aggregation_is_unordered) {
    hton->secondary_engine_flags =
        MakeSecondaryEngineFlags(SecondaryEngineFlag::SUPPORTS_HASH_JOIN,
                                 SecondaryEngineFlag::AGGREGATION_IS_UNORDERED);
  } else {
    hton->secondary_engine_flags =
        MakeSecondaryEngineFlags(SecondaryEngineFlag::SUPPORTS_HASH_JOIN);
  }
  hton->secondary_engine_modify_access_path_cost = nullptr;

  for (const auto &[name, table] : m_fake_tables) {
    table->file->ht = hton;
    static_cast<Fake_TABLE_SHARE *>(table->s)->set_secondary_engine(true);
    ON_CALL(table->mock_handler, table_type())
        .WillByDefault(testing::Return("unit test"));
  }

  m_thd->lex->m_sql_cmd->use_secondary_storage_engine(hton);

  return hton;
}

inline void OptimizerTestBase::SetUpJoinTabs(
    Query_block *query_block, int num_tables,
    std::vector<optimizer_test::Table> tables) {
  // This is done during optimization. We prepare JOIN_TABs to be used
  // by make_cond_for_table() and substitute_for_best_equal_field()
  JOIN *join = query_block->join;
  join->join_tab = m_thd->mem_root->ArrayAlloc<JOIN_TAB>(join->tables);
  for (int i = 0; i < num_tables; i++) {
    JOIN_TAB *join_tab = &join->join_tab[i];
    join_tab->set_qs(new (m_thd->mem_root) QEP_shared);
    join_tab->set_idx(tables[i].m_plan_idx);
    join_tab->set_prefix_tables(tables[i].m_prefix_tables,
                                TablesBetween(0, tables[i].m_plan_idx));
    join_tab->set_table(m_fake_tables[tables[i].m_table_name]);
    join_tab->table_ref =
        m_fake_tables[tables[i].m_table_name]->pos_in_table_list;
    join_tab->set_join(join);
    join_tab->set_condition(join_tab->table_ref->join_cond_optim());
  }
}

inline void OptimizerTestBase::SetUpQEPTabs(
    Query_block *query_block, int num_tables,
    std::vector<optimizer_test::Table> tables) {
  // Usually set up by the optimizer at the end of optimization to be used by
  // the executor.
  JOIN *join = query_block->join;
  join->qep_tab = m_thd->mem_root->ArrayAlloc<QEP_TAB>(join->tables);
  std::string tbl_name;
  for (int i = 0; i < num_tables; i++) {
    QEP_TAB *qep_tab = &join->qep_tab[i];
    qep_tab->set_qs(new (m_thd->mem_root) QEP_shared);
    qep_tab->set_idx(tables[i].m_plan_idx);
    qep_tab->set_table(m_fake_tables[tables[i].m_table_name]);
    qep_tab->table_ref =
        m_fake_tables[tables[i].m_table_name]->pos_in_table_list;
    qep_tab->set_join(join);
    // We set operation type as BNL so that executor can choose hash joins.
    qep_tab->op_type = QEP_TAB::OT_BNL;
    qep_tab->set_type(JT_ALL);
  }
}

#endif  // UNITTEST_GUNIT_OPTIMIZER_TEST_H