/* SPDX-FileCopyrightText: 2024 Google LLC */
/* SPDX-License-Identifier: Apache-2.0 */

#include "board/board.h"
#include "drivers/flash/flash_impl.h"
#include "drivers/flash/qspi_flash.h"
#include "drivers/flash/qspi_flash_part_definitions.h"
#include "flash_region/flash_region.h"
#include "system/passert.h"
#include "system/status_codes.h"
#include "util/math.h"

#define STM32F4_COMPATIBLE
#define STM32F7_COMPATIBLE
#include <mcu.h>

static QSPIFlashPart QSPI_FLASH_PART = {
    .instructions =
        {
            .fast_read = 0x0B,
            .fast_read_ddr = 0x0D,
            .pp = 0x02,
            .erase_sector_4k = 0x20,
            .erase_block_64k = 0xD8,
            .write_enable = 0x06,
            .write_disable = 0x04,
            .rdsr1 = 0x05,
            .rdsr2 = 0x70,
            .erase_suspend = 0x75,
            .erase_resume = 0x7A,
            .enter_low_power = 0xB9,
            .exit_low_power = 0xAB,
            .enter_quad_mode = 0x35,
            .reset_enable = 0x66,
            .reset = 0x99,
            .qspi_id = 0xAF,

            .block_lock = 0xE5,
            .block_lock_status = 0xE8,
        },
    .status_bit_masks =
        {
            .busy = 1 << 0,
            .write_enable = 1 << 1,
        },
    .flag_status_bit_masks =
        {
            .erase_suspend = 1 << 6,
        },
    .dummy_cycles =
        {
            .fast_read = 10,
            .fast_read_ddr = 8,
        },
    .block_lock =
        {
            .has_lock_data = true,
            .lock_data = 0x1,
            .locked_check = 0x1,
        },
    .reset_latency_ms = 51,
    .suspend_to_read_latency_us = 0,
    .standby_to_low_power_latency_us = 3,
    .low_power_to_standby_latency_us = 30,
    .supports_fast_read_ddr = true,
    .supports_block_lock = true,
#if BOARD_ROBERT_BB || BOARD_ROBERT_BB2 || BOARD_CUTTS_BB
    .qspi_id_value = 0x19BB20,
    .name = "MT25Q256",
#elif BOARD_ROBERT_EVT
    .qspi_id_value = 0x18BB20,
    .name = "MT25Q128",
#else
#error "Unsupported board"
#endif
};

bool flash_check_whoami(void) { return qspi_flash_check_whoami(QSPI_FLASH); }

FlashAddress flash_impl_get_sector_base_address(FlashAddress addr) {
  return (addr & SECTOR_ADDR_MASK);
}

FlashAddress flash_impl_get_subsector_base_address(FlashAddress addr) {
  return (addr & SUBSECTOR_ADDR_MASK);
}

void flash_impl_enable_write_protection(void) {}

status_t flash_impl_write_protect(FlashAddress start_sector, FlashAddress end_sector) {
  FlashAddress block_addr = start_sector;
  while (block_addr <= end_sector) {
    uint32_t block_size;
    if (WITHIN(block_addr, SECTOR_SIZE_BYTES, BOARD_NOR_FLASH_SIZE - SECTOR_SIZE_BYTES - 1)) {
      // Middle of flash has 64k lock units
      block_addr = flash_impl_get_sector_base_address(block_addr);
      block_size = SECTOR_SIZE_BYTES;
    } else {
      // Start and end of flash have 1 sector of 4k lock units
      block_addr = flash_impl_get_subsector_base_address(block_addr);
      block_size = SUBSECTOR_SIZE_BYTES;
    }
    const status_t sc = qspi_flash_lock_sector(QSPI_FLASH, block_addr);
    if (FAILED(sc)) {
      return sc;
    }
    block_addr += block_size;
  }

  return S_SUCCESS;
}

status_t flash_impl_unprotect(void) {
  // No way to unprotect all of flash. This requires a full reset of the mt25q
  qspi_flash_init(QSPI_FLASH, &QSPI_FLASH_PART, qspi_flash_is_in_coredump_mode(QSPI_FLASH));
  return S_SUCCESS;
}

static void prv_set_high_drive_strength(void) {
  // Match the impedance of the traces (~50 ohms) by configuring the drive strength of the data
  // output pins on the MT25Q to the 45 ohm setting This avoids signal integreity issues. This is
  // done by setting bits 2:0 in the "Enhanced Volatile Configuration Register" to 101b.
  const uint8_t read_instruction = 0x65;
  const uint8_t write_instruction = 0x61;
  const uint8_t mask = 0x7;
  const uint8_t value = 0x5;
  qspi_flash_ll_set_register_bits(QSPI_FLASH, read_instruction, write_instruction, value, mask);
}

status_t flash_impl_init(bool coredump_mode) {
  qspi_flash_init(QSPI_FLASH, &QSPI_FLASH_PART, coredump_mode);
#if BOARD_CUTTS_BB || BOARD_ROBERT_BB || BOARD_ROBERT_EVT
  prv_set_high_drive_strength();
#endif
  return S_SUCCESS;
}

status_t flash_impl_get_erase_status(void) { return qspi_flash_is_erase_complete(QSPI_FLASH); }

status_t flash_impl_erase_subsector_begin(FlashAddress subsector_addr) {
  return qspi_flash_erase_begin(QSPI_FLASH, subsector_addr, true /* is_subsector */);
}
status_t flash_impl_erase_sector_begin(FlashAddress sector_addr) {
  return qspi_flash_erase_begin(QSPI_FLASH, sector_addr, false /* !is_subsector */);
}

status_t flash_impl_erase_suspend(FlashAddress sector_addr) {
  return qspi_flash_erase_suspend(QSPI_FLASH, sector_addr);
}

status_t flash_impl_erase_resume(FlashAddress sector_addr) {
  qspi_flash_erase_resume(QSPI_FLASH, sector_addr);
  return S_SUCCESS;
}

status_t flash_impl_read_sync(void *buffer_ptr, FlashAddress start_addr, size_t buffer_size) {
  PBL_ASSERT(buffer_size > 0, "flash_impl_read_sync() called with 0 bytes to read");
  qspi_flash_read_blocking(QSPI_FLASH, start_addr, buffer_ptr, buffer_size);
  return S_SUCCESS;
}

int flash_impl_write_page_begin(const void *buffer, const FlashAddress start_addr, size_t len) {
  return qspi_flash_write_page_begin(QSPI_FLASH, buffer, start_addr, len);
}

status_t flash_impl_get_write_status(void) { return qspi_flash_get_write_status(QSPI_FLASH); }

status_t flash_impl_enter_low_power_mode(void) {
  qspi_flash_set_lower_power_mode(QSPI_FLASH, true);
  return S_SUCCESS;
}
status_t flash_impl_exit_low_power_mode(void) {
  qspi_flash_set_lower_power_mode(QSPI_FLASH, false);
  return S_SUCCESS;
}

status_t flash_impl_set_burst_mode(bool burst_mode) {
  // NYI
  return S_SUCCESS;
}

status_t flash_impl_blank_check_sector(FlashAddress addr) {
  return qspi_flash_blank_check(QSPI_FLASH, addr, false /* !is_subsector */);
}
status_t flash_impl_blank_check_subsector(FlashAddress addr) {
  return qspi_flash_blank_check(QSPI_FLASH, addr, true /* is_subsector */);
}

uint32_t flash_impl_get_typical_sector_erase_duration_ms(void) { return 150; }

uint32_t flash_impl_get_typical_subsector_erase_duration_ms(void) { return 50; }
