#include <Arduino.h>
#include <ArduinoJson.h>

// ========== ЛЁГКАЯ ЗАМЕНА КРИПТОГРАФИИ (вместо mbedtls) ==========
class LightGOST {
private:
    uint8_t secretKey[16] = {0x12, 0x34, 0x56, 0x78, 0x9A, 0xBC, 0xDE, 0xF0,
                              0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88};
    
    void simpleHash(const String& message, uint8_t* output) {
        uint32_t hash = 0xDEADBEEF;
        
        for(int i = 0; i < message.length(); i++) {
            hash ^= (message[i] << ((i % 4) * 8));
            hash ^= (secretKey[i % 16] << (((i + 7) % 4) * 8));
            hash = (hash << 5) | (hash >> 27);
            hash ^= 0x9E3779B9;
        }
        
        for(int i = 0; i < 32; i++) {
            output[i] = (hash >> ((i % 4) * 8)) ^ secretKey[i % 16];
        }
    }
    
public:
    void sign(const String& message, uint8_t* signature) {
        simpleHash(message, signature);
    }
    
    // ИСПРАВЛЕНО: теперь принимает const uint8_t*
    bool verify(const String& message, const uint8_t* signature) {
        uint8_t computed[32];
        simpleHash(message, computed);
        
        for(int i = 0; i < 32; i++) {
            if(signature[i] != computed[i]) return false;
        }
        return true;
    }
    
    void printSignature(const uint8_t* sig, int len = 32) {
        for(int i = 0; i < len; i++) {
            if(sig[i] < 0x10) Serial.print("0");
            Serial.print(sig[i], HEX);
            if((i + 1) % 8 == 0 && i < len - 1) Serial.print(" ");
        }
    }
};

// ========== СТРУКТУРЫ ДАННЫХ ==========
struct Signature {
    uint8_t hash[32];
    uint32_t timestamp;
};

struct Message {
    String json;
    Signature signature;
    bool verified;
    int retryCount;
};

// ========== ГЕНЕРАТОР JSON ==========
class JSONGenerator {
public:
    static String generateSensorData(int msgId) {
        StaticJsonDocument<128> doc;
        doc["msg_id"] = msgId;
        doc["dev"] = "STM32-001";
        doc["ts"] = millis();
        doc["temp"] = random(200, 350) / 10.0;
        doc["hum"] = random(300, 800) / 10.0;
        
        String output;
        serializeJson(doc, output);
        return output;
    }
    
    static String generateConfirmation(bool verified, int msgId) {
        StaticJsonDocument<64> doc;
        doc["type"] = "ack";
        doc["id"] = msgId;
        doc["ver"] = verified ? 1 : 0;
        doc["ts"] = millis();
        
        String output;
        serializeJson(doc, output);
        return output;
    }
    
    static String generateRetryRequest(int msgId) {
        StaticJsonDocument<64> doc;
        doc["type"] = "retry";
        doc["id"] = msgId;
        doc["reason"] = "bad_sig";
        doc["ts"] = millis();
        
        String output;
        serializeJson(doc, output);
        return output;
    }
};

// ========== ОСНОВНАЯ ЛОГИКА ОБМЕНА ==========
class SecureTransceiver {
private:
    LightGOST crypto;
    Message lastSent;
    Message lastReceived;
    int msgCounter;
    int failCount;
    
public:
    SecureTransceiver() : msgCounter(0), failCount(0) {
        randomSeed(analogRead(0));
    }
    
    Message transmit(const String& jsonData, bool corrupt = false) {
        Message msg;
        msg.json = corrupt ? corruptData(jsonData) : jsonData;
        msg.signature.timestamp = millis();
        
        // ИСПРАВЛЕНО: передаём указатель на hash
        crypto.sign(msg.json, msg.signature.hash);
        
        msg.verified = false;
        msg.retryCount = failCount;
        
        lastSent = msg;
        msgCounter++;
        
        Serial.println(F("\n[TX] Message #"));
        Serial.print(msgCounter);
        Serial.print(F(" JSON: "));
        Serial.println(msg.json);
        Serial.print(F("     Sig: "));
        crypto.printSignature(msg.signature.hash);
        Serial.println();
        
        if(corrupt) {
            Serial.println(F("     ⚠️ DATA CORRUPTED!"));
        } else {
            Serial.println(F("     ✅ Data OK"));
        }
        
        return msg;
    }
    
