#include "mbedtls/base64.h"
#include "rom/crc.h"

#include <iostream>
#include <sstream>
#include <cstring>
#include <vector>
#include <string>
#include <cmath>
#include <chrono>

/**
 * \file test_serial_ota_decode.cpp
 * \brief Tests for serial I/O OTA update functions.
 * \author Adam Robichaud
 * 
 * These test functions are meant to be used either with qemu or via 
 * the wokwi esp32 emulator ().
 * 
 * TEST 1 : OTA packet decoder & CRC32 validation.
 * --------------------------------------------------------------------
 * The new Serial OTA update functionality transmits firmware image
 * data over the serial port in 512-byte chunks of data, encoded to 
 * base64, at high speeds (115200 baud). To prevent data corruption, 
 * we've decided to CRC32 checksum the data before writing it to 
 * flash memory. This test verifies that the base64 decoder and CRC32
 * checksum functions are working correctly.
 *  
 */

/**
 * \brief decode_base64()
 * \param input_b64 A base64-encoded payload string to decode. Payload cannot exceed 512 bytes (as per the OTA serial spec).
 * \return A vector of int8_t containing the decoded payload
 */
std::vector<char> decode_base64(std::string input_b64) 
{
    // Calculate the required buffer size for the decoded data
    const size_t buffer_lng = 1025; // One more than protocol allows for
    auto output = std::vector<char>(buffer_lng);

    // Decode the base64 string into the output buffer
    size_t bytesOut = 0;
    auto rc = mbedtls_base64_decode((unsigned char*)output.data(), output.size(), 
      &bytesOut, 
      (const unsigned char*)input_b64.c_str(), input_b64.length());
    
    output.resize(bytesOut);
    if (rc != 0) output.clear();
    return output;
}

/**
 * \brief decode_crc32()
 * \param str the hex-encoded CRC32 checksum to decode
 * \param result uint32_t storage variable for the decoded CRC32 checksum
 * \return true if successful, false otherwise
*/
bool decode_crc32(const std::string& str, uint32_t& result) {
    // Use stringstream for conversion
    std::istringstream iss(str);
    uint64_t temp;
    if (!(iss >> std::hex >> temp)) {
        return false; // Conversion failed
    }
    if (temp > UINT32_MAX) {
        return false; // Value exceeds uint32_t range
    }
    result = static_cast<uint32_t>(temp); // Conversion successful
    return true;
}

/**
 * \brief test_valid_packet()
 * \param[in] input_b64 The base64 encoded data to checksum
 * \param[in] user_crc32_encoded The expected CRC32 checksum, hex-encoded, calculated using zlib.crc32()
 * \returns true if the packet is valid, false otherwise
*/
bool test_valid_packet(std::string payload_b64, std::string user_crc32_encoded) {
  uint32_t user_crc32 = 0;
  decode_crc32(user_crc32_encoded, user_crc32);

  auto start = std::chrono::high_resolution_clock::now();
  auto payload = decode_base64(payload_b64);
  uint32_t payload_crc32 = crc32_le(0, (uint8_t*)payload.data(), payload.size());
  auto stop = std::chrono::high_resolution_clock::now();
  auto dt = std::chrono::duration_cast<std::chrono::microseconds>(stop - start).count();
  Serial.println(payload_b64.length(), DEC);
  Serial.println(dt, DEC);

  return user_crc32 == payload_crc32;
}

void setup() {
  std::string payload_b64 = "aSN5MyBhQRCAQAb3/zLN8EkDWRNpI3kzMsX0OAMyw+CJA5kTqSO5MyCiQSCAQAbu/xCAQNLN8AkNGR0pLTk9MsH0OAPSzRAyw9BJA1kTaSN5M4lDmVOpY7lz8IBAIONBMIBAxuD/AAAQgEAgSAMAEkAMEgAioSBJExAgAAwCDfAMEg3wAADwIhEbIjBIA/fiCgsz8CIR92L3MDA0ABNAADIRMCKBICD0IEkTMAMDLAIwMsAwSBMQIAAMIg3wAAAANiEAYSevUOYDTQBgJRA7IiDmExAgAEXt/w0EUOYTECAAHfAANiEAIOMTHfA4AjDn8zgSMBATOCIwERM4MjAEEzhCMAwTOFIwIBM4YjAhEzhyMCITOIIwIxM4kjDq8ziiMOvzOLIw7PMN8AAAcD7jOQIwEAM5EjARAzkiMAQDOTIwDAM5QjAgAzlSMCEDOWIwIgM5cjAjAzmCoD7jOZKwPuM5osA+4zmyDfAAADYhAAwDMORhECAAICMgIOQTECAALQMd8DYhAAwEQORhECAAIDQgIDMwMOQTECAALQQd8AAAAAAAAAAAAAAAAAAAAAD/JaEN9tHnsjELw0NV4K80ZBHy2zrZsHCAD9ESrRKy+OE=";
  std::string user_crc32_hex = "e38ac2b2";

  Serial.begin(115200);
  Serial.println(test_valid_packet(payload_b64, user_crc32_hex) ? "PASS" : "FAIL");
}

void loop() {
  // put your main code here, to run repeatedly:
  delay(10); // this speeds up the simulation
}