#include "tlbmc/collector/fru_collector.h"

#include <cstddef>
#include <filesystem>
#include <fstream>
#include <ios>
#include <memory>
#include <optional>
#include <string>
#include <system_error>
#include <utility>
#include <vector>

#include "absl/base/thread_annotations.h"
#include "absl/container/flat_hash_map.h"
#include "absl/container/flat_hash_set.h"
#include "absl/functional/any_invocable.h"
#include "absl/log/log.h"
#include "absl/memory/memory.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "absl/synchronization/mutex.h"
#include "g3/macros.h"
#include "nlohmann/json_fwd.hpp"
#include "json_utils.h"
#include "ad_hoc_fru_config.pb.h"
#include "tlbmc/configs/entity_config.h"
#include "hal_common_config.pb.h"
#include "tlbmc/hal/fru_scanner.h"
#include "fru.pb.h"
#include "resource.pb.h"
#include "tlbmc/scheduler/scheduler.h"
#include "tlbmc/time/time.h"
#include "tlbmc/utils/fram_utils.h"
#include "tlbmc/utils/fru_utils.h"
#include "google/protobuf/json/json.h"
#include "google/protobuf/util/json_util.h"

namespace milotic_tlbmc {

using AdHocFruScanContext = FruCollector::AdHocFruScanContext;

// These ad-hoc scanning configs will be in
// json["ProbeV2"]["AdHocFruConfig"]
std::vector<AdHocFruConfig> FruCollector::Options::ParseAdHocFruConfigs(
    const std::vector<nlohmann::json>& config_list) {
  std::vector<AdHocFruConfig> ad_hoc_fru_configs;
  for (const auto& config : config_list) {
    const nlohmann::json* tlbmc_config =
        GetValueAsJson(config, kTlbmcConfigKey);
    if (tlbmc_config == nullptr) {
      continue;
    }
    const nlohmann::json* ad_hoc_fru_scanning_config =
        GetValueAsJson(*tlbmc_config, kAdHocFruConfigKey);
    if (ad_hoc_fru_scanning_config == nullptr) {
      continue;
    }
    AdHocFruConfig ad_hoc_fru_config_proto;
    ::google::protobuf::json::ParseOptions opts;
    if (!::google::protobuf::json::JsonStringToMessage(ad_hoc_fru_scanning_config->dump(),
                                             &ad_hoc_fru_config_proto, opts)
             .ok()) {
      const std::string* config_name =
          milotic::authz::GetValueAsString(config, "Name");
      LOG(WARNING) << absl::StrCat(
          "Failed to convert ad-hoc FRU scanning config to proto for ",
          (config_name == nullptr ? "unnamed config"
                                  : std::string(*config_name)));
      continue;
    }
    ad_hoc_fru_configs.push_back(ad_hoc_fru_config_proto);
  }
  return ad_hoc_fru_configs;
}

nlohmann::json FruCollector::ToJson() const {
  nlohmann::json response;
  std::string json_string;
  ::google::protobuf::util::JsonPrintOptions opts;
  opts.preserve_proto_field_names = true;
  absl::MutexLock lock(&fru_table_mutex_);
  if (!::google::protobuf::json::MessageToJsonString(fru_table_, &json_string, opts)
           .ok()) {
    LOG(ERROR) << "Failed to convert FRU table to JSON";
    return response;
  }
  return nlohmann::json::parse(json_string, nullptr, false);
}

// Returns the next scan info if there is a next scan period to schedule.
// Returns nullopt if there is no next scan period to schedule.
std::optional<AdHocFruScanContext> FruCollector::GetNextScanInfo(
    const AdHocFruConfig& ad_hoc_fru_config,
    const AdHocFruScanContext& scan_context) {
  if (ad_hoc_fru_config.periodic_scan_config_size() <=
      scan_context.cur_periodic_scan_config_index) {
    return std::nullopt;
  }

  // 0 indexed so if we are at the last scan period, we need to move onto next
  // period.
  if (ad_hoc_fru_config
          .periodic_scan_config(
              static_cast<int>(scan_context.cur_periodic_scan_config_index))
          .scan_count() ==
      scan_context.cur_periodic_scan_config_scan_count + 1) {
    // If we are at the last scan period, then no more scans.
    if (scan_context.cur_periodic_scan_config_index + 1 ==
        ad_hoc_fru_config.periodic_scan_config_size()) {
      return std::nullopt;
    }
    return AdHocFruScanContext{
        .cur_periodic_scan_config_index =
            scan_context.cur_periodic_scan_config_index + 1,
        .cur_periodic_scan_config_scan_count = 0,
        .found = scan_context.found};
  }

  return AdHocFruScanContext{
      .cur_periodic_scan_config_index =
          scan_context.cur_periodic_scan_config_index,
      .cur_periodic_scan_config_scan_count =
          scan_context.cur_periodic_scan_config_scan_count + 1,
      .found = scan_context.found};
}

void FruCollector::RescheduleAdHocFruScan(
    const AdHocFruConfig& ad_hoc_fru_config, size_t ad_hoc_fru_scan_context_id,
    const AdHocFruScanContext& scan_context) {
  // Schedule task with new period
  task_scheduler_->ScheduleOneShotAsync(
      [this, ad_hoc_fru_config,
       ad_hoc_fru_scan_context_id](absl::AnyInvocable<void()> on_done) {
        AttemptAdHocFruScan(ad_hoc_fru_config, ad_hoc_fru_scan_context_id);
        on_done();
      },
      DecodeGoogleApiProto(ad_hoc_fru_config
                               .periodic_scan_config(static_cast<int>(
                                   scan_context.cur_periodic_scan_config_index))
                               .scan_interval()));
}

void FruCollector::AttemptAdHocFruScan(const AdHocFruConfig& ad_hoc_fru_config,
                                       size_t ad_hoc_fru_scan_context_id) {
  AdHocFruScanContext scan_context =
      GetAdHocFruScanContext(ad_hoc_fru_scan_context_id);

  LOG(INFO) << "Attempting ad-hoc FRU scan for " << ad_hoc_fru_config.name();

  std::shared_ptr<EntityConfig> entity_config = GetEntityConfig();
  if (entity_config == nullptr) {
    LOG(WARNING) << "Entity config is not yet set. Will rescan later for "
                 << ad_hoc_fru_config.name();
    RescheduleAdHocFruScan(ad_hoc_fru_config, ad_hoc_fru_scan_context_id,
                           scan_context);
    return;
  }

  AdHocScannerType scanner_type = ad_hoc_fru_config.scanner_type();

  LOG(INFO) << "Ad-hoc config: " << ad_hoc_fru_config.name()
            << " requests scanner_type: "
            << AdHocScannerType_Name(scanner_type);
  auto it = options_.fru_scanners.find(scanner_type);
  if (it == options_.fru_scanners.end()) {
    LOG(ERROR) << "FruScanner not found for type "
               << AdHocScannerType_Name(scanner_type);
    return;
  }
  FruScanner* fru_scanner = it->second;
  if (fru_scanner == nullptr) {
    LOG(ERROR) << "FruScanner found for type "
               << AdHocScannerType_Name(scanner_type) << " but it is NULL";
    return;
  }

  // Try to scan FRU
  absl::StatusOr<std::unique_ptr<I2cFruInfo>> ad_hoc_fru_info =
      fru_scanner->GetI2cFruInfoFromBus(
          static_cast<int>(ad_hoc_fru_config.hal_common_config().bus()),
          static_cast<int>(ad_hoc_fru_config.hal_common_config().address()),
          DecodeGoogleApiProto(ad_hoc_fru_config.delay_between_reads()));

  // If scan failed, next scan is already scheduled, so update scan
  // context and reschedule with different period if needed
  if (!ad_hoc_fru_info.ok()) {
    LOG(WARNING) << "Failed to scan ad-hoc FRU: " << ad_hoc_fru_config.name()
                 << " with status: " << ad_hoc_fru_info.status();
    ScanNextAdHocIfNeeded(ad_hoc_fru_config, ad_hoc_fru_scan_context_id,
                          scan_context);
    return;
  }
  LOG(INFO) << "Successfully scanned ad-hoc FRU: " << ad_hoc_fru_config.name();

  // If scan succeeded, update FRU table
  absl::StatusOr<RawFru> fru =
      CreateRawFruFromScanner(scanner_type, *ad_hoc_fru_info);
  if (!fru.ok()) {
    LOG(WARNING) << "Failed to create raw FRU from data: " << fru.status()
                 << " for FRU: " << ad_hoc_fru_config.name();
    ScanNextAdHocIfNeeded(ad_hoc_fru_config, ad_hoc_fru_scan_context_id,
                          scan_context);
    return;
  }
  AddFru(*fru);
  scan_context.found = true;
  UpdateAdHocFruScanContext(ad_hoc_fru_scan_context_id, scan_context);

  LOG(WARNING) << "Updated FRU table for FRU key: " << fru->key()
               << " for ad-hoc FRU: " << ad_hoc_fru_config.name();

  entity_config->UpdateFruAndTopology(GetCopyOfCurrentScannedFrus(), *fru);
}

void FruCollector::ScanNextAdHocIfNeeded(
    const AdHocFruConfig& ad_hoc_fru_config, size_t ad_hoc_fru_scan_context_id,
    AdHocFruScanContext& scan_context) {
  std::optional<AdHocFruScanContext> next_scan_info =
      GetNextScanInfo(ad_hoc_fru_config, scan_context);

  // If there is no next scan info, then we are done with scanning.
  if (!next_scan_info.has_value()) {
    return;
  }

  // Sanity check periodic scan config size, fail safe.
  if (ad_hoc_fru_config.periodic_scan_config_size() <
      next_scan_info->cur_periodic_scan_config_index) {
    LOG(ERROR) << "Error in ad-hoc FRU scanning logic. No longer scanning "
               << ad_hoc_fru_config.name();
    return;
  }
  scan_context.cur_periodic_scan_config_index =
      next_scan_info->cur_periodic_scan_config_index;
  scan_context.cur_periodic_scan_config_scan_count =
      next_scan_info->cur_periodic_scan_config_scan_count;
  UpdateAdHocFruScanContext(ad_hoc_fru_scan_context_id, scan_context);
  RescheduleAdHocFruScan(ad_hoc_fru_config, ad_hoc_fru_scan_context_id,
                         scan_context);
  LOG(INFO) << "Rescheduled ad-hoc FRU scan for " << ad_hoc_fru_config.name();
}

void FruCollector::ScheduleAdHocFruScan(
    const AdHocFruConfig& ad_hoc_fru_config) {
  // Index to add to ad_hoc_fru_scan_contexts_
  size_t ad_hoc_fru_scan_context_id = GetAdHocFruScanContextsSize();
  AddAdHocFruScanContext(
      AdHocFruScanContext{.cur_periodic_scan_config_index = 0,
                          .cur_periodic_scan_config_scan_count = 0,
                          .found = false});

  // Scheduler automatically reruns task at defined periods
  task_scheduler_->ScheduleOneShotAsync(
      [this, ad_hoc_fru_config,
       ad_hoc_fru_scan_context_id](absl::AnyInvocable<void()> on_done) {
        AttemptAdHocFruScan(ad_hoc_fru_config, ad_hoc_fru_scan_context_id);
        on_done();
      },
      DecodeGoogleApiProto(
          ad_hoc_fru_config.periodic_scan_config(0).scan_interval()));
}

absl::StatusOr<RawFru> FruCollector::CreateRawFruFromScanner(
    AdHocScannerType scanner_type,
    const std::unique_ptr<I2cFruInfo>& i2c_fru_info) {
  switch (scanner_type) {
    case SCANNER_TYPE_DEFAULT_IPMI_I2C:
      return CreateRawFruFromIpmiFruData(i2c_fru_info);
    case SCANNER_TYPE_CAVIUM_HSM_FRAM:
      return CreateRawFruFromFramData(i2c_fru_info);
    default:
      LOG(ERROR) << "Unexpected value for AdHocScannerType: "
                 << static_cast<int>(scanner_type);
      return absl::InvalidArgumentError(
          absl::StrCat("Unexpected scanner type: ", scanner_type));
  }
}

absl::StatusOr<RawFru> FruCollector::CreateRawFruFromIpmiFruData(
    const std::unique_ptr<I2cFruInfo>& i2c_fru) {
  absl::flat_hash_map<std::string, std::string> fru_data_fields;

  // Get Fru data fields.
  size_t unknown_bus_object_count = 0;

  std::optional<std::string> product_name =
      getProductName(i2c_fru->data, fru_data_fields, i2c_fru->bus,
                     i2c_fru->address, unknown_bus_object_count);

  if (!product_name.has_value()) {
    return absl::InternalError(absl::StrCat(
        "Failed to get product name for FRU at bus: ", i2c_fru->bus,
        " address: ", i2c_fru->address));
  }

  // Create FRU resource.
  RawFru fru;

  // Assuming we can have only one FRU per I2C bus/address pair
  fru.set_key(absl::StrCat(i2c_fru->bus, ":", i2c_fru->address));

  // Populate I2C info.
  fru.mutable_i2c_info()->set_bus(i2c_fru->bus);
  fru.mutable_i2c_info()->set_address(i2c_fru->address);

  // Populate FRU data.
  FruData* fru_data = fru.mutable_data();
  for (const auto& fru_data_field : fru_data_fields) {
    (*fru_data->mutable_fields())[fru_data_field.first] = fru_data_field.second;
    if (fru_data_field.first == "BOARD_MANUFACTURER") {
      fru_data->mutable_fru_info()->set_board_manufacturer(
          fru_data_field.second);
    } else if (fru_data_field.first == "BOARD_PRODUCT_NAME") {
      fru_data->mutable_fru_info()->set_board_product_name(
          fru_data_field.second);
    } else if (fru_data_field.first == "BOARD_SERIAL_NUMBER") {
      fru_data->mutable_fru_info()->set_board_serial_number(
          fru_data_field.second);
    } else if (fru_data_field.first == "BOARD_PART_NUMBER") {
      fru_data->mutable_fru_info()->set_board_part_number(
          fru_data_field.second);
    } else if (fru_data_field.first == "PRODUCT_MANUFACTURER") {
      fru_data->mutable_fru_info()->set_product_manufacturer(
          fru_data_field.second);
    } else if (fru_data_field.first == "PRODUCT_PRODUCT_NAME") {
      fru_data->mutable_fru_info()->set_product_product_name(
          fru_data_field.second);
    } else if (fru_data_field.first == "PRODUCT_SERIAL_NUMBER") {
      fru_data->mutable_fru_info()->set_product_serial_number(
          fru_data_field.second);
    } else if (fru_data_field.first == "PRODUCT_PART_NUMBER") {
      fru_data->mutable_fru_info()->set_product_part_number(
          fru_data_field.second);
    }
  }

  return fru;
}

std::optional<RawFruTable> FruCollector::DetectsAndReadsCachedFru(
    const std::string& cached_fru_table_path) {
  std::ifstream file(cached_fru_table_path, std::ios::binary);
  RawFruTable fru_table;
  if (!fru_table.ParseFromIstream(&file)) {
    LOG(WARNING) << "Error parsing cached FRU table file at path: "
                 << cached_fru_table_path;
    LOG(WARNING) << "Check the file exists and is in the correct format. Treat "
                    "this as no cached FRU table";
    return std::nullopt;
  }
  LOG(WARNING) << "Successfully read cached FRU table from file at path: "
               << cached_fru_table_path;
  return fru_table;
}

void FruCollector::WriteCachedFru(const RawFruTable& fru_table,
                                  const std::string& cached_fru_table_path) {
  std::filesystem::path path(cached_fru_table_path);
  std::error_code error_code;
  std::filesystem::path parent_path = path.parent_path();
  if (!std::filesystem::exists(parent_path)) {
    std::filesystem::create_directories(path.parent_path(), error_code);
    if (error_code) {
      LOG(WARNING) << "Error creating directories for cached FRU table file at "
                      "path: "
                   << path.parent_path();
      return;
    }
  }

  std::ofstream file(cached_fru_table_path, std::ios::binary);
  if (!fru_table.SerializeToOstream(&file)) {
    LOG(WARNING) << "Error serializing cached FRU table to file at path: "
                 << cached_fru_table_path;
    return;
  }
  LOG(WARNING) << "Successfully wrote cached FRU table to file at path: "
               << cached_fru_table_path;
}

absl::StatusOr<std::unique_ptr<FruCollector>> FruCollector::Create(
    Options options) {
  std::optional<RawFruTable> cached_fru_table =
      DetectsAndReadsCachedFru(options.cached_fru_table_path);
  // If cached FRU table is present, use it, and will not scan any fixed FRUs,
  // nor write cached FRU to file.
  if (cached_fru_table.has_value()) {
    return absl::WrapUnique(
        new FruCollector(std::move(options), std::move(*cached_fru_table),
                         std::make_shared<TaskScheduler>()));
  }

  if (options.fru_scanners.empty()) {
    return absl::InvalidArgumentError("FruScanners is empty");
  }

  // Validate AdHocFruConfigs against available scanners
  for (const auto& ad_hoc_config : options.ad_hoc_fru_scanning_configs) {
    AdHocScannerType scanner_type = ad_hoc_config.scanner_type();
    if (!options.fru_scanners.contains(scanner_type)) {
      return absl::InvalidArgumentError(
          absl::StrCat("Scanner type ", AdHocScannerType_Name(scanner_type),
                       " required by ad-hoc config '", ad_hoc_config.name(),
                       "' not found in provided fru_scanners map."));
    }
  }

  RawFruTable fru_table;
  for (const auto& [scanner_type, scanner] : options.fru_scanners) {
    if (scanner == nullptr) {
      return absl::InvalidArgumentError(
          absl::StrCat("Scanner for type ", AdHocScannerType_Name(scanner_type),
                       " is null."));
    }
    LOG(INFO) << "Scanning FRUs with scanner type: "
              << AdHocScannerType_Name(scanner_type);

    ECCLESIA_ASSIGN_OR_RETURN(std::vector<std::unique_ptr<I2cFruInfo>> frus,
                              scanner->ScanAllI2cFrus());

    for (const auto& fru_info : frus) {
      ECCLESIA_ASSIGN_OR_RETURN(
          RawFru fru, CreateRawFruFromScanner(scanner_type, fru_info));
      fru_table.mutable_key_to_raw_fru()->insert({fru.key(), fru});
      LOG(WARNING) << "Updated FRU table for FRU key: " << fru.key()
                   << " from scanner " << AdHocScannerType_Name(scanner_type);
    }
  }

  WriteCachedFru(fru_table, options.cached_fru_table_path);

  return absl::WrapUnique(new FruCollector(std::move(options),
                                           std::move(fru_table),
                                           std::make_shared<TaskScheduler>()));
}

void FruCollector::SetUpAdHocFruScanning() {
  if (options_.fru_scanners.empty()) {
    LOG(WARNING) << "FruScanners is empty";
    return;
  }

  LOG(INFO) << "Number of available fru_scanners: "
            << options_.fru_scanners.size();
  for (const auto& ad_hoc_fru_config : options_.ad_hoc_fru_scanning_configs) {
    std::string fru_key =
        absl::StrCat(ad_hoc_fru_config.hal_common_config().bus(), ":",
                     ad_hoc_fru_config.hal_common_config().address());
    // If the frutable already has the ad-hoc fru, skip it.
    if (ContainsFru(fru_key)) {
      LOG(WARNING) << "FRU key " << fru_key
                   << " already exists, skipping ad-hoc setup for "
                   << ad_hoc_fru_config.name();
      continue;
    }

    // If Ad-hoc FRU config has no periodic scan config, skip it.
    if (ad_hoc_fru_config.periodic_scan_config_size() == 0) {
      LOG(WARNING) << "No periodic_scan_config for " << ad_hoc_fru_config.name()
                   << ", skipping.";
      continue;
    }

    LOG(INFO) << "Scheduling AdHocFruScan for " << ad_hoc_fru_config.name();

    ScheduleAdHocFruScan(ad_hoc_fru_config);
  }
}

FruCollector::FruCollector(Options options, RawFruTable fru_table,
                           std::shared_ptr<TaskScheduler> task_scheduler)
    : task_scheduler_(std::move(task_scheduler)),
      options_(std::move(options)),
      fru_table_(std::move(fru_table)) {}

std::unique_ptr<FruCollector> EmptyFruCollector::Create() {
  return std::make_unique<EmptyFruCollector>();
}

RawFruTable EmptyFruCollector::GetCopyOfCurrentScannedFrus() const
    ABSL_LOCKS_EXCLUDED(fru_table_mutex_) {
  LOG(WARNING) << "EmptyFruCollector::GetCopyOfCurrentScannedFrus is called. "
                  "This is a no-op.";
  return RawFruTable();
}

nlohmann::json EmptyFruCollector::ToJson() const {
  return nlohmann::json::parse("{\"Warning\": \"EmptyFruCollector used.\"}");
}

void EmptyFruCollector::SetEntityConfig(
    const std::shared_ptr<EntityConfig>& entity_config)
    ABSL_LOCKS_EXCLUDED(entity_config_mutex_) {
  LOG(WARNING) << "EmptyFruCollector::SetEntityConfig is called. This is a "
                  "no-op.";
}

std::shared_ptr<EntityConfig> EmptyFruCollector::GetEntityConfig() const {
  return nullptr;
}

nlohmann::json EmptyFruCollector::GetSchedulerStats() const {
  return nlohmann::json::parse("{\"Warning\": \"EmptyFruCollector used.\"}");
}

std::size_t EmptyFruCollector::GetAllUserTaskCount() const { return 0; }

}  // namespace milotic_tlbmc