/**
 * DID Mesh Single File Test for Wokwi
 * Combined all modules into one file for online simulation
 * Tests only basic DID crypto operations
 */

#include "Arduino.h"
#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>
#include "mbedtls/ecdsa.h"
#include "mbedtls/entropy.h"
#include "mbedtls/ctr_drbg.h"
#include "mbedtls/pk.h"
#include "esp_timer.h"

// Display Configuration
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
#define OLED_RESET -1
#define OLED_ADDRESS 0x3C
#define SDA_PIN 9
#define SCL_PIN 8

// Global variables
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);
mbedtls_pk_context nodeKey;
mbedtls_entropy_context entropy;
mbedtls_ctr_drbg_context ctr_drbg;
bool displayAvailable = false;

// DID Identity Structure
struct DIDIdentity {
    String did;
    uint8_t hash[16];
    bool valid;
};

DIDIdentity nodeIdentity;

void displayMessage(String line1, String line2 = "", String line3 = "", String line4 = "") {
    if (!displayAvailable) return;
    
    display.clearDisplay();
    display.setTextSize(1);
    display.setTextColor(SSD1306_WHITE);
    
    display.setCursor(0, 0);
    display.println(line1);
    
    if (line2.length() > 0) {
        display.setCursor(0, 16);
        display.println(line2);
    }
    
    if (line3.length() > 0) {
        display.setCursor(0, 32);
        display.println(line3);
    }
    
    if (line4.length() > 0) {
        display.setCursor(0, 48);
        display.println(line4);
    }
    
    display.display();
}

void displayResult(String test, float timeMs, bool success) {
    if (!displayAvailable) return;
    
    String status = success ? "PASS" : "FAIL";
    String timeStr = String(timeMs, 1) + "ms";
    
    display.clearDisplay();
    display.setTextSize(1);
    display.setTextColor(SSD1306_WHITE);
    
    display.setCursor(0, 0);
    display.println("DID Crypto Test");
    
    display.setCursor(0, 16);
    display.println(test);
    
    display.setCursor(0, 32);
    display.println("Time: " + timeStr);
    
    display.setCursor(0, 48);
    display.setTextSize(2);
    display.println(status);
    
    display.display();
}

bool initDisplay() {
    Wire.begin(SDA_PIN, SCL_PIN);
    delay(100);
    
    if (!display.begin(SSD1306_SWITCHCAPVCC, OLED_ADDRESS)) {
        Serial.println("Display init failed");
        return false;
    }
    
    display.clearDisplay();
    display.setTextSize(1);
    display.setTextColor(SSD1306_WHITE);
    display.display();
    
    return true;
}

bool initCrypto() {
    mbedtls_entropy_init(&entropy);
    mbedtls_ctr_drbg_init(&ctr_drbg);
    mbedtls_pk_init(&nodeKey);
    
    const char* pers = "did_mesh_test";
    int ret = mbedtls_ctr_drbg_seed(&ctr_drbg, mbedtls_entropy_func, &entropy,
                                   (const unsigned char*)pers, strlen(pers));
    
    if (ret != 0) {
        Serial.printf("Crypto init failed: %d\n", ret);
        return false;
    }
    
    return true;
}

bool generateDIDIdentity() {
    Serial.println("Generating DID identity...");
    
    uint64_t start = esp_timer_get_time();
    
    // Set up ECDSA key
    int ret = mbedtls_pk_setup(&nodeKey, mbedtls_pk_info_from_type(MBEDTLS_PK_ECKEY));
    if (ret != 0) {
        Serial.printf("Key setup failed: %d\n", ret);
        return false;
    }
    
    // Generate key pair
    ret = mbedtls_ecp_gen_key(MBEDTLS_ECP_DP_SECP256R1,
                             mbedtls_pk_ec(nodeKey),
                             mbedtls_ctr_drbg_random,
                             &ctr_drbg);
    
    if (ret != 0) {
        Serial.printf("Key generation failed: %d\n", ret);
        return false;
    }
    
    float timeMs = (esp_timer_get_time() - start) / 1000.0;
    
    // Create DID identifier
    nodeIdentity.did = "did:key:z6Mk" + String(random(100000, 999999));
    
    // Create simple hash
    uint32_t seed = 0;
    for (int i = 0; i < nodeIdentity.did.length(); i++) {
        seed = seed * 31 + nodeIdentity.did.charAt(i);
    }
    
    for (int i = 0; i < 16; i++) {
        nodeIdentity.hash[i] = (seed + i) & 0xFF;
    }
    
    nodeIdentity.valid = true;
    
    Serial.printf("DID identity generated in %.2f ms\n", timeMs);
    Serial.println("DID: " + nodeIdentity.did);
    
    return true;
}

