#include "tlbmc/hal/fru_scanner_i2c.h"

#include <endian.h>
#include <fcntl.h>
#include <linux/i2c-dev.h>
#include <linux/i2c.h>
#include <sys/ioctl.h>
#include <unistd.h>

#include <array>
#include <cerrno>
#include <charconv>
#include <cstddef>
#include <cstdint>
#include <cstdio>
#include <cstring>
#include <filesystem>
#include <iomanip>
#include <ios>
#include <memory>
#include <optional>
#include <sstream>
#include <string>
#include <system_error>
#include <utility>
#include <vector>

#include "absl/container/flat_hash_map.h"
#include "absl/container/flat_hash_set.h"
#include "absl/log/log.h"
#include "absl/memory/memory.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/numbers.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "absl/strings/substitute.h"
#include "absl/time/time.h"
#include "absl/types/span.h"
#include "apifs/apifs.h"
#include "io/smbus/smbus.h"
#include "g3/macros.h"
#include "tlbmc/hal/fru_scanner.h"
#include "tlbmc/utils/fru_reader.h"
#include "tlbmc/utils/fru_utils.h"
#include "re2/re2.h"

namespace milotic_tlbmc {

using SmbusLocation = ecclesia::SmbusLocation;

namespace {

namespace fs = std::filesystem;

// Regex to extract the address from i2c file path.
// Example: /sys/bus/i2c/devices/i2c-12/12-0050 -> 0050
constexpr LazyRE2 kAddressRegex = {".+\\d+-([0-9abcdef]+$)"};

absl::StatusOr<absl::flat_hash_set<int>> GetI2cBusList(
    absl::string_view dir_path) {
  absl::flat_hash_set<int> bus_list;
  ecclesia::ApifsDirectory apifs((std::string(dir_path)));
  ECCLESIA_ASSIGN_OR_RETURN(auto entries, apifs.ListEntryPaths());
  for (const auto& entry : entries) {
    auto iter = entry.find("i2c-");
    if (iter == std::string::npos) {
      continue;
    }
    std::string bus_string = entry.substr(iter + 4);
    int bus = 0;
    if (!absl::SimpleAtoi(bus_string, &bus)) {
      return absl::InternalError("Failed to convert bus string to int");
    }
    bus_list.insert(bus);
  }
  if (bus_list.empty()) {
    return absl::InternalError("No i2c buses found in directory");
  }
  return bus_list;
}

int64_t ReadFromEeprom(int fd, off_t offset, size_t len, uint8_t* buf) {
  auto result = lseek(fd, offset, SEEK_SET);
  if (result < 0) {
    LOG(ERROR) << "Failed to seek to offset " << offset << " in eeprom file";
    return -1;
  };
  return read(fd, buf, len);
}

std::string GetEepromPath(size_t bus, size_t address,
                          absl::string_view root_dir) {
  std::stringstream output;
  output << root_dir << "/sys/bus/i2c/devices/" << bus << "-" << std::right
         << std::setfill('0') << std::setw(4) << std::hex << address
         << "/eeprom";
  return output.str();
}

std::vector<uint8_t> ProcessEeprom(int bus, size_t address,
                                   absl::string_view root_dir) {
  auto path = GetEepromPath(bus, address, root_dir);
  int file = open(path.c_str(), O_RDONLY);
  if (file < 0) {
    LOG(ERROR) << "Failed to open eeprom file: " << path << " "
               << strerror(errno);
    return {};
  }

  std::string error_message = "eeprom at " + std::to_string(bus) + " address " +
                              std::to_string(address);
  FRUReader reader(
      [file](off_t offset, size_t length, uint8_t* outbuf) -> int64_t {
        absl::StatusOr<int64_t> bytes_read =
            ReadFromEeprom(file, offset, length, outbuf);
        if (!bytes_read.ok()) {
          return -1;
        }
        return bytes_read.value();
      });
  std::pair<std::vector<uint8_t>, bool> pair =
      readFRUContents(reader, error_message);
  close(file);
  if (!pair.second) {
    LOG(ERROR) << "Failed to read eeprom at " << bus << " address " << address;
    return {};
  }
  return pair.first;
}

int GetRootBus(size_t bus, absl::string_view root_dir) {
  auto ec = std::error_code();
  auto path = std::filesystem::read_symlink(
      std::filesystem::path(absl::StrCat(root_dir, "/sys/bus/i2c/devices/i2c-",
                                         bus, "/mux_device")),
      ec);
  if (ec) {
    return -1;
  }

  std::string filename = path.filename();
  auto bus_iter = filename.find('-');
  if (bus_iter == std::string::npos) {
    return -1;
  }

  int root_bus = 0;
  return absl::SimpleAtoi(filename.substr(0, bus_iter), &root_bus) ? root_bus
                                                                   : -1;
}

}  // namespace

std::pair<std::vector<uint8_t>, bool> FruScannerI2c::ReadFruContentsFromI2c(
    int i2c_bus, int address, absl::Duration delay_between_reads) {
  std::optional<SmbusLocation> location =
      SmbusLocation::TryMake(i2c_bus, address);
  if (!location.has_value()) {
    return {{}, false};
  }

  uint8_t data;
  // Probe to see if this address is valid
  if (!smbus_access_->ReceiveByte(*location, &data).ok()) {
    return {{}, false};
  }

  std::string error_message =
      absl::Substitute("i2c scan at $0 address $1", i2c_bus, address);
  FRUReader reader(
      [this, location, delay_between_reads](off_t offset, size_t length,
                                            uint8_t* outbuf) -> int64_t {
        DLOG(INFO) << "Reading from i2c at offset " << offset << " length "
                   << length;
        absl::Span<uint8_t> outbuf_span(outbuf, length);
        for (int i = 0; i < length; ++i) {
          // TODO(b/385740524): Add 16bit scanning logic
          absl::Status smbus_status = smbus_access_->Read8(
              *location, static_cast<int>(offset + i), &outbuf_span[i]);
          if (!smbus_status.ok()) {
            LOG(INFO) << "Smbus read at offset " << offset
                      << " failed: " << smbus_status;
            return static_cast<int64_t>(outbuf_span[i]);
          }

          if (delay_between_reads > absl::ZeroDuration()) {
            clock_->Sleep(delay_between_reads);
          }
        }
        return static_cast<int64_t>(length);
      });

  return readFRUContents(reader, error_message);
}

absl::StatusOr<std::unique_ptr<I2cFruInfo>> FruScannerI2c::GetI2cFruInfoFromBus(
    int bus, int address, absl::Duration delay_between_reads) {
  auto fru_info = std::make_unique<I2cFruInfo>();
  fru_info->bus = bus;
  fru_info->address = address;
  std::vector<uint8_t> device = ProcessEeprom(bus, address, root_dir_);
  // If we get a device, then we don't need to manually scan.
  if (!device.empty()) {
    fru_info->data = std::move(device);
    return fru_info;
  }
  return ScanDirectI2cFru(bus, address, delay_between_reads);
}

absl::StatusOr<std::unique_ptr<I2cFruInfo>> FruScannerI2c::ScanDirectI2cFru(
    int bus, int address, absl::Duration delay_between_reads) {
  std::pair<std::vector<uint8_t>, bool> pair =
      ReadFruContentsFromI2c(bus, address, delay_between_reads);
  if (!pair.second) {
    return absl::InternalError(
        absl::StrCat("Failed to read eeprom at ", bus, " address ", address));
  }

  auto fru_info = std::make_unique<I2cFruInfo>();
  fru_info->bus = bus;
  fru_info->address = address;
  fru_info->data = pair.first;
  return fru_info;
}

absl::flat_hash_set<size_t> FruScannerI2c::FindI2cEeproms(
    int i2cBus, const std::shared_ptr<DeviceMap>& devices,
    absl::string_view root_dir,
    const absl::flat_hash_set<size_t>& root_blocked_addresses,
    absl::flat_hash_set<size_t>& root_found_addresses) {
  absl::flat_hash_set<size_t> found_list;
  std::string path =
      absl::StrCat(root_dir, "/sys/bus/i2c/devices/i2c-", i2cBus);

  // For each file listed under the i2c device
  // NOTE: This should be faster than just checking for each possible address
  // path.
  auto ec = std::error_code();
  for (const auto& p : fs::directory_iterator(path, ec)) {
    if (ec) {
      LOG(ERROR) << "directory_iterator err " << ec.message();
      break;
    }
    const std::string node = p.path().string();
    std::string address_string;
    if (!RE2::PartialMatch(node, *kAddressRegex, &address_string)) {
      continue;
    }
    absl::string_view address_view(address_string);
    size_t address = 0;
    std::from_chars(address_view.begin(), address_view.end(), address, 16);

    if (root_blocked_addresses.contains(address) ||
        root_found_addresses.contains(address)) {
      continue;
    }

    const std::string eeprom = node + "/eeprom";

    if (!fs::exists(eeprom, ec)) {
      continue;
    }

    // There is an eeprom file at this address, it may have invalid
    // contents, but we found it.
    found_list.insert(address);

    std::vector<uint8_t> device = ProcessEeprom(i2cBus, address, root_dir);
    if (!device.empty()) {
      devices->emplace(address, device);
    }
  }

  // Scan every bus from 0x50 to 0x57 only
  // SoT:
  // https://source.corp.google.com/h/gbmc/codesearch/+/main:meta-gbmc-staging/recipes-phosphor/configuration/entity-manager/0002-fru_device-limit-the-fru-scan-range-to-0x50-0x57.patch?q=fru%20device%20limit%20the%20fru%20scan%20range
  for (int address = kI2cFruScanStart; address <= kI2cFruScanEnd; address++) {
    if (found_list.contains(address) ||
        root_blocked_addresses.contains(address) ||
        root_found_addresses.contains(address)) {
      continue;
    }

    std::pair<std::vector<uint8_t>, bool> pair =
        ReadFruContentsFromI2c(i2cBus, address, absl::ZeroDuration());

    if (!pair.second) {
      LOG(INFO) << "Failed to read eeprom at " << i2cBus << " address "
                << address;
      continue;
    }

    found_list.insert(address);
    devices->emplace(address, pair.first);
  }

  return found_list;
}

void FruScannerI2c::FindI2CDevices(const absl::flat_hash_set<int>& i2c_buses,
                                   BusMap& bus_map) {
  const absl::flat_hash_map<size_t, std::optional<absl::flat_hash_set<size_t>>>&
      bus_block_list = scan_context_.bus_block_list;
  DLOG(INFO) << "Scanning i2c buses! ";
  for (const auto& bus : i2c_buses) {
    // Will need to store all addresses in blocklist for a bus.
    // If this is an extended bus, need to skip addresses found at root bus and
    // inherit the root bus block list.
    // blocked_addresses: addresses that are blocked on this bus or root bus.
    // root_found_addresses: addresses that are found at root bus.
    absl::flat_hash_set<size_t> blocked_addresses;
    absl::flat_hash_set<size_t> root_found_addresses;

    if (scan_context_.fru_addresses_list.contains(bus)) {
      // Bus is a root bus that has already been scanned, continue
      continue;
    }
    auto bus_iter = bus_block_list.find(bus);
    if (bus_iter != bus_block_list.end()) {
      if (bus_iter->second == std::nullopt) {
        DLOG(INFO) << "Skipping blocked bus " << bus;
        continue;  // Skip blocked busses.
      }
      for (size_t address : *(bus_iter->second)) {
        DLOG(INFO) << "Adding blocked address " << address << " on bus " << bus;
        blocked_addresses.insert(address);
      }
    }

    int root_bus = GetRootBus(bus, root_dir_);
    if (root_bus >= 0) {
      auto root_bus_iter = bus_block_list.find(root_bus);
      if (root_bus_iter != bus_block_list.end()) {
        if (root_bus_iter->second != std::nullopt) {
          for (auto& root_address : *(root_bus_iter->second)) {
            DLOG(INFO) << "Skipping root address " << root_address;
            blocked_addresses.insert(root_address);
          }
        }
      }

      if (auto find_root_bus = scan_context_.fru_addresses_list.find(root_bus);
          find_root_bus != scan_context_.fru_addresses_list.end()) {
        root_found_addresses = find_root_bus->second;
      }
    }

    auto device = std::make_shared<DeviceMap>();
    DLOG(INFO) << "Scanning bus " << bus;
    if (root_bus >= 0 && scan_context_.fru_addresses_list.find(root_bus) ==
                             scan_context_.fru_addresses_list.end()) {
      // Root bus has not been scanned, scan it before extended bus.
      DLOG(INFO) << "Scanning root bus " << root_bus;
      auto root_device = std::make_shared<DeviceMap>();

      absl::flat_hash_set<size_t> root_found_list =
          FindI2cEeproms(root_bus, root_device, root_dir_,
                         blocked_addresses, root_found_addresses);
      bus_map.emplace(root_bus, std::move(root_device));
      scan_context_.fru_addresses_list[root_bus].insert(root_found_list.begin(),
                                                        root_found_list.end());
      root_found_addresses = root_found_list;
    }
    absl::flat_hash_set<size_t> found_list = FindI2cEeproms(
        bus, device, root_dir_, blocked_addresses, root_found_addresses);
    bus_map.emplace(bus, std::move(device));
    scan_context_.fru_addresses_list[bus].insert(found_list.begin(),
                                                 found_list.end());
  }
}

absl::StatusOr<std::vector<std::unique_ptr<I2cFruInfo>>>
FruScannerI2c::ScanAllI2cFrus() {
  ECCLESIA_ASSIGN_OR_RETURN(auto bus_list, GetI2cBusList(root_dir_ + "/dev"));
  BusMap bus_map;
  FindI2CDevices(bus_list, bus_map);
  std::vector<std::unique_ptr<I2cFruInfo>> i2c_frus;
  for (const auto& [bus, device_map] : bus_map) {
    for (const auto& [address, data] : *device_map) {
      auto fru_info = std::make_unique<I2cFruInfo>();
      fru_info->bus = bus;
      fru_info->address = address;
      fru_info->data = data;
      i2c_frus.push_back(std::move(fru_info));
    }
  }
  return i2c_frus;
}

std::unique_ptr<FruScannerI2c> FruScannerI2c::Create(const Options& options) {
  return absl::WrapUnique(new FruScannerI2c(
      options.root_dir, options.i2c_dev_dir, options.bus_block_list));
}

}  // namespace milotic_tlbmc