    bool receiveAndVerify(const Message& msg) {
        Serial.println(F("\n[RX] Received: "));
        Serial.print(F("     JSON: "));
        Serial.println(msg.json);
        
        // ИСПРАВЛЕНО: передаём константный указатель
        bool isValid = crypto.verify(msg.json, msg.signature.hash);
        
        if(isValid) {
            Serial.println(F("     🔐 Signature: VALID"));
            Serial.println(F("     ✅ Message integrity OK"));
        } else {
            Serial.println(F("     🔐 Signature: INVALID"));
            Serial.println(F("     ❌ Message COMPROMISED!"));
        }
        
        // ИСПРАВЛЕНО: убираем const для сохранения
        lastReceived.json = msg.json;
        memcpy(lastReceived.signature.hash, msg.signature.hash, 32);
        lastReceived.signature.timestamp = msg.signature.timestamp;
        lastReceived.verified = isValid;
        lastReceived.retryCount = msg.retryCount;
        
        return isValid;
    }
    
    void sendConfirmation(bool verified, int maxRetries = 3) {
        if(verified) {
            String ack = JSONGenerator::generateConfirmation(true, msgCounter);
            Serial.println(F("\n[ACK] Sending: "));
            Serial.print(F("      "));
            Serial.println(ack);
            Serial.println(F("      ✅ Verified successfully!"));
            failCount = 0;
        } else {
            if(failCount < maxRetries) {
                failCount++;
                String retry = JSONGenerator::generateRetryRequest(msgCounter);
                Serial.println(F("\n[RETRY] Requesting: "));
                Serial.print(F("        "));
                Serial.println(retry);
                Serial.print(F("        Attempt "));
                Serial.print(failCount);
                Serial.print(F("/"));
                Serial.println(maxRetries);
                
                delay(1000);
                
                Serial.println(F("\n[RETRY] Resending..."));
                Message retryMsg = transmit(lastSent.json, false);
                bool retryVerified = receiveAndVerify(retryMsg);
                sendConfirmation(retryVerified, maxRetries);
            } else {
                Serial.println(F("\n[ERROR] Max retries exceeded!"));
                Serial.println(F("        Manual intervention required"));
                failCount = 0;
            }
        }
    }
    
    void runCycle(bool forceCorrupt = false) {
        String data = JSONGenerator::generateSensorData(msgCounter + 1);
        bool corrupt = forceCorrupt ? true : (random(100) < 30);
        
        Message sent = transmit(data, corrupt);
        bool verified = receiveAndVerify(sent);
        sendConfirmation(verified);
    }
    
    void demo(int cycles = 5) {
        Serial.println(F("   SECURE TRANSCEIVER DEMO (GOST emulation)"));
        Serial.println(F("Mode: Random corruption 30%"));
        Serial.println(F("Max retries: 3"));
        
        for(int i = 0; i < cycles; i++) {
            Serial.print(F("CYCLE "));
            Serial.print(i + 1);
            Serial.print(F("/"));
            Serial.println(cycles);
            
            runCycle();
            delay(1500);
        }
        
        Serial.println(F("   DEMO COMPLETED"));
    }
    
private:
    String corruptData(const String& original) {
        if(original.length() == 0) return original;
        String corrupted = original;
        int pos = random(original.length());
        corrupted.setCharAt(pos, 'X');
        return corrupted;
    }
};

// ========== MAIN ==========
SecureTransceiver transceiver;

void setup() {
    Serial.begin(115200);
    delay(1000);
    
    Serial.println(F("\n"));
    Serial.println(F("║     SECURE DATA EXCHANGE SYSTEM              ║"));
    Serial.println(F("║     GOST R 34.10-2012 (emulation)            ║"));
    
    Serial.println(F("\n[INIT] Crypto subsystem ready"));
    Serial.println(F("[INIT] LightGOST active (demo mode)"));
    Serial.println(F("[INIT] JSON generator ready"));
    Serial.println(F("[INIT] Starting demo...\n"));
    
    delay(500);
}

void loop() {
    transceiver.demo(5);
    
    Serial.println(F("\n[WAIT] 10 seconds until next session..."));
    Serial.println(F("[WAIT] Press RESET in Wokwi to restart\n"));
    delay(10000);
}