#include "sw/device/lib/spi_flash.h"

#include <string.h>

#include "hw/top_matcha/sw/autogen/top_matcha.h"
#include "sw/device/lib/dif/dif_spi_host.h"
#include "sw/device/lib/tar.h"
#include "sw/device/lib/testing/test_framework/check.h"

static dif_spi_host_t spi_host0;
static const int kSpiFlashBytes = (64 * 1024 * 1024);

dif_result_t spi_flash_read_page(uint32_t page, uint8_t* buf) {
  /* clang-format off */
  dif_spi_host_segment_t segments[] = {
      {
          .type = kDifSpiHostSegmentTypeOpcode,
          .opcode = 0x13,  // Page Read, 4b Addr
      },
      {.type = kDifSpiHostSegmentTypeAddress,
       .address = {
               .width = kDifSpiHostWidthStandard,
               .mode = kDifSpiHostAddrMode4b,
               .address = page,
           }},
      // Two RX segments, due to the RX fifo only being 256B.
      {.type = kDifSpiHostSegmentTypeRx,
       .rx = {.width = kDifSpiHostWidthStandard,
              .buf = buf,
              .length = SPI_PAGE_SIZE >> 1}},
      {.type = kDifSpiHostSegmentTypeRx,
       .rx = {.width = kDifSpiHostWidthStandard,
              .buf = buf + (SPI_PAGE_SIZE >> 1),
              .length = SPI_PAGE_SIZE >> 1}},
  };
  /* clang-format on */
  CHECK_DIF_OK(dif_spi_host_transaction(&spi_host0, /*csid=*/0, segments,
                                        ARRAYSIZE(segments)));
  return kDifOk;
}

dif_result_t spi_flash_init(void) {
  CHECK_DIF_OK(dif_spi_host_init(
      mmio_region_from_addr(TOP_MATCHA_SPI_HOST0_BASE_ADDR), &spi_host0));
  dif_spi_host_config_t config = {
      // TODO(atv): Add spi clock into arch definition
      .spi_clock = 1000000,
      .peripheral_clock_freq_hz = 25 * 100 * 1000,
  };
  CHECK_DIF_OK(dif_spi_host_configure(&spi_host0, config));
  CHECK_DIF_OK(dif_spi_host_output_set_enabled(&spi_host0, /*enabled=*/true));
  return kDifOk;
}

static int parse_octal(uint8_t* cursor) {
  int size = 0;
  while (*cursor) size = (size << 3) | ((*cursor++) - '0');
  return size;
}

int contains(const void* src, const char* str2) {
  const char* str1 = (const char*)src;
  const char* cursor = str2;
  while (*str1 && *cursor) {
    cursor = (*str1 == *cursor) ? cursor + 1 : str2;
    str1++;
  }
  return *cursor ? 0 : 1;
}

// Returns the offset of the file in the tar.
// `size_out` will be filled with the size of the file.
static size_t find_file_in_tar(const char* filename, size_t* size_out) {
  if (!size_out || !filename) {
    return 0;
  }

  tar_header tar;
  size_t cursor = 0;
  CHECK_DIF_OK(spi_flash_read_page(cursor, (uint8_t*)&tar));
  while (cursor < kSpiFlashBytes) {
    int size = parse_octal((uint8_t*)&tar.size);
    cursor += 512;

    if (contains(&tar.name, filename)) {
      if (size_out) {
        *size_out = size;
      }
      return cursor;
    }

    cursor += (size + 511) & ~511;
    CHECK_DIF_OK(spi_flash_read_page(cursor, (uint8_t*)&tar));
  }

  return 0;
}

dif_result_t load_file_from_tar(const char* filename, void* addr) {
  size_t size = 0;
  size_t bin_offset = find_file_in_tar(filename, &size);
  if (!bin_offset) return kDifError;

  uint8_t page[SPI_PAGE_SIZE];
  size_t remaining = size;
  while (remaining > 0) {
    size_t page_used = remaining >= SPI_PAGE_SIZE ? SPI_PAGE_SIZE : remaining;
    size_t offset = size - remaining;
    CHECK_DIF_OK(spi_flash_read_page(bin_offset + offset, page));
    void* dst = (uint8_t*)addr + offset;
    if (memcpy(dst, page, page_used) != dst) {
      return kDifError;
    }
    remaining -= page_used;
  }

  return kDifOk;
}