#include "ota_update.h"
#include "wled.h"
#include "wled_metadata.h"

#ifndef WLED_VERSION
  #warning WLED_VERSION was not set - using default value of 'dev'
  #define WLED_VERSION dev
#endif
#ifndef WLED_RELEASE_NAME
  #warning WLED_RELEASE_NAME was not set - using default value of 'Custom'
  #define WLED_RELEASE_NAME "Custom"
#endif
#ifndef WLED_REPO
  // No warning for this one: integrators are not always on GitHub
  #define WLED_REPO "unknown"
#endif

constexpr uint32_t WLED_CUSTOM_DESC_MAGIC = 0x57535453;  // "WSTS" (WLED System Tag Structure)
constexpr uint32_t WLED_CUSTOM_DESC_VERSION = 2;    // v1 - original PR; v2 - "safe to update from" version

// Compile-time validation that release name doesn't exceed maximum length
static_assert(sizeof(WLED_RELEASE_NAME) <= WLED_RELEASE_NAME_MAX_LEN, 
              "WLED_RELEASE_NAME exceeds maximum length of WLED_RELEASE_NAME_MAX_LEN characters");


/**
 * DJB2 hash function (C++11 compatible constexpr)
 * Used for compile-time hash computation to validate structure contents
 * Recursive for compile time: not usable at runtime due to stack depth
 * 
 * Note that this only works on strings; there is no way to produce a compile-time
 * hash of a struct in C++11 without explicitly listing all the struct members.
 * So for now, we hash only the release name.  This suffices for a "did you find 
 * valid structure" check.
 * 
 */
constexpr uint32_t djb2_hash_constexpr(const char* str, uint32_t hash = 5381) {
    return (*str == '\0') ? hash : djb2_hash_constexpr(str + 1, ((hash << 5) + hash) + *str);
}

/**
 * Runtime DJB2 hash function for validation
 */
inline uint32_t djb2_hash_runtime(const char* str) {
    uint32_t hash = 5381;
    while (*str) {
        hash = ((hash << 5) + hash) + *str++;
    }
    return hash;
}

// ------------------------------------
// GLOBAL VARIABLES
// ------------------------------------
// Structure instantiation for this build 
const wled_metadata_t __attribute__((section(BUILD_METADATA_SECTION))) WLED_BUILD_DESCRIPTION = {
    WLED_CUSTOM_DESC_MAGIC,                   // magic
    /*WLED_CUSTOM_DESC_VERSION*/ 1,           // structure version.  Currently set to 1 to allow OTA from broken original version. FIXME before 0.16 release.
    TOSTRING(WLED_VERSION),
    WLED_RELEASE_NAME,                        // release_name
    std::integral_constant<uint32_t, djb2_hash_constexpr(WLED_RELEASE_NAME)>::value, // hash - computed at compile time; integral_constant enforces this
#if defined(ESP32) && defined(CONFIG_IDF_TARGET_ESP32)
    { 0, 15, 3 },  // Some older ESP32 might have bootloader issues; assume we'll have it sorted by 0.15.3
#else    
    { 0, 15, 2 },  // All other platforms can update safely
#endif
};

static const char repoString_s[] PROGMEM = WLED_REPO;
const __FlashStringHelper* repoString = FPSTR(repoString_s);

static const char productString_s[] PROGMEM = WLED_PRODUCT_NAME;
const __FlashStringHelper* productString = FPSTR(productString_s);

static const char brandString_s [] PROGMEM = WLED_BRAND;
const __FlashStringHelper* brandString = FPSTR(brandString_s);



/**
 * Extract WLED custom description structure from binary
 * @param binaryData Pointer to binary file data
 * @param dataSize Size of binary data in bytes
 * @param extractedDesc Buffer to store extracted custom description structure
 * @return true if structure was found and extracted, false otherwise
 */
