// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you 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.
//
// Ksck, a tool to run a Kudu System Check.
#pragma once

#include <array>
#include <atomic>
#include <cstdint>
#include <iosfwd>
#include <map>
#include <memory>
#include <optional>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>

#include <glog/logging.h>
#include <gtest/gtest_prod.h>

#include "kudu/common/partition.h"
#include "kudu/common/schema.h"
#include "kudu/consensus/metadata.pb.h"
#include "kudu/gutil/macros.h"
#include "kudu/gutil/strings/substitute.h"
#include "kudu/rebalance/cluster_status.h" // IWYU pragma: keep
#include "kudu/server/server_base.pb.h" // IWYU pragma: keep
#include "kudu/tablet/metadata.pb.h"
#include "kudu/tablet/tablet.pb.h"  // IWYU pragma: keep
#include "kudu/tools/ksck_results.h"
#include "kudu/util/status.h"
#include "kudu/util/threadpool.h"

namespace kudu {

namespace rpc {
class Messenger;
} // namespace rpc

namespace tools {

class KsckChecksumManager;
class KsckTable;
struct KsckChecksumOptions;

// Representation of a tablet replica on a tablet server.
class KsckTabletReplica {
 public:
  KsckTabletReplica(std::string ts_uuid, bool is_leader, bool is_voter)
      : ts_uuid_(std::move(ts_uuid)),
        is_leader_(is_leader),
        is_voter_(is_voter) {
  }

  const std::string& ts_uuid() const {
    return ts_uuid_;
  }

  bool is_leader() const {
    return is_leader_;
  }

  bool is_voter() const {
    return is_voter_;
  }

 private:
  const std::string ts_uuid_;
  const bool is_leader_;
  const bool is_voter_;

  DISALLOW_COPY_AND_ASSIGN(KsckTabletReplica);
};

// Representation of a tablet belonging to a table. The tablet is composed of replicas.
class KsckTablet {
 public:
  KsckTablet(KsckTable* table, std::string id, Partition partition)
      : table_(table),
        id_(std::move(id)),
        partition_(std::move(partition)) {
  }

  const std::string& id() const {
    return id_;
  }

  const std::vector<std::shared_ptr<KsckTabletReplica>>& replicas() const {
    return replicas_;
  }

  const Partition& partition() const {
    return partition_;
  }

  void set_replicas(std::vector<std::shared_ptr<KsckTabletReplica>> replicas) {
    replicas_.swap(replicas);
  }

  KsckTable* table() {
    return table_;
  }

 private:
  KsckTable* table_;
  const std::string id_;
  const Partition partition_;
  std::vector<std::shared_ptr<KsckTabletReplica>> replicas_;
  DISALLOW_COPY_AND_ASSIGN(KsckTablet);
};

// Representation of a table. Composed of tablets.
class KsckTable {
 public:
  KsckTable(std::string id, std::string name, Schema schema, int num_replicas)
      : id_(std::move(id)),
        name_(std::move(name)),
        schema_(std::move(schema)),
        num_replicas_(num_replicas) {
  }

  const std::string& id() const {
    return id_;
  }

  const std::string& name() const {
    return name_;
  }

  const Schema& schema() const {
    return schema_;
  }

  int num_replicas() const {
    return num_replicas_;
  }

  void set_tablets(std::vector<std::shared_ptr<KsckTablet>> tablets) {
    tablets_ = std::move(tablets);
  }

  const std::vector<std::shared_ptr<KsckTablet>>& tablets() const {
    return tablets_;
  }

