#include <Wire.h>
#include <LiquidCrystal_I2C.h>
#include <TinyGPS++.h>
#include "mbedtls/aes.h"

// ==========================================
// 1. KHAI BÁO 3 MÀN HÌNH ĐỘC LẬP (Đã đổi địa chỉ I2C)
// ==========================================
LiquidCrystal_I2C lcd_tau(0x27, 16, 2);    // Màn ở giữa (Tàu)
LiquidCrystal_I2C lcd_tram(0x26, 16, 2);   // Màn bên trái (Trạm)
LiquidCrystal_I2C lcd_hacker(0x28, 16, 2); // Màn bên phải (Hacker)

// ==========================================
// 2. KHAI BÁO PIN ĐIỀU KHIỂN (Chuẩn theo JSON mới)
// ==========================================
// Nhóm Trạm & Hacker (Kéo mượn sang Tàu)
#define BTN_TRAM 4
#define BTN_HACKER 15
#define SW_JAMMING 13

// Nhóm linh kiện của Tàu
#define GPS_RX 35
#define TRIG_PIN 5
#define ECHO_PIN 18
#define LINE_L 32
#define LINE_R 33
#define LED_L 26
#define LED_R 14
#define BUZZER 19

// ==========================================
// 3. BIẾN TOÀN CỤC & GIẢ LẬP LORA
// ==========================================
unsigned char simulated_lora_packet[16];
bool is_packet_flying = false; // Báo hiệu có sóng bay trong không trung
bool is_jamming_active = false; // Báo hiệu đang bị phá sóng
bool is_train_running = false;

unsigned char aes_key[16] = "UIT-RAILWAY-SEC"; 
uint32_t tram_counter = 0;
uint32_t tau_last_counter = 0;
unsigned long tau_last_receive_time = 0;

HardwareSerial SerialGPS(1);
TinyGPSPlus gps;

// ==========================================
// 4. CÁC HÀM MÃ HÓA AES-128
// ==========================================
void encrypt_aes(unsigned char *plainText, unsigned char *cipherText) {
  mbedtls_aes_context aes; mbedtls_aes_init(&aes);
  mbedtls_aes_setkey_enc(&aes, aes_key, 128);
  mbedtls_aes_crypt_ecb(&aes, MBEDTLS_AES_ENCRYPT, plainText, cipherText);
  mbedtls_aes_free(&aes);
}

void decrypt_aes(unsigned char *cipherText, unsigned char *plainText) {
  mbedtls_aes_context aes; mbedtls_aes_init(&aes);
  mbedtls_aes_setkey_dec(&aes, aes_key, 128);
  mbedtls_aes_crypt_ecb(&aes, MBEDTLS_AES_DECRYPT, cipherText, plainText);
  mbedtls_aes_free(&aes);
}

// ==========================================
// LUỒNG 1: TRẠM OCC (Gửi lệnh hợp lệ)
// ==========================================
void Task_Tram(void *pvParameters) {
  for(;;) {
    if (digitalRead(BTN_TRAM) == LOW && !is_jamming_active) {
      tram_counter++;
      uint32_t timestamp = millis() / 1000;
      
      unsigned char plaintext[16] = {0};
      plaintext[0] = (tram_counter >> 24) & 0xFF; plaintext[1] = (tram_counter >> 16) & 0xFF;
      plaintext[2] = (tram_counter >> 8) & 0xFF;  plaintext[3] = tram_counter & 0xFF;
      plaintext[4] = (timestamp >> 24) & 0xFF;    plaintext[5] = (timestamp >> 16) & 0xFF;
      plaintext[6] = (timestamp >> 8) & 0xFF;     plaintext[7] = timestamp & 0xFF;
      plaintext[8] = 1; // Lệnh chạy

      encrypt_aes(plaintext, simulated_lora_packet);
      is_packet_flying = true; // Phát sóng!

      lcd_tram.clear(); lcd_tram.print(">> DA PHAT LENH");
      lcd_tram.setCursor(0, 1); lcd_tram.print("Ma: ");
      for(int i=0; i<4; i++) { if(simulated_lora_packet[i]<16) lcd_tram.print("0"); lcd_tram.print(simulated_lora_packet[i], HEX); }
      vTaskDelay(1000 / portTICK_PERIOD_MS); // Chống dội phím
      lcd_tram.clear(); lcd_tram.print("OCC: San Sang");
    }
    vTaskDelay(50 / portTICK_PERIOD_MS);
  }
}

