/*
 Esp.cpp - ESP31B-specific APIs
 Copyright (c) 2015 Ivan Grokhotkov. All rights reserved.

 This library is free software; you can redistribute it and/or
 modify it under the terms of the GNU Lesser General Public
 License as published by the Free Software Foundation; either
 version 2.1 of the License, or (at your option) any later version.

 This library is distributed in the hope that it will be useful,
 but WITHOUT ANY WARRANTY; without even the implied warranty of
 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 Lesser General Public License for more details.

 You should have received a copy of the GNU Lesser General Public
 License along with this library; if not, write to the Free Software
 Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA
 */

#include "Arduino.h"
#include "Esp.h"
#include "esp_sleep.h"
#include "spi_flash_mmap.h"
#include "esp_idf_version.h"
#include <memory>
#include <soc/soc.h>
#include <esp_partition.h>
extern "C" {
#include "esp_ota_ops.h"
#include "esp_image_format.h"
}
#include <MD5Builder.h>

#include "soc/spi_reg.h"
#include "esp_system.h"
#include "esp_chip_info.h"
#include "esp_mac.h"
#include "esp_flash.h"

// Include HAL layer for flash clock access
#include "hal/spi_flash_ll.h"

#if CONFIG_IDF_TARGET_ESP32
#include "soc/spi_struct.h"
#else
// All modern chips (S2, S3, C2, C3, C5, C6, H2, P4) use spimem
#include "hal/spimem_flash_ll.h"
// Try to include the newer c_struct header first, fall back to regular struct
#if __has_include("soc/spi_mem_c_struct.h")
#include "soc/spi_mem_c_struct.h"
#else
#include "soc/spi_mem_struct.h"
#endif
#endif

#ifdef ESP_IDF_VERSION_MAJOR  // IDF 4+
#if CONFIG_IDF_TARGET_ESP32   // ESP32/PICO-D4
#include "esp32/rom/spi_flash.h"
#include "soc/efuse_reg.h"
#define ESP_FLASH_IMAGE_BASE 0x1000  // Flash offset containing flash size and spi mode
#elif CONFIG_IDF_TARGET_ESP32S2
#include "esp32s2/rom/spi_flash.h"
#include "soc/efuse_reg.h"
#define ESP_FLASH_IMAGE_BASE 0x1000
#elif CONFIG_IDF_TARGET_ESP32S3
#include "esp32s3/rom/spi_flash.h"
#include "soc/efuse_reg.h"
#define ESP_FLASH_IMAGE_BASE 0x0000  // Esp32s3 is located at 0x0000
#elif CONFIG_IDF_TARGET_ESP32C2
#include "esp32c2/rom/spi_flash.h"
#define ESP_FLASH_IMAGE_BASE 0x0000  // Esp32c2 is located at 0x0000
#elif CONFIG_IDF_TARGET_ESP32C3
#include "esp32c3/rom/spi_flash.h"
#define ESP_FLASH_IMAGE_BASE 0x0000  // Esp32c3 is located at 0x0000
#elif CONFIG_IDF_TARGET_ESP32C6
#include "esp32c6/rom/spi_flash.h"
#define ESP_FLASH_IMAGE_BASE 0x0000  // Esp32c6 is located at 0x0000
#elif CONFIG_IDF_TARGET_ESP32H2
#include "esp32h2/rom/spi_flash.h"
#define ESP_FLASH_IMAGE_BASE 0x0000  // Esp32h2 is located at 0x0000
#elif CONFIG_IDF_TARGET_ESP32P4
#include "esp32p4/rom/spi_flash.h"
#define ESP_FLASH_IMAGE_BASE 0x2000  // Esp32p4 is located at 0x2000
#elif CONFIG_IDF_TARGET_ESP32C5
#include "esp32c5/rom/spi_flash.h"
#define ESP_FLASH_IMAGE_BASE 0x2000  // Esp32c5 is located at 0x2000
#elif CONFIG_IDF_TARGET_ESP32C61
#include "esp32c61/rom/spi_flash.h"
#define ESP_FLASH_IMAGE_BASE 0x0000  // Esp32c61 is located at 0x0000
#else
#error Target CONFIG_IDF_TARGET is not supported
#endif
#else  // ESP32 Before IDF 4.0
#include "rom/spi_flash.h"
#define ESP_FLASH_IMAGE_BASE 0x1000
#endif

// REG_SPI_BASE is not defined for S3/C3 ??