 private:
  const std::string id_;
  const std::string name_;
  const Schema schema_;
  const int num_replicas_;
  std::vector<std::shared_ptr<KsckTablet>> tablets_;
  DISALLOW_COPY_AND_ASSIGN(KsckTable);
};

// Enum representing the fetch status of a ksck master or tablet server.
enum class KsckFetchState {
  // Information has not yet been fetched.
  UNINITIALIZED,
  // The attempt to fetch information failed.
  FETCH_FAILED,
  // Information was fetched successfully.
  FETCHED,
};

// Required for logging in case of CHECK failures.
std::ostream& operator<<(std::ostream& lhs, KsckFetchState state);

// Enum representing logical categories of flags used in a Kudu cluster.
enum FlagsCategory {
  // Flags specific to the time source, clock, and alike.
  TIME_SOURCE = 0,

  // Flags tagged hidden, experimental, or unsafe.
  UNUSUAL = 1,

  // Utility members used for range-related expressions.
  // NOTE: update the MAX member upon updates
  MIN = 0,
  MAX = UNUSUAL
};

// Information on flags fetched using GetFlags() RPC. Consists of the status of
// the GetFlags() operation and the result flags wrapped into std::optional,
// where the latter is std::nullopt unless the RPC was successful.
struct FetchedFlags {
  KsckFetchState state = KsckFetchState::UNINITIALIZED;
  std::optional<server::GetFlagsResponsePB> flags;
};

// Flags retrieved using GetFlags(), indexed by FlagsCategory.
typedef std::array<FetchedFlags, FlagsCategory::MAX + 1> FetchedFlagsByCategory;

// Structure to represent a filter to build appropriate request for GetFlags()
// RPC.
struct FlagsFetchFilter {
  const std::vector<std::string> flags;
  const std::vector<std::string> tags;
};

// Get filter to build GetFlags() RPC for the specified flags category.
const FlagsFetchFilter& GetFlagsCategoryFilter(FlagsCategory category);

// Get string representation for the specified flag category.
const char* FlagsCategoryToString(FlagsCategory category);

// Convert string representation of flag category into corresponding enum field.
// On success, returns Status::OK() and outputs the result category into the
// 'category' output parameter. On failure returns Status::InvalidParameter().
Status StringToFlagsCategory(const std::string& str,
                             FlagsCategory* category);

// Convert comma-separated string into vector of flag categories. The result
// is appended to the 'categories' output parameter. In addition, this function
// removes duplicates in resulting 'categories'.
Status StringToFlagsCategories(const std::string& str,
                               std::vector<FlagsCategory>* categories);

// The following three classes must be extended in order to communicate with their respective
// components. The two main use cases envisioned for this are:
// - To be able to mock a cluster to more easily test the ksck checks.
// - To be able to communicate with a real Kudu cluster.

// Class that must be extended to represent a master.
class KsckMaster {
 public:
  explicit KsckMaster(std::string address) :
      address_(std::move(address)),
      uuid_(strings::Substitute("$0 ($1)", kDummyUuid, address_)) {}

  virtual ~KsckMaster() = default;

  virtual Status Init() = 0;

  // Connects to the master, checking if it is healthy and gathering basic info.
  virtual Status FetchInfo() = 0;

  // Connects to the master and populates the consensus map.
  virtual Status FetchConsensusState() = 0;

  // Fetch flags for the requested categories.
  virtual Status FetchFlags(const std::vector<FlagsCategory>& categories) = 0;

  // Since masters are provided by address, FetchInfo() must be called before
  // calling this method.
  virtual const std::string& uuid() const {
    CHECK_NE(KsckFetchState::UNINITIALIZED, state_);
    return uuid_;
  }

  virtual const std::string& address() const {
    return address_;
  }

  virtual const std::optional<std::string>& version() const {
    CHECK_NE(KsckFetchState::UNINITIALIZED, state_);
    return version_;
  }

  virtual const std::optional<consensus::ConsensusStatePB>& cstate() const {
    CHECK_NE(KsckFetchState::UNINITIALIZED, state_);
    return cstate_;
  }

  virtual const std::optional<server::GetFlagsResponsePB>& flags(
      FlagsCategory category) const {
    CHECK_GE(category, FlagsCategory::MIN);
    CHECK_LE(category, FlagsCategory::MAX);
    CHECK_NE(KsckFetchState::UNINITIALIZED, flags_by_category_[category].state);
    return flags_by_category_[category].flags;
  }