// ==========================================
// LUỒNG 2: HACKER (Nghe lén, Replay, Jamming)
// ==========================================
void Task_Hacker(void *pvParameters) {
  bool has_sniffed_data = false;
  unsigned char sniffed_buffer[16];

  for(;;) {
    // 1. Phá sóng (Jamming) - Gạt switch sang phía 3V3
    if (digitalRead(SW_JAMMING) == LOW) {
      is_jamming_active = true;
      lcd_hacker.setCursor(0, 0); lcd_hacker.print("!! JAMMING !!   ");
      lcd_hacker.setCursor(0, 1); lcd_hacker.print("Xa rac 433MHz...");
    } else {
      is_jamming_active = false;
      
      // 2. Nghe lén (Sniffing)
      if (is_packet_flying) {
        for(int i=0; i<16; i++) sniffed_buffer[i] = simulated_lora_packet[i];
        has_sniffed_data = true;
        lcd_hacker.clear(); lcd_hacker.print("Bat duoc tin!");
        lcd_hacker.setCursor(0, 1); lcd_hacker.print("San sang Replay");
        vTaskDelay(1500 / portTICK_PERIOD_MS);
        lcd_hacker.clear(); lcd_hacker.print("Radar: Dang cho");
      }

      // 3. Tấn công phát lại (Replay Attack)
      if (digitalRead(BTN_HACKER) == LOW && has_sniffed_data) {
        for(int i=0; i<16; i++) simulated_lora_packet[i] = sniffed_buffer[i];
        is_packet_flying = true; // Bắn lại gói cũ
        lcd_hacker.clear(); lcd_hacker.print(">> REPLAY ATTACK");
        vTaskDelay(1000 / portTICK_PERIOD_MS);
        lcd_hacker.clear(); lcd_hacker.print("Radar: Dang cho");
      }
    }
    vTaskDelay(50 / portTICK_PERIOD_MS);
  }
}