bool findWledMetadata(const uint8_t* binaryData, size_t dataSize, wled_metadata_t* extractedDesc) {
  if (!binaryData || !extractedDesc || dataSize < sizeof(wled_metadata_t)) {
    return false;
  }

  for (size_t offset = 0; offset <= dataSize - sizeof(wled_metadata_t); offset++) {
    if ((binaryData[offset]) == static_cast<char>(WLED_CUSTOM_DESC_MAGIC)) {
      // First byte matched; check next in an alignment-safe way
      uint32_t data_magic;
      memcpy(&data_magic, binaryData + offset, sizeof(data_magic));
      
      // Check for magic number
      if (data_magic == WLED_CUSTOM_DESC_MAGIC) {            
        wled_metadata_t candidate;
        memcpy(&candidate, binaryData + offset, sizeof(candidate));

        // Validate hash using runtime function
        uint32_t expected_hash = djb2_hash_runtime(candidate.release_name);
        if (candidate.hash != expected_hash) {
          DEBUG_PRINTF_P(PSTR("Found WLED structure at offset %u but hash mismatch\n"), offset);
          continue;
        }
        
        // Valid structure found - copy entire structure
        *extractedDesc = candidate;
        
        DEBUG_PRINTF_P(PSTR("Extracted WLED structure at offset %u: '%s'\n"), 
                      offset, extractedDesc->release_name);
        return true;
      }
    }
  }
  
  DEBUG_PRINTLN(F("No WLED custom description found in binary"));
  return false;
}


// Strip "_V4" suffix from a release name to allow upgrading between IDF v4 and newer IDF builds.
static String normalizeReleaseName(const String& name) {
  if (name.endsWith("_V4")) return name.substring(0, name.length() - 3);
  return name;
}

template<size_t len>
static inline String bufToString(const char (&buf)[len]) {
  char sbuf[len+1];
  size_t real_len = strnlen(buf, len);
  memcpy(sbuf, buf, real_len);
  sbuf[len] = '\0';
  return sbuf;
}

/**
 * Check if OTA should be allowed based on release compatibility using custom description
 * @param firmwareDescription Description object from proposed new firmware
 * @param errorMessage Buffer to store error message if validation fails 
 * @param errorMessageLen Maximum length of error message buffer
 * @return true if OTA should proceed, false if it should be blocked
 */
bool shouldAllowOTA(const wled_metadata_t& firmwareDescription, char* errorMessage, size_t errorMessageLen) {
  // Clear error message
  if (errorMessage && errorMessageLen > 0) {
    errorMessage[0] = '\0';
  }

  const String uploadedRelease = bufToString(firmwareDescription.release_name);
  
  if (normalizeReleaseName(uploadedRelease) != normalizeReleaseName(releaseString)) {
    if (errorMessage && errorMessageLen > 0) {
      snprintf_P(errorMessage, errorMessageLen, PSTR("Firmware release name mismatch: current='%s', uploaded='%s'."),
                releaseString, uploadedRelease.c_str());
      errorMessage[errorMessageLen - 1] = '\0'; // Ensure null termination
    }
    return false;
  }

  if (firmwareDescription.desc_version > 1) {
    // Add safe version check
    // Parse our version (x.y.z) and compare it to the "safe version" array
    const char* our_version = versionString;
    for(unsigned v_index = 0; v_index < 3; ++v_index) {
      char* our_version_end = nullptr;
      long our_v_parsed = strtol(our_version, &our_version_end, 10); 
      if (!our_version_end || (our_version_end == our_version)) {
        // We were built with a malformed version string
        // We blame the integrator and attempt the update anyways - nothing the user can do to fix this
        break;
      }

      if (firmwareDescription.safe_update_version[v_index] > our_v_parsed) {
        if (errorMessage && errorMessageLen > 0) {
          snprintf_P(errorMessage, errorMessageLen, PSTR("Cannot update from this version: requires at least %d.%d.%d, current='%s'."), 
                  firmwareDescription.safe_update_version[0], firmwareDescription.safe_update_version[1], firmwareDescription.safe_update_version[2],
                  versionString);
          errorMessage[errorMessageLen - 1] = '\0'; // Ensure null termination
        }
        return false;
      } else if (firmwareDescription.safe_update_version[v_index] < our_v_parsed) {
        break;  // no need to check the other components
      }

      if (*our_version_end == '.') ++our_version_end;
      our_version = our_version_end;
    }  
  }

  // TODO: additional checks go here

  return true;
}