  std::string ToString() const {
    return strings::Substitute("$0 ($1)", uuid(), address());
  }

  bool is_healthy() const {
    CHECK_NE(KsckFetchState::UNINITIALIZED, state_);
    return state_ == KsckFetchState::FETCHED;
  }

  // Masters that haven't been fetched from or that were unavailable have a
  // dummy uuid.
  static constexpr const char* const kDummyUuid = "<unknown>";

 protected:
  friend class KsckTest;

  const std::string address_;
  std::string uuid_;

  // state_ reflects whether the fetch of the core ksck info has been done, and
  // if it succeeded or failed.
  KsckFetchState state_ = KsckFetchState::UNINITIALIZED;

  // May be none if fetching info from the master fails.
  std::optional<std::string> version_;

  // Fetched flags indexed by category.
  FetchedFlagsByCategory flags_by_category_;

  // May be none if consensus state fetch fails.
  std::optional<consensus::ConsensusStatePB> cstate_;

 private:
  DISALLOW_COPY_AND_ASSIGN(KsckMaster);
};

// Class that must be extended to represent a tablet server.
class KsckTabletServer {
 public:
  // Map from tablet id to tablet replicas.
  typedef std::unordered_map<std::string, tablet::TabletStatusPB> TabletStatusMap;

  // Map from (tserver id, tablet id) to tablet consensus information.
  typedef std::map<std::pair<std::string, std::string>,
                   consensus::ConsensusStatePB>
      TabletConsensusStateMap;

  explicit KsckTabletServer(std::string uuid, std::string location = "")
      : uuid_(std::move(uuid)),
        location_(std::move(location)) {}

  virtual ~KsckTabletServer() {}

  // Connects to the configured tablet server and populates the fields of this class. 'health' must
  // not be nullptr.
  //
  // If Status is OK, 'health' will be HEALTHY
  // If the UUID is not what ksck expects, 'health' will be WRONG_SERVER_UUID
  // Otherwise 'health' will be UNAVAILABLE
  virtual Status FetchInfo(cluster_summary::ServerHealth* health) = 0;

  // Connects to the configured tablet server and populates the consensus map. 'health' must not be
  // nullptr.
  //
  // If Status is OK, 'health' will be HEALTHY
  // Otherwise 'health' will be UNAVAILABLE
  virtual Status FetchConsensusState(cluster_summary::ServerHealth* health) = 0;

  // Fetches flags for the requested categories.
  virtual Status FetchFlags(const std::vector<FlagsCategory>& categories) = 0;

  // Fetches and updates the current timestamp from the tablet server.
  virtual void FetchCurrentTimestampAsync() = 0;
  virtual Status FetchCurrentTimestamp() = 0;

  // Fetches the quiescing information for the tablet server. This is best
  // effort; a message is logged if it is unsuccessful (e.g. if the server
  // doesn't support the quiescing RPC).
  virtual void FetchQuiescingInfo() = 0;

  // Executes a checksum scan on a tablet and reports the result to 'manager'.
  virtual void RunTabletChecksumScanAsync(
                  const std::string& tablet_id,
                  const Schema& schema,
                  const KsckChecksumOptions& options,
                  std::shared_ptr<KsckChecksumManager> manager) = 0;

  virtual const std::string& uuid() const {
    return uuid_;
  }

  virtual const std::string& location() const {
    return location_;
  }

  std::string ToString() const {
    return strings::Substitute("$0 ($1)", uuid(), address());
  }

  virtual std::string address() const = 0;

  bool is_healthy() const {
    CHECK_NE(KsckFetchState::UNINITIALIZED, state_);
    return state_ == KsckFetchState::FETCHED;
  }