// ==========================================
// LUỒNG 3: TÀU (Bảo mật, Cảm biến & Động cơ)
// ==========================================
void Task_Tau(void *pvParameters) {
  for(;;) {
    // 1. FAILSAFE CHỐNG JAMMING
    if (is_jamming_active) {
      is_train_running = false;
      lcd_tau.setCursor(0, 0); lcd_tau.print("Mat Song->Phanh ");
      digitalWrite(LED_L, LOW); digitalWrite(LED_R, LOW);
      digitalWrite(BUZZER, HIGH); // Hú còi báo động đỏ
    } 
    else {
      digitalWrite(BUZZER, LOW); // Tắt còi nếu hết nhiễu
      
      // 2. NHẬN LỆNH & GIẢI MÃ
      if (is_packet_flying) {
        is_packet_flying = false; // Xác nhận đã bú xong gói tin
        unsigned char decrypted[16];
        decrypt_aes(simulated_lora_packet, decrypted);

        uint32_t rcv_counter = (decrypted[0]<<24) | (decrypted[1]<<16) | (decrypted[2]<<8) | decrypted[3];
        uint32_t rcv_time = (decrypted[4]<<24) | (decrypted[5]<<16) | (decrypted[6]<<8) | decrypted[7];
        uint32_t current_time = millis() / 1000;

        // KIỂM TRA BẢO MẬT: Phân biệt Lệnh thật và Lệnh Hack
        if (rcv_counter <= tau_last_counter || abs((long)(current_time - rcv_time)) > 3) {
          // Bắt được Hacker!
          lcd_tau.clear(); lcd_tau.print("! PHAT HIEN HACK");
          lcd_tau.setCursor(0, 1); lcd_tau.print("Khoa dong co!   ");
          digitalWrite(BUZZER, HIGH); vTaskDelay(800 / portTICK_PERIOD_MS); digitalWrite(BUZZER, LOW);
          is_train_running = false;
        } else {
          // Lệnh chuẩn từ OCC
          tau_last_counter = rcv_counter;
          tau_last_receive_time = millis();
          is_train_running = true;
          lcd_tau.clear(); lcd_tau.print("Lenh Hop Le: OK ");
        }
      }

      // 3. ĐIỀU KHIỂN CHẠY TÀU & QUÉT VẬT CẢN NẾU ĐƯỢC CHẠY
      if (is_train_running) {
        // Bắn siêu âm đo khoảng cách
        digitalWrite(TRIG_PIN, LOW); delayMicroseconds(2);
        digitalWrite(TRIG_PIN, HIGH); delayMicroseconds(10);
        digitalWrite(TRIG_PIN, LOW);
        long duration = pulseIn(ECHO_PIN, HIGH, 30000); 
        int distance = duration * 0.034 / 2;

        if (distance > 0 && distance < 15) {
          // Gặp vật cản < 15cm -> Phanh gấp
          digitalWrite(LED_L, LOW); digitalWrite(LED_R, LOW);
          lcd_tau.setCursor(0, 1); lcd_tau.print("Vat can! Phanh! ");
          digitalWrite(BUZZER, HIGH); vTaskDelay(100 / portTICK_PERIOD_MS); digitalWrite(BUZZER, LOW);
        } else {
          // Đường thoáng -> Động cơ chạy (Sáng 2 LED)
          digitalWrite(LED_L, HIGH); digitalWrite(LED_R, HIGH);
          lcd_tau.setCursor(0, 1); lcd_tau.print("Tau Dang Chay...");
        }
      } else {
        digitalWrite(LED_L, LOW); digitalWrite(LED_R, LOW);
        if(!is_jamming_active) { lcd_tau.setCursor(0, 1); lcd_tau.print("Tau Dang Dung..."); }
      }
    }
    vTaskDelay(50 / portTICK_PERIOD_MS);
  }
}

// ==========================================
// HÀM MAIN SETUP
// ==========================================
void setup() {
  Serial.begin(115200);
  SerialGPS.begin(9600, SERIAL_8N1, GPS_RX, -1);
  
  // Set PinMode
  pinMode(BTN_TRAM, INPUT_PULLUP);
  pinMode(BTN_HACKER, INPUT_PULLUP);
  pinMode(SW_JAMMING, INPUT_PULLUP);
  pinMode(TRIG_PIN, OUTPUT); pinMode(ECHO_PIN, INPUT);
  pinMode(LED_L, OUTPUT); pinMode(LED_R, OUTPUT);
  pinMode(BUZZER, OUTPUT);
  
  // Khởi động 3 màn hình cùng lúc
  lcd_tau.init(); lcd_tau.backlight(); lcd_tau.print("TAU: KHOI DONG");
  lcd_tram.init(); lcd_tram.backlight(); lcd_tram.print("OCC: KHOI DONG");
  lcd_hacker.init(); lcd_hacker.backlight(); lcd_hacker.print("HACK: KHOI DONG");
  
  delay(1500);
  lcd_tau.clear(); lcd_tram.clear(); lcd_hacker.clear();
  lcd_tram.print("OCC: San Sang");
  lcd_hacker.print("Radar: Dang cho");
  lcd_tau.print("Tau: Cho Lenh");

  // Đẩy 3 tác vụ chạy song song trên FreeRTOS
  xTaskCreate(Task_Tram, "Tram", 2048, NULL, 1, NULL);
  xTaskCreate(Task_Hacker, "Hacker", 2048, NULL, 1, NULL);
  xTaskCreate(Task_Tau, "Tau", 2048, NULL, 1, NULL);
}

void loop() {
  // Chỉ dùng loop() để nuốt dữ liệu GPS, không làm kẹt FreeRTOS
  while (SerialGPS.available() > 0) { gps.encode(SerialGPS.read()); }
}