#if CONFIG_IDF_TARGET_ESP32S3 || CONFIG_IDF_TARGET_ESP32C3
#ifdef REG_SPI_BASE
#undef REG_SPI_BASE
#endif  // REG_SPI_BASE
#define REG_SPI_BASE(i) (DR_REG_SPI1_BASE + (((i) > 1) ? (((i) * 0x1000) + 0x20000) : (((~(i)) & 1) * 0x1000)))
#endif  // TARGET

/**
 * User-defined Literals
 *  usage:
 *
 *   uint32_t = test = 10_MHz; // --> 10000000
 */

unsigned long long operator""_kHz(unsigned long long x) {
  return x * 1000;
}

unsigned long long operator""_MHz(unsigned long long x) {
  return x * 1000 * 1000;
}

unsigned long long operator""_GHz(unsigned long long x) {
  return x * 1000 * 1000 * 1000;
}

unsigned long long operator""_kBit(unsigned long long x) {
  return x * 1024;
}

unsigned long long operator""_MBit(unsigned long long x) {
  return x * 1024 * 1024;
}

unsigned long long operator""_GBit(unsigned long long x) {
  return x * 1024 * 1024 * 1024;
}

unsigned long long operator""_kB(unsigned long long x) {
  return x * 1024;
}

unsigned long long operator""_MB(unsigned long long x) {
  return x * 1024 * 1024;
}

unsigned long long operator""_GB(unsigned long long x) {
  return x * 1024 * 1024 * 1024;
}

EspClass ESP;

void EspClass::deepSleep(uint64_t time_us) {
  esp_deep_sleep(time_us);
}

void EspClass::restart(void) {
  esp_restart();
}

uint32_t EspClass::getHeapSize(void) {
  return heap_caps_get_total_size(MALLOC_CAP_INTERNAL);
}

uint32_t EspClass::getFreeHeap(void) {
  return heap_caps_get_free_size(MALLOC_CAP_INTERNAL);
}

uint32_t EspClass::getMinFreeHeap(void) {
  return heap_caps_get_minimum_free_size(MALLOC_CAP_INTERNAL);
}

uint32_t EspClass::getMaxAllocHeap(void) {
  return heap_caps_get_largest_free_block(MALLOC_CAP_INTERNAL);
}

uint32_t EspClass::getPsramSize(void) {
  if (psramFound()) {
    return heap_caps_get_total_size(MALLOC_CAP_SPIRAM);
  }
  return 0;
}

uint32_t EspClass::getFreePsram(void) {
  if (psramFound()) {
    return heap_caps_get_free_size(MALLOC_CAP_SPIRAM);
  }
  return 0;
}

uint32_t EspClass::getMinFreePsram(void) {
  if (psramFound()) {
    return heap_caps_get_minimum_free_size(MALLOC_CAP_SPIRAM);
  }
  return 0;
}

uint32_t EspClass::getMaxAllocPsram(void) {
  if (psramFound()) {
    return heap_caps_get_largest_free_block(MALLOC_CAP_SPIRAM);
  }
  return 0;
}

static uint32_t sketchSize(sketchSize_t response) {
  esp_image_metadata_t data;
  const esp_partition_t *running = esp_ota_get_running_partition();
  if (!running) {
    return 0;
  }
  const esp_partition_pos_t running_pos = {
    .offset = running->address,
    .size = running->size,
  };
  data.start_addr = running_pos.offset;
  esp_image_verify(ESP_IMAGE_VERIFY, &running_pos, &data);
  if (response) {
    return running_pos.size - data.image_len;
  } else {
    return data.image_len;
  }
}

uint32_t EspClass::getSketchSize() {
  return sketchSize(SKETCH_SIZE_TOTAL);
}

String EspClass::getSketchMD5() {
  static String result;
  if (result.length()) {
    return result;
  }
  uint32_t lengthLeft = getSketchSize();

  const esp_partition_t *running = esp_ota_get_running_partition();
  if (!running) {
    log_e("Partition could not be found");
    return String();
  }

  const size_t bufSize = SPI_FLASH_SEC_SIZE;
  uint8_t *pb = (uint8_t *)malloc(bufSize);
  if (!pb) {
    log_e("Not enough memory to allocate buffer");
    return String();
  }
  uint32_t offset = 0;

  MD5Builder md5;
  md5.begin();
  while (lengthLeft > 0) {
    size_t readBytes = (lengthLeft < bufSize) ? lengthLeft : bufSize;
    if (!ESP.flashRead(running->address + offset, (uint32_t *)pb, (readBytes + 3) & ~3)) {
      free(pb);
      log_e("Could not read buffer from flash");
      return String();
    }
    md5.add(pb, readBytes);
    lengthLeft -= readBytes;
    offset += readBytes;

#if CONFIG_FREERTOS_UNICORE
    delay(1);  // Fix solo WDT
#endif
  }
  free(pb);
  md5.calculate();
  result = md5.toString();
  return result;
}

