#include <Wire.h>
#include <LiquidCrystal_I2C.h>
#include <Keypad.h>
#include <Adafruit_GFX.h>
#include <Adafruit_ILI9341.h>
#include <math.h>

// ================= LCD =================
LiquidCrystal_I2C lcd(0x27, 16, 2);

// ================= TFT =================
#define TFT_CS   5
#define TFT_DC   2
#define TFT_RST  4
Adafruit_ILI9341 tft(TFT_CS, TFT_DC, TFT_RST);

// ================= RGB =================
#define LED_R 16
#define LED_G 17
#define LED_B 15

// ================= BUZZER ==============
#define BUZZ 0

// ================= KEYPAD ==============
const byte ROWS = 4, COLS = 4;
char keys[ROWS][COLS] = {
  {'1','2','3','A'},
  {'4','5','6','B'},
  {'7','8','9','C'},
  {'*','0','#','D'}
};
byte rowPins[ROWS] = {13,12,14,27};
byte colPins[COLS] = {26,25,33,32};
Keypad keypad(makeKeymap(keys), rowPins, colPins, ROWS, COLS);

// ================= SYSTEM STATES =======
enum SystemState { STANDBY, INPUT_MODE, JAMMING, SPECTRUM_ANALYSIS };
SystemState currentState = STANDBY;

// ================= TEXT =================
String typedText = "";
String confirmedText = "";
String targetFreq = "2412";
String detectedDevices[5] = {"WiFi AP1", "WiFi AP2", "BT Device", "Unknown", "WiFi AP3"};
int deviceCount = 5;
int scrollPos = 0;
unsigned long lastScroll = 0;

// ================= JAMMER PARAMS =======
bool jammerActive = false;
int signalStrength = 0;
int freqIndex = 0;
float noiseLevel = 0.0;
unsigned long lastUpdate = 0;
unsigned long startTime = 0;
unsigned long batteryTimer = 0;
int batteryLevel = 98;

// ================= SPECTRUM DATA =======
float spectrumData[64];
float waterfallData[32][64];
int waterfallPos = 0;

// ================= CUSTOM COLORS =======
#define ILI9341_GRAY    0x7BEF
#define ILI9341_DARKGRAY 0x4208
#define ILI9341_NAVY    0x000F
#define ILI9341_DARKGREEN 0x03E0
#define ILI9341_DARKCYAN 0x03EF
#define ILI9341_MAROON  0x7800
#define ILI9341_PURPLE  0x780F
#define ILI9341_OLIVE   0x7BE0
#define ILI9341_LIGHTGRAY 0xC618
#define ILI9341_DARKRED 0x8000

// ================= SETUP ================
void setup() {
  Serial.begin(115200);
  Wire.begin(21, 22);

  lcd.init();
  lcd.backlight();
  lcd.print("SYSTEM INIT...");
  delay(500);
  lcd.clear();
  lcd.print("FREQ SCANNER v2.1");
  lcd.setCursor(0, 1);
  lcd.print("MODEL: RF-JM4");

  pinMode(LED_R, OUTPUT);
  pinMode(LED_G, OUTPUT);
  pinMode(LED_B, OUTPUT);
  pinMode(BUZZ, OUTPUT);

  tft.begin();
  tft.setRotation(1);
  tft.fillScreen(ILI9341_BLACK);
  
  // Initialize spectrum data with random noise
  for (int i = 0; i < 64; i++) {
    spectrumData[i] = random(10, 40);
  }
  
  // Initialize waterfall
  for (int y = 0; y < 32; y++) {
    for (int x = 0; x < 64; x++) {
      waterfallData[y][x] = random(5, 30);
    }
  }
  
  drawMainUI();
  beepStartup();
}

// ================= MAIN LOOP ============
void loop() {
  char key = keypad.getKey();
  
  if (key) {
    handleKeypad(key);
    beepKey();
  }
  
  switch (currentState) {
    case STANDBY:
      updateStandbyMode();
      break;
    case INPUT_MODE:
      updateInputMode();
      break;
    case JAMMING:
      updateJammingMode();
      break;
    case SPECTRUM_ANALYSIS:
      updateSpectrumMode();
      break;
  }
  
  updateBattery();
  updateStatusLED();
}