bool testDIDSigning() {
    Serial.println("Testing DID message signing...");
    
    const char* testMessage = "Hello DID from ESP32!";
    uint8_t signature[128];
    size_t sigLen;
    
    Serial.printf("Signing message: \"%s\"\n", testMessage);
    
    // Test signing
    uint64_t start = esp_timer_get_time();
    
    mbedtls_pk_context* key = const_cast<mbedtls_pk_context*>(&nodeKey);
    
    int ret = mbedtls_pk_sign(key,
                             MBEDTLS_MD_SHA256,
                             (const unsigned char*)testMessage, strlen(testMessage),
                             signature, &sigLen,
                             mbedtls_ctr_drbg_random,
                             &ctr_drbg);
    
    uint64_t sign_duration = esp_timer_get_time() - start;
    
    if (ret != 0) {
        Serial.printf("Signing failed: %d\n", ret);
        return false;
    }
    
    Serial.printf("Message signed in %.2f ms\n", sign_duration / 1000.0);
    Serial.printf("Signature length: %d bytes\n", sigLen);
    
    // Test verification
    start = esp_timer_get_time();
    
    ret = mbedtls_pk_verify(key,
                           MBEDTLS_MD_SHA256,
                           (const unsigned char*)testMessage, strlen(testMessage),
                           signature, sigLen);
    
    uint64_t verify_duration = esp_timer_get_time() - start;
    float totalTimeMs = (sign_duration + verify_duration) / 1000.0;
    
    if (ret == 0) {
        Serial.printf("Signature verified in %.2f ms\n", verify_duration / 1000.0);
        Serial.printf("Total sign+verify: %.2f ms\n", totalTimeMs);
        Serial.println("Message integrity and authenticity confirmed");
        
        displayResult("Sign+Verify", totalTimeMs, true);
        return true;
    } else {
        Serial.printf("Verification failed: %d\n", ret);
        displayResult("Verify Failed", 0, false);
        return false;
    }
}

void setup() {
    Serial.begin(115200);
    delay(2000);
    
    Serial.println("================================");
    Serial.println("DID Mesh Crypto Test (Wokwi)");
    Serial.println("================================");
    
    // Initialize display
    displayAvailable = initDisplay();
    if (displayAvailable) {
        displayMessage("DID Mesh", "Crypto Test", "Starting...");
        Serial.println("Display initialized");
    } else {
        Serial.println("Display not available");
    }
    
    delay(2000);
    
    // Initialize crypto
    if (!initCrypto()) {
        Serial.println("Crypto initialization failed");
        if (displayAvailable) {
            displayMessage("ERROR", "Crypto Init", "Failed");
        }
        return;
    }
    
    Serial.println("Crypto initialized");
    if (displayAvailable) {
        displayMessage("Crypto Ready", "Generating DID...");
    }
    
    delay(1000);
    
    // Generate DID identity
    if (!generateDIDIdentity()) {
        Serial.println("DID generation failed");
        if (displayAvailable) {
            displayMessage("ERROR", "DID Generation", "Failed");
        }
        return;
    }
    
    Serial.println("DID identity created");
    if (displayAvailable) {
        displayMessage("DID Created", "Testing crypto...");
    }
    
    delay(1000);
    
    // Test DID signing
    bool success = testDIDSigning();
    
    // Final result
    Serial.println("\n================================");
    if (success) {
        Serial.println("DID CRYPTO TEST PASSED!");
        Serial.println("DID mesh security validated");
        if (displayAvailable) {
            displayMessage("TEST COMPLETE", "DID Crypto", "PASSED", "Ready for mesh!");
        }
    } else {
        Serial.println("DID CRYPTO TEST FAILED");
        if (displayAvailable) {
            displayMessage("TEST FAILED", "Check serial", "for details");
        }
    }
    Serial.println("================================");
}

void loop() {
    delay(10000);
    Serial.println("Test complete - system idle");
}