#include <Arduino.h>
#include <SPI.h>
#include <bitset>
#include <cstddef>
#include <stdint.h>

#include "gonk/spi_flash.h"

#define PW_LOG_LEVEL PW_LOG_LEVEL_DEBUG
#define PW_LOG_MODULE_NAME "SpiFlash"

#include "pw_log/log.h"
#include "pw_result/result.h"
#include "pw_span/span.h"
#include "pw_status/status.h"

namespace gonk::spi_flash {

SpiFlash::SpiFlash(uint16_t flash_cs, int baudrate, SPIClass &flash_spi)
    : flash_cs_(flash_cs), spi_settings_(baudrate, MSBFIRST, SPI_MODE0),
      flash_spi_(flash_spi) {}

Status SpiFlash::WriteEnable() {
  flash_spi_.beginTransaction(spi_settings_);
  digitalWrite(flash_cs_, LOW);
  flash_spi_.transfer(0x06);
  digitalWrite(flash_cs_, HIGH);
  flash_spi_.endTransaction();

  return pw::OkStatus();
}

Status SpiFlash::WriteDisable() {
  flash_spi_.beginTransaction(spi_settings_);
  digitalWrite(flash_cs_, LOW);
  flash_spi_.transfer(0x04);
  digitalWrite(flash_cs_, HIGH);
  flash_spi_.endTransaction();

  return pw::OkStatus();
}

std::bitset<8> SpiFlash::StatusRegister() {
  uint8_t status_register[2] = {
      0x05, // Read Status Register
      0,    // Result value
  };

  flash_spi_.beginTransaction(spi_settings_);
  digitalWrite(flash_cs_, LOW);
  flash_spi_.transfer(status_register, 2);
  digitalWrite(flash_cs_, HIGH);
  flash_spi_.endTransaction();

  return (std::bitset<8>)status_register[1];
}

bool SpiFlash::IsBusy() {
  // 1 in the first bit of the status register indicates the device is busy.
  return StatusRegister()[0];
}

bool SpiFlash::WritingIsEnabled() {
  // 1 in the second bit of the status register indicates the device is busy.
  return StatusRegister()[1];
}

Status SpiFlash::Erase() {
  if (!WritingIsEnabled()) {
    return pw::Status::FailedPrecondition();
  }

  PW_LOG_DEBUG("Start Flash Erase");
  // Begin chip erase.
  flash_spi_.beginTransaction(spi_settings_);
  digitalWrite(flash_cs_, LOW);
  flash_spi_.transfer(0x60);
  digitalWrite(flash_cs_, HIGH);
  flash_spi_.endTransaction();

  int wait_time = 100; // Wait at most 1 second.
  while (IsBusy() && wait_time > 0) {
    PW_LOG_DEBUG("Device is busy");
    delay(10);
    wait_time--;
  }

  if (wait_time < 0) {
    return pw::Status::DeadlineExceeded();
  }

  return pw::OkStatus();
}

pw::Result<SpiFlash::DeviceId> SpiFlash::GetDeviceIds() {
  uint8_t manufacturer_id[4] = {
      0x9f,    // Read Manufacturer and Device ID Command
      0, 0, 0, // Result values
  };

  // Read the SPI Flash JEDEC ID.
  flash_spi_.beginTransaction(spi_settings_);
  digitalWrite(flash_cs_, LOW);
  flash_spi_.transfer(manufacturer_id, 4);
  digitalWrite(flash_cs_, HIGH);
  flash_spi_.endTransaction();

  PW_LOG_DEBUG("JEDEC ID: %x %x %x", manufacturer_id[1], manufacturer_id[2],
               manufacturer_id[3]);

  // Check for expected values
  if (manufacturer_id[1] == 0xff && manufacturer_id[2] == 0xff &&
      manufacturer_id[3] == 0xff) {
    PW_LOG_ERROR("Error: Unexpected device ID values. SPI bus may not be "
                 "released by the FPGA.");
    return pw::Status::NotFound();
  }

  return SpiFlash::DeviceId{
      .manufacturer_id = manufacturer_id[1],
      .family_code = manufacturer_id[2],
      .product_version = manufacturer_id[3],
  };
}

void SpiFlash::ResumeFromPowerDown() {
  // Resume from Deep Power-Down
  flash_spi_.beginTransaction(spi_settings_);
  digitalWrite(flash_cs_, LOW);
  flash_spi_.transfer(0xAB);
  digitalWrite(flash_cs_, HIGH);
  flash_spi_.endTransaction();
}

void SpiFlash::StartFlashSpi() {
  flash_spi_.begin();
  delay(10);
  ResumeFromPowerDown();
}

void SpiFlash::StopFlashSpi() { flash_spi_.end(); }

} // namespace gonk::spi_flash