// ================= KEYPAD HANDLER =======
void handleKeypad(char key) {
  switch (key) {
    case 'A':
      currentState = SPECTRUM_ANALYSIS;
      tft.fillScreen(ILI9341_BLACK);
      drawSpectrumUI();
      break;
      
    case 'B':
      if (currentState != JAMMING) {
        currentState = INPUT_MODE;
        typedText = targetFreq;
        drawFrequencyUI();
      }
      break;
      
    case 'C':
      if (currentState == JAMMING) {
        jammerActive = false;
        currentState = STANDBY;
        tft.fillScreen(ILI9341_BLACK);
        drawMainUI();
        lcd.clear();
        lcd.print("JAMMING STOPPED");
        lcd.setCursor(0, 1);
        lcd.print("SYSTEM STANDBY");
      }
      break;
      
    case 'D':
      if (typedText.length() > 0 && !typedText.equals(targetFreq)) {
        targetFreq = typedText;
        lcd.clear();
        lcd.print("TARGET SET: ");
        lcd.print(targetFreq);
        lcd.setCursor(0, 1);
        lcd.print("MHz");
      }
      break;
      
    case '#':
      if (typedText.length() == 4 && !jammerActive) {
        currentState = JAMMING;
        jammerActive = true;
        startTime = millis();
        signalStrength = 75;
        tft.fillScreen(ILI9341_BLACK);
        drawJammingUI();
        lcd.clear();
        lcd.print("ACTIVE JAM: ");
        lcd.print(targetFreq);
        lcd.setCursor(0, 1);
        lcd.print("SIG: ");
        lcd.print(signalStrength);
        lcd.print("%");
      }
      break;
      
    case '*':
      if (typedText.length() > 0) {
        typedText.remove(typedText.length() - 1);
      }
      break;
      
    default:
      if (isdigit(key) && typedText.length() < 4) {
        typedText += key;
      }
      break;
  }
  
  if (currentState == INPUT_MODE) {
    updateFrequencyDisplay();
  }
}

// ================= UI DRAWING ===========
void drawMainUI() {
  tft.setTextColor(ILI9341_CYAN);
  tft.setTextSize(2);
  tft.setCursor(40, 10);
  tft.print("RF SCANNER");
  
  tft.setTextColor(ILI9341_GREEN);
  tft.setTextSize(1);
  tft.setCursor(10, 40);
  tft.print("STATUS: SYSTEM READY");
  
  tft.setCursor(10, 60);
  tft.print("TARGET FREQ: ");
  tft.print(targetFreq);
  tft.print(" MHz");
  
  tft.setCursor(10, 80);
  tft.print("BATTERY: ");
  tft.print(batteryLevel);
  tft.print("%");
  
  // Draw signal strength indicator
  tft.drawRect(10, 110, 150, 20, ILI9341_DARKGRAY);
  tft.fillRect(12, 112, 146, 16, ILI9341_BLACK);
  tft.fillRect(12, 112, random(30, 146), 16, ILI9341_GREEN);
  
  tft.setTextSize(1);
  tft.setCursor(10, 135);
  tft.print("SIGNAL: ");
  tft.print(random(20, 95));
  tft.print("%");
  
  // Draw control labels
  tft.setTextColor(ILI9341_YELLOW);
  tft.setCursor(10, 180);
  tft.print("[A] SPECTRUM  [B] SET FREQ");
  tft.setCursor(10, 200);
  tft.print("[C] STOP      [#] ACTIVATE");
}

void drawFrequencyUI() {
  tft.fillScreen(ILI9341_BLACK);
  tft.setTextColor(ILI9341_YELLOW);
  tft.setTextSize(2);
  tft.setCursor(50, 20);
  tft.print("FREQUENCY");
  tft.setCursor(50, 45);
  tft.print("SETTING");
  
  tft.drawRect(60, 90, 200, 60, ILI9341_WHITE);
  updateFrequencyDisplay();
  
  tft.setTextColor(ILI9341_CYAN);
  tft.setTextSize(1);
  tft.setCursor(10, 180);
  tft.print("ENTER 4-DIGIT FREQ (MHz)");
  tft.setCursor(10, 200);
  tft.print("[D] CONFIRM  [*] DELETE");
}