  // Gets the mapping of tablet id to tablet replica for this tablet server.
  const TabletStatusMap& tablet_status_map() const {
    CHECK_EQ(KsckFetchState::FETCHED, state_);
    return tablet_status_map_;
  }

  // Gets the mapping of tablet id to tablet consensus info for this tablet server.
  const TabletConsensusStateMap& tablet_consensus_state_map() const {
    CHECK_EQ(KsckFetchState::FETCHED, state_);
    return tablet_consensus_state_map_;
  }

  tablet::TabletStatePB ReplicaState(const std::string& tablet_id) const;

  virtual const std::optional<std::string>& version() const {
    CHECK_NE(KsckFetchState::UNINITIALIZED, state_);
    return version_;
  }

  virtual const std::optional<server::GetFlagsResponsePB>& flags(
      FlagsCategory category) const {
    DCHECK_GE(category, FlagsCategory::MIN);
    DCHECK_LE(category, FlagsCategory::MAX);
    CHECK_NE(KsckFetchState::UNINITIALIZED, flags_by_category_[category].state);
    return flags_by_category_[category].flags;
  }

  virtual const std::optional<cluster_summary::QuiescingInfo>& quiescing_info() const {
    CHECK_NE(KsckFetchState::UNINITIALIZED, state_);
    return quiescing_info_;
  }

  uint64_t current_timestamp() const {
    CHECK_EQ(KsckFetchState::FETCHED, state_);
    return timestamp_;
  }

 protected:
  friend class KsckTest;
  FRIEND_TEST(KsckTest, TestConsensusConflictExtraPeer);
  FRIEND_TEST(KsckTest, TestConsensusConflictDifferentLeader);
  FRIEND_TEST(KsckTest, TestConsensusConflictMissingPeer);
  FRIEND_TEST(KsckTest, TestMasterNotReportingTabletServerWithConsensusConflict);
  FRIEND_TEST(KsckTest, TestMismatchedAssignments);
  FRIEND_TEST(KsckTest, TestTabletCopying);

  // state_ reflects whether the fetch of the core info has been done, and if
  // it succeeded or failed.
  KsckFetchState state_ = KsckFetchState::UNINITIALIZED;

  // May be none if fetching info from the tablet server fails.
  std::optional<std::string> version_;

  // Fetched flags indexed by category.
  FetchedFlagsByCategory flags_by_category_;

  TabletStatusMap tablet_status_map_;
  TabletConsensusStateMap tablet_consensus_state_map_;

  // May be none if the quiescing request fails.
  std::optional<cluster_summary::QuiescingInfo> quiescing_info_;

  const std::string uuid_;
  std::string location_;
  std::atomic<uint64_t> timestamp_;

 private:
  DISALLOW_COPY_AND_ASSIGN(KsckTabletServer);
};

// Class used to communicate with a cluster.
class KsckCluster {
 public:
  virtual ~KsckCluster() = default;

  // A list of masters.
  typedef std::vector<std::shared_ptr<KsckMaster>> MasterList;

  // Map of KsckTabletServer objects keyed by tablet server uuid.
  typedef std::map<std::string, std::shared_ptr<KsckTabletServer>> TSMap;

  // Fetches the lists of tables, tablets, and tablet servers from the master.
  Status FetchTableAndTabletInfo() {
    RETURN_NOT_OK(Connect());
    RETURN_NOT_OK(RetrieveTablesList());
    RETURN_NOT_OK(RetrieveTabletServers());
    RETURN_NOT_OK(RetrieveAllTablets());
    return Status::OK();
  }

  // Connects to the cluster (i.e. to the leader master).
  virtual Status Connect() = 0;

  // Fetches the list of tablet servers.
  virtual Status RetrieveTabletServers() = 0;

  // Fetches the list of tables.
  virtual Status RetrieveTablesList() = 0;

  // Fetches all tablets in the cluster.
  virtual Status RetrieveAllTablets() = 0;

  // Fetches the list of tablets for the given table.
  // The table's tablet list is modified only if this method returns OK.
  virtual Status RetrieveTabletsList(const std::shared_ptr<KsckTable>& table) = 0;

  const MasterList& masters() const {
    return masters_;
  }

  const TSMap& tablet_servers() const {
    return tablet_servers_;
  }

  const KsckTServerStateMap& ts_states() const {
    return ts_states_;
  }

  const std::shared_ptr<KsckTable>& txn_sys_table() const {
    return txn_sys_table_;
  }

  const std::vector<std::shared_ptr<KsckTable>>& tables() const {
    return tables_;
  }

  // Returns a reference to the messenger used by this instance.
  // Returns nullptr if no messenger is used.
  virtual std::shared_ptr<rpc::Messenger> messenger() const {
    return nullptr;
  }

  // Setters for filtering the tables/tablets to be checked.
  // Equivalent to the same functions in class 'Ksck'.
  void set_table_filters(std::vector<std::string> table_names) {
    table_filters_ = std::move(table_names);
  }

  // See above.
  void set_tablet_id_filters(std::vector<std::string> tablet_ids) {
    tablet_id_filters_ = std::move(tablet_ids);
  }

  const std::vector<std::string>& table_filters() const {
    return table_filters_;
  }

  const std::vector<std::string>& tablet_id_filters() const {
    return tablet_id_filters_;
  }

  int filtered_tables_count() const {
    return filtered_tables_count_;
  }

  int filtered_tablets_count() const {
    return filtered_tablets_count_;
  }

 protected:
  KsckCluster() : filtered_tables_count_(0), filtered_tablets_count_(0) {}
  MasterList masters_;
  TSMap tablet_servers_;
  KsckTServerStateMap ts_states_;
  std::shared_ptr<KsckTable> txn_sys_table_;
  std::vector<std::shared_ptr<KsckTable>> tables_;
  std::unique_ptr<ThreadPool> pool_;

  std::vector<std::string> table_filters_;
  std::vector<std::string> tablet_id_filters_;

  // The count of tables/tablets filtered out.
  // Used to determine whether all tables/tablets have been filtered out.
  std::atomic<int> filtered_tables_count_;
  std::atomic<int> filtered_tablets_count_;

 private:
  DISALLOW_COPY_AND_ASSIGN(KsckCluster);
};

// Externally facing class to run checks against the provided cluster.
class Ksck {
 public:
  explicit Ksck(std::shared_ptr<KsckCluster> cluster,
                std::ostream* out = nullptr);

  ~Ksck() = default;

  // Set whether ksck should verify that each of the tablet's raft configurations
  // has the same number of replicas that is specified by the tablet metadata.
  // (default: true)
  void set_check_replica_count(bool check) {
    check_replica_count_ = check;
  }

  // Setters for filtering the tables/tablets to be checked.
  //
  // Filter strings are glob-style patterns. For example, 'Foo*' matches
  // all tables whose name begins with 'Foo'.
  //
  // If tables is not empty, checks only the named tables.
  // If tablets is not empty, checks only the specified tablet IDs.
  // If both are specified, takes the intersection.
  // If both are empty (unset), all tables and tablets are checked.
  void set_table_filters(std::vector<std::string> table_names) {
    table_filters_ = std::move(table_names);
  }

  // See above.
  void set_tablet_id_filters(std::vector<std::string> tablet_ids) {
    tablet_id_filters_ = std::move(tablet_ids);
  }

  void set_print_sections(const std::vector<std::string>& sections);

  const KsckResults& results() const {
    return results_;
  }

  // Check that all masters are healthy.
  Status CheckMasterHealth();

  // Check that the masters' consensus information is consistent.
  Status CheckMasterConsensus();

  // Check for "unusual" flags on masters.
  // "Unusual" flags are ones tagged hidden, experimental, or unsafe and set
  // to a non-default value.
  // Must first call CheckMasterHealth().
  Status CheckMasterUnusualFlags();

  // Check for the consistency of flag values across all masters in the cluster.
  // Must first call CheckMasterHealth().
  Status CheckMasterDivergedFlags();

  // Verifies that it can connect to the cluster, i.e. that it can contact a
  // leader master.
  Status CheckClusterRunning();

  // Populates all the cluster table and tablet info from the master.
  // Must first call CheckClusterRunning().
  Status FetchTableAndTabletInfo();

  // Connects to all tablet servers, checks that they are alive, and fetches
  // their current status and tablet information.
  // Must first call FetchTableAndTabletInfo().
  Status FetchInfoFromTabletServers();

  // Check for "unusual" flags on tablet servers.
  // "Unusual" flags are ones tagged hidden, experimental, or unsafe and set
  // to a non-default value.
  // Must first call FetchInfoFromTabletServers().
  Status CheckTabletServerUnusualFlags();

  // Check for the consistency of flag values across all tablet servers
  // in the cluster.
  // Must first call FetchInfoFromTabletServers().
  Status CheckTabletServerDivergedFlags();

  // Check for the consistency of flag values across all tablet servers and
  // masters in the cluster.
  // Must first call CheckMasterDivergedFlags() and
  // CheckTabletServerDivergedFlags().
  Status CheckDivergedFlags();

  // Check for version inconsistencies among all servers.
  Status CheckServerVersions();

  // Verifies that all the tablets in all tables matching the filters have
  // enough replicas, and that each tablet's view of the tablet's consensus
  // matches every other tablet's and the master's.
  // Must first call FetchTableAndTabletInfo() and, if doing checks against
  // tablet servers (the default), must first call FetchInfoFromTabletServers().
  Status CheckTablesConsistency();

  // Verifies data checksums on all tablets by doing a scan of the database on each replica.
  // Must first call FetchTableAndTabletInfo().
  Status ChecksumData(const KsckChecksumOptions& opts);

  // Runs all the checks of ksck in the proper order, including checksum scans,
  // if enabled. Returns OK if and only if all checks succeed.
  Status Run();

  // Prints the results of ksck.
  Status PrintResults();

  // Performs all checks and prints the results.
  // Returns OK if and only if the ksck finds the cluster completely healthy,
  // and printing succeeds.
  Status RunAndPrintResults();

 private:
  friend class KsckTest;

  // Accumulate information about flags from a server into a FlagToServersMap
  // and, optionally, flag's tags into a FlagTagsMap. 'flags_to_server_map' must
  // not be null; 'flag_tags_map' may be null: in such case no information on
  // flag's tags is accumulated.
  static void AddFlagsToFlagMaps(const server::GetFlagsResponsePB& flags,
                                 const std::string& server_address,
                                 KsckFlagToServersMap* flags_to_servers_map,
                                 KsckFlagTagsMap* flag_tags_map = nullptr);

  // Returns true if the table is healthy, creating a health summary for it and
  // adding the table to 'table_summaries'. Only adds a summary if there are
  // tablets in the table.
  bool VerifyTable(const std::shared_ptr<KsckTable>& table,
                   std::vector<cluster_summary::TableSummary>* table_summaries);

  cluster_summary::HealthCheckResult VerifyTablet(
      const std::shared_ptr<KsckTablet>& tablet,
      int table_num_replicas);

  const std::shared_ptr<KsckCluster> cluster_;
  std::unique_ptr<ThreadPool> pool_;

  bool check_replica_count_ = true;
  std::vector<std::string> table_filters_;
  std::vector<std::string> tablet_id_filters_;

  std::ostream* const out_;

  // The output sections, could be a composite of PrintSections::Values
  int print_sections_flags_ = PrintSections::ALL_SECTIONS;

  KsckResults results_;

  DISALLOW_COPY_AND_ASSIGN(Ksck);
};

} // namespace tools
} // namespace kudu