uint32_t EspClass::getFreeSketchSpace() {
  const esp_partition_t *_partition = esp_ota_get_next_update_partition(NULL);
  if (!_partition) {
    return 0;
  }

  return _partition->size;
}

uint16_t EspClass::getChipRevision(void) {
  esp_chip_info_t chip_info;
  esp_chip_info(&chip_info);
  return chip_info.revision;
}

const char *EspClass::getChipModel(void) {
#if CONFIG_IDF_TARGET_ESP32
  uint32_t chip_ver = REG_GET_FIELD(EFUSE_BLK0_RDATA3_REG, EFUSE_RD_CHIP_PACKAGE);
  uint32_t pkg_ver = chip_ver & 0x7;
  switch (pkg_ver) {
    case EFUSE_RD_CHIP_VER_PKG_ESP32D0WDQ6:
      if ((getChipRevision() / 100) == 3) {
        return "ESP32-D0WDQ6-V3";
      } else {
        return "ESP32-D0WDQ6";
      }
    case EFUSE_RD_CHIP_VER_PKG_ESP32D0WDQ5:
      if ((getChipRevision() / 100) == 3) {
        return "ESP32-D0WD-V3";
      } else {
        return "ESP32-D0WD";
      }
    case EFUSE_RD_CHIP_VER_PKG_ESP32D2WDQ5:   return "ESP32-D2WD";
    case EFUSE_RD_CHIP_VER_PKG_ESP32U4WDH:    return "ESP32-U4WDH";
    case EFUSE_RD_CHIP_VER_PKG_ESP32PICOD4:   return "ESP32-PICO-D4";
    case EFUSE_RD_CHIP_VER_PKG_ESP32PICOV302: return "ESP32-PICO-V3-02";
    case EFUSE_RD_CHIP_VER_PKG_ESP32D0WDR2V3: return "ESP32-D0WDR2-V3";
    default:                                  return "Unknown";
  }
#elif CONFIG_IDF_TARGET_ESP32S2
  uint32_t pkg_ver = REG_GET_FIELD(EFUSE_RD_MAC_SPI_SYS_3_REG, EFUSE_PKG_VERSION);
  switch (pkg_ver) {
    case 0:  return "ESP32-S2";
    case 1:  return "ESP32-S2FH16";
    case 2:  return "ESP32-S2FH32";
    default: return "ESP32-S2 (Unknown)";
  }
#else
  esp_chip_info_t chip_info;
  esp_chip_info(&chip_info);
  switch (chip_info.model) {
    case CHIP_ESP32S3: return "ESP32-S3";
    case CHIP_ESP32C3: return "ESP32-C3";
    case CHIP_ESP32C2: return "ESP32-C2";
    case CHIP_ESP32C6: return "ESP32-C6";
    case CHIP_ESP32H2: return "ESP32-H2";
    case CHIP_ESP32P4: return "ESP32-P4";
#if ESP_IDF_VERSION >= ESP_IDF_VERSION_VAL(5, 5, 0)
    case CHIP_ESP32C5:  return "ESP32-C5";
    case CHIP_ESP32C61: return "ESP32-C61";
    case CHIP_ESP32H21: return "ESP32-H21";
#endif
    default: return "UNKNOWN";
  }
#endif
}

uint8_t EspClass::getChipCores(void) {
  esp_chip_info_t chip_info;
  esp_chip_info(&chip_info);
  return chip_info.cores;
}

const char *EspClass::getSdkVersion(void) {
  return esp_get_idf_version();
}

const char *EspClass::getCoreVersion(void) {
  return ESP_ARDUINO_VERSION_STR;
}

uint32_t ESP_getFlashChipId(void) {
  uint32_t id = g_rom_flashchip.device_id;
  id = ((id & 0xff) << 16) | ((id >> 16) & 0xff) | (id & 0xff00);
  return id;
}

uint32_t EspClass::getFlashChipSize(void) {
  uint32_t id = (ESP_getFlashChipId() >> 16) & 0xFF;
  return 2 << (id - 1);
}

uint32_t EspClass::getFlashChipSpeed(void) {
  esp_image_header_t fhdr;
  if (esp_flash_read(esp_flash_default_chip, (void *)&fhdr, ESP_FLASH_IMAGE_BASE, sizeof(esp_image_header_t)) && fhdr.magic != ESP_IMAGE_HEADER_MAGIC) {
    return 0;
  }
  return magicFlashChipSpeed(fhdr.spi_speed);
}

FlashMode_t EspClass::getFlashChipMode(void) {
#if CONFIG_IDF_TARGET_ESP32S2 || CONFIG_IDF_TARGET_ESP32P4 || CONFIG_IDF_TARGET_ESP32C5 || CONFIG_IDF_TARGET_ESP32C61
  uint32_t spi_ctrl = REG_READ(PERIPHS_SPI_FLASH_CTRL);
#elif CONFIG_IDF_TARGET_ESP32H2 || CONFIG_IDF_TARGET_ESP32C2 || CONFIG_IDF_TARGET_ESP32C6
  uint32_t spi_ctrl = REG_READ(DR_REG_SPI0_BASE + 0x8);
#else
  uint32_t spi_ctrl = REG_READ(SPI_CTRL_REG(0));
#endif
  /* Not all of the following constants are already defined in older versions of spi_reg.h, so do it manually for now*/
  if (spi_ctrl & BIT(24)) {  //SPI_FREAD_QIO
    return (FM_QIO);
  } else if (spi_ctrl & BIT(20)) {  //SPI_FREAD_QUAD
    return (FM_QOUT);
  } else if (spi_ctrl & BIT(23)) {  //SPI_FREAD_DIO
    return (FM_DIO);
  } else if (spi_ctrl & BIT(14)) {  // SPI_FREAD_DUAL
    return (FM_DOUT);
  } else if (spi_ctrl & BIT(13)) {  //SPI_FASTRD_MODE
    return (FM_FAST_READ);
  } else {
    return (FM_SLOW_READ);
  }
}

uint32_t EspClass::magicFlashChipSize(uint8_t flashByte) {
  /*
    FLASH_SIZES = {
        "1MB": 0x00,
        "2MB": 0x10,
        "4MB": 0x20,
        "8MB": 0x30,
        "16MB": 0x40,
        "32MB": 0x50,
        "64MB": 0x60,
        "128MB": 0x70,
    }
*/
  switch (flashByte & 0x0F) {
    case 0x0: return (1_MB);    // 8 MBit (1MB)
    case 0x1: return (2_MB);    // 16 MBit (2MB)
    case 0x2: return (4_MB);    // 32 MBit (4MB)
    case 0x3: return (8_MB);    // 64 MBit (8MB)
    case 0x4: return (16_MB);   // 128 MBit (16MB)
    case 0x5: return (32_MB);   // 256 MBit (32MB)
    case 0x6: return (64_MB);   // 512 MBit (64MB)
    case 0x7: return (128_MB);  // 1 GBit (128MB)
    default:                    // fail?
      return 0;
  }
}

uint32_t EspClass::magicFlashChipSpeed(uint8_t flashByte) {
#if CONFIG_IDF_TARGET_ESP32C2
  /*
    FLASH_FREQUENCY = {
        "60m": 0xF,
        "30m": 0x0,
        "20m": 0x1,
        "15m": 0x2,
    }
*/
  switch (flashByte & 0x0F) {
    case 0xF: return (60_MHz);
    case 0x0: return (30_MHz);
    case 0x1: return (20_MHz);
    case 0x2: return (15_MHz);
    default:  // fail?
      return 0;
  }

#elif CONFIG_IDF_TARGET_ESP32C6
  /*
   FLASH_FREQUENCY = {
        "80m": 0x0,  # workaround for wrong mspi HS div value in ROM
        "40m": 0x0,
        "20m": 0x2,
    }
*/
  switch (flashByte & 0x0F) {
    case 0x0: return (80_MHz);
    case 0x2: return (20_MHz);
    default:  // fail?
      return 0;
  }

#elif CONFIG_IDF_TARGET_ESP32C61
  /*
    FLASH_FREQUENCY = {
        "80m": 0xF,
        "40m": 0x0,
        "20m": 0x2,
    }
*/
  switch (flashByte & 0x0F) {
    case 0xF: return (80_MHz);
    case 0x0: return (40_MHz);
    case 0x2: return (20_MHz);
    default:  // fail?
      return 0;
  }

#elif CONFIG_IDF_TARGET_ESP32H2

  /*
    FLASH_FREQUENCY = {
        "48m": 0xF,
        "24m": 0x0,
        "16m": 0x1,
        "12m": 0x2,
    }
*/
  switch (flashByte & 0x0F) {
    case 0xF: return (48_MHz);
    case 0x0: return (24_MHz);
    case 0x1: return (16_MHz);
    case 0x2: return (12_MHz);
    default:  // fail?
      return 0;
  }

#else
  /*
    FLASH_FREQUENCY = {
        "80m": 0xF,
        "40m": 0x0,
        "26m": 0x1,
        "20m": 0x2,
    }
*/
  switch (flashByte & 0x0F) {
    case 0xF: return (80_MHz);
    case 0x0: return (40_MHz);
    case 0x1: return (26_MHz);
    case 0x2: return (20_MHz);
    default:  // fail?
      return 0;
  }
#endif
}

FlashMode_t EspClass::magicFlashChipMode(uint8_t flashByte) {
  FlashMode_t mode = (FlashMode_t)flashByte;
  if (mode > FM_SLOW_READ) {
    mode = FM_UNKNOWN;
  }
  return mode;
}

bool EspClass::flashEraseSector(uint32_t sector) {
  return esp_flash_erase_region(esp_flash_default_chip, sector * SPI_FLASH_SEC_SIZE, SPI_FLASH_SEC_SIZE) == ESP_OK;
}

// Warning: These functions do not work with encrypted flash
bool EspClass::flashWrite(uint32_t offset, uint32_t *data, size_t size) {
  return esp_flash_write(esp_flash_default_chip, (const void *)data, offset, size) == ESP_OK;
}

bool EspClass::flashRead(uint32_t offset, uint32_t *data, size_t size) {
  return esp_flash_read(esp_flash_default_chip, (void *)data, offset, size) == ESP_OK;
}

bool EspClass::partitionEraseRange(const esp_partition_t *partition, uint32_t offset, size_t size) {
  return esp_partition_erase_range(partition, offset, size) == ESP_OK;
}

bool EspClass::partitionWrite(const esp_partition_t *partition, uint32_t offset, uint32_t *data, size_t size) {
  return esp_partition_write(partition, offset, data, size) == ESP_OK;
}

bool EspClass::partitionRead(const esp_partition_t *partition, uint32_t offset, uint32_t *data, size_t size) {
  return esp_partition_read(partition, offset, data, size) == ESP_OK;
}

uint64_t EspClass::getEfuseMac(void) {
  uint64_t _chipmacid = 0LL;
  esp_efuse_mac_get_default((uint8_t *)(&_chipmacid));
  return _chipmacid;
}

// ============================================================================
// Flash Frequency Runtime Detection
// ============================================================================

/**
 * @brief Read the source clock frequency using ESP-IDF HAL functions
 * @return Source clock frequency in MHz (80, 120, 160, or 240)
 */
uint8_t EspClass::getFlashSourceFrequencyMHz(void) {
#if CONFIG_IDF_TARGET_ESP32
  // ESP32: Use HAL function
  return spi_flash_ll_get_source_clock_freq_mhz(0);  // host_id = 0 for SPI0
#else
  // All modern MCUs: Use spimem HAL function
  return spimem_flash_ll_get_source_freq_mhz();
#endif
}

/**
 * @brief Read the clock divider from hardware using HAL structures
 * Based on ESP-IDF HAL implementation:
 * - ESP32: Uses SPI1.clock (typedef in spi_flash_ll.h)
 * - All newer MCUs: Use SPIMEM1.clock (typedef in spimem_flash_ll.h)
 * @return Clock divider value (1 = no division, 2 = divide by 2, etc.)
 */
uint8_t EspClass::getFlashClockDivider(void) {
#if CONFIG_IDF_TARGET_ESP32
  // ESP32: Flash uses SPI1
  // See: line 52: esp-idf/components/hal/esp32/include/hal/spi_flash_ll.h
  if (SPI1.clock.clk_equ_sysclk) {
    return 1;  // 1:1 clock
  }
  return SPI1.clock.clkcnt_n + 1;
#else
  // All newer MCUs: Flash uses SPIMEM1
  // See: esp-idf/components/hal/esp32*/include/hal/spimem_flash_ll.h
  // Example S3: line 38: typedef typeof(SPIMEM1.clock.val) spimem_flash_ll_clock_reg_t;
  // Example C5: lines 97-99: esp-idf/components/soc/esp32c5/mp/include/soc/spi_mem_struct.h
  if (SPIMEM1.clock.clk_equ_sysclk) {
    return 1;  // 1:1 clock
  }
  return SPIMEM1.clock.clkcnt_n + 1;
#endif
}

/**
 * @brief Get the actual flash frequency in MHz
 * @return Flash frequency in MHz (80, 120, 160, or 240)
 */
uint32_t EspClass::getFlashFrequencyMHz(void) {
  uint8_t source = getFlashSourceFrequencyMHz();
  uint8_t divider = getFlashClockDivider();

  if (divider == 0) {
    divider = 1;  // Safety check
  }

  return source / divider;
}