void drawJammingUI() {
  tft.setTextColor(ILI9341_RED);
  tft.setTextSize(3);
  tft.setCursor(40, 10);
  tft.print("ACTIVE JAM");
  
  tft.setTextColor(ILI9341_WHITE);
  tft.setTextSize(4);
  tft.setCursor(80, 60);
  tft.print(targetFreq);
  
  tft.setTextSize(2);
  tft.setCursor(110, 100);
  tft.print("MHz");
  
  // Draw power meter
  tft.drawRect(20, 140, 280, 30, ILI9341_WHITE);
}

void drawSpectrumUI() {
  tft.setTextColor(ILI9341_CYAN);
  tft.setTextSize(2);
  tft.setCursor(40, 5);
  tft.print("SPECTRUM ANALYZER");
  
  // Draw frequency scale
  tft.setTextColor(ILI9341_GREEN);
  tft.setTextSize(1);
  for (int i = 0; i < 6; i++) {
    int x = 20 + i * 50;
    tft.setCursor(x, 220);
    tft.print(2400 + i * 100);
  }
  
  tft.setCursor(10, 230);
  tft.print("MHz");
}

// ================= MODE UPDATES =========
void updateStandbyMode() {
  static unsigned long lastStandbyUpdate = 0;
  if (millis() - lastStandbyUpdate > 1000) {
    lastStandbyUpdate = millis();
    
    // Update signal strength randomly
    tft.fillRect(12, 112, 146, 16, ILI9341_BLACK);
    int strength = random(30, 146);
    uint16_t color = ILI9341_GREEN;
    if (strength < 50) color = ILI9341_RED;
    else if (strength < 100) color = ILI9341_YELLOW;
    tft.fillRect(12, 112, strength, 16, color);
    
    tft.fillRect(80, 135, 50, 10, ILI9341_BLACK);
    tft.setCursor(10, 135);
    tft.print("SIGNAL: ");
    tft.print(map(strength, 30, 146, 20, 95));
    tft.print("%");
  }
}

void updateInputMode() {
  // Just wait for input
}

void updateJammingMode() {
  static unsigned long lastJammUpdate = 0;
  if (millis() - lastJammUpdate > 100) {
    lastJammUpdate = millis();
    
    // Animate power meter
    int powerWidth = map(sin(millis() / 200.0) * 100 + random(-10, 10), -100, 100, 50, 270);
    tft.fillRect(22, 142, 276, 26, ILI9341_BLACK);
    uint16_t color = ILI9341_GREEN;
    if (powerWidth < 100) color = ILI9341_RED;
    else if (powerWidth < 200) color = ILI9341_YELLOW;
    tft.fillRect(22, 142, powerWidth, 26, color);
    
    // Update signal strength
    signalStrength = constrain(signalStrength + random(-3, 4), 20, 95);
    
    // Update LCD
    lcd.clear();
    lcd.print("JAM: ");
    lcd.print(targetFreq);
    lcd.print("MHz");
    lcd.setCursor(0, 1);
    lcd.print("PWR:");
    lcd.print(map(powerWidth, 50, 270, 0, 100));
    lcd.print("% T:");
    unsigned long elapsed = (millis() - startTime) / 1000;
    lcd.print(elapsed);
    lcd.print("s");
    
    // Occasionally show "device blocked" message
    if (random(100) < 5) {
      lcd.clear();
      lcd.print("DEVICE BLOCKED!");
      lcd.setCursor(0, 1);
      lcd.print(detectedDevices[random(5)]);
      delay(800);
    }
  }
}

void updateSpectrumMode() {
  static unsigned long lastSpectrumUpdate = 0;
  if (millis() - lastSpectrumUpdate > 150) {
    lastSpectrumUpdate = millis();
    
    // Generate fake spectrum data
    for (int i = 0; i < 64; i++) {
      float base = sin(i / 10.0 + millis() / 1000.0) * 15 + 20;
      float spike = (i == freqIndex % 64) ? 40 : 0;
      float noise = random(-5, 5);
      spectrumData[i] = constrain(base + spike + noise, 5, 60);
    }
    
    // Shift waterfall down
    for (int y = 31; y > 0; y--) {
      for (int x = 0; x < 64; x++) {
        waterfallData[y][x] = waterfallData[y-1][x];
      }
    }
    
    // Add new row
    for (int x = 0; x < 64; x++) {
      waterfallData[0][x] = spectrumData[x];
    }
    
    // Draw spectrum
    for (int i = 0; i < 64; i++) {
      int height = map(spectrumData[i], 5, 60, 0, 100);
      int x = 20 + i * 4;
      
      // Clear previous
      tft.drawFastVLine(x, 40, 100, ILI9341_BLACK);
      
      // Draw new bar
      uint16_t color = ILI9341_GREEN;
      if (spectrumData[i] > 45) color = ILI9341_RED;
      else if (spectrumData[i] > 35) color = ILI9341_YELLOW;
      
      tft.drawFastVLine(x, 140 - height, height, color);
    }
    
    // Draw waterfall
    for (int y = 0; y < 32; y++) {
      for (int x = 0; x < 64; x++) {
        int value = waterfallData[y][x];
        uint16_t color = ILI9341_BLACK;
        if (value > 50) color = ILI9341_RED;
        else if (value > 40) color = ILI9341_YELLOW;
        else if (value > 30) color = ILI9341_GREEN;
        else if (value > 20) color = ILI9341_BLUE;
        else if (value > 10) color = ILI9341_NAVY;
        
        tft.drawPixel(20 + x * 4, 150 + y, color);
      }
    }
    
    freqIndex++;
  }
}

// ================= DISPLAY UPDATES ======
void updateFrequencyDisplay() {
  tft.fillRect(62, 92, 196, 56, ILI9341_BLACK);
  tft.setTextColor(ILI9341_WHITE);
  tft.setTextSize(4);
  
  if (typedText.length() == 0) {
    tft.setCursor(120, 105);
    tft.print("____");
  } else {
    tft.setCursor(140 - typedText.length() * 20, 105);
    tft.print(typedText);
    // Draw underscores for missing digits
    for (int i = typedText.length(); i < 4; i++) {
      tft.setCursor(140 + i * 20, 105);
      tft.print("_");
    }
  }
  
  tft.setTextSize(2);
  tft.setCursor(220, 110);
  tft.print("MHz");
}

void updateBattery() {
  static unsigned long lastBatteryUpdate = 0;
  if (millis() - lastBatteryUpdate > 10000) {
    lastBatteryUpdate = millis();
    batteryLevel = constrain(batteryLevel - 1, 10, 100);
    
    if (currentState == STANDBY) {
      tft.fillRect(70, 80, 60, 10, ILI9341_BLACK);
      tft.setCursor(10, 80);
      tft.print("BATTERY: ");
      tft.print(batteryLevel);
      tft.print("%   ");
    }
    
    if (batteryLevel <= 20) {
      lcd.clear();
      lcd.print("LOW BATTERY!");
      lcd.setCursor(0, 1);
      lcd.print(batteryLevel);
      lcd.print("% REMAINING");
      delay(1000);
    }
  }
}

void updateStatusLED() {
  static unsigned long lastLEDUpdate = 0;
  if (millis() - lastLEDUpdate > 500) {
    lastLEDUpdate = millis();
    
    switch (currentState) {
      case STANDBY:
        digitalWrite(LED_G, HIGH);
        digitalWrite(LED_R, LOW);
        digitalWrite(LED_B, LOW);
        break;
      case INPUT_MODE:
        digitalWrite(LED_G, LOW);
        digitalWrite(LED_R, LOW);
        digitalWrite(LED_B, HIGH);
        break;
      case JAMMING:
        digitalWrite(LED_G, LOW);
        digitalWrite(LED_R, HIGH);
        digitalWrite(LED_B, LOW);
        break;
      case SPECTRUM_ANALYSIS:
        digitalWrite(LED_G, LOW);
        digitalWrite(LED_R, LOW);
        digitalWrite(LED_B, !digitalRead(LED_B));
        break;
    }
  }
}

// ================= AUDIO ================
void beepKey() {
  tone(BUZZ, 800, 50);
}

void beepStartup() {
  tone(BUZZ, 1000, 100);
  delay(120);
  tone(BUZZ, 1200, 100);
  delay(120);
  tone(BUZZ, 800, 100);
}

void beepConfirm() {
  tone(BUZZ, 1200, 80);
  delay(100);
  tone(BUZZ, 1400, 80);
}