#include <Adafruit_GFX.h>
#include <Adafruit_ILI9341.h>
#include <XPT2046_Touchscreen.h>
#include <WiFi.h>
#include <SPI.h>
extern "C" {
    #include "esp_wifi.h"
}

// ILI9341 Display pins (SPI)
#define TFT_CS     15   // CS pin
#define TFT_DC     2    // DC pin  
#define TFT_RST    4    // RST pin
#define TFT_MOSI   23   // MOSI pin (SDA for display)
#define TFT_CLK    18   // SCK pin
#define TFT_MISO   19   // MISO pin
#define TFT_LED    5    // LED backlight pin (optional)

// Touch pins (I2C/SPI shared)
#define TOUCH_CS   21   // Touch CS (separate from display CS)
// Touch uses same MOSI, MISO, SCK as display
// SCL and SDA are for I2C touch (if using I2C mode)

// Display dimensions
#define SCREEN_WIDTH  320
#define SCREEN_HEIGHT 240

// Display and touch setup
Adafruit_ILI9341 tft = Adafruit_ILI9341(TFT_CS, TFT_DC, TFT_MOSI, TFT_CLK, TFT_RST, TFT_MISO);
XPT2046_Touchscreen touch(TOUCH_CS);  // Touch shares SPI bus with display

// Touch calibration (adjust these for your display)
#define TOUCH_MIN_X 200
#define TOUCH_MAX_X 3700
#define TOUCH_MIN_Y 240
#define TOUCH_MAX_Y 3800

// Colors
#define COLOR_BACKGROUND  ILI9341_BLACK
#define COLOR_TEXT        ILI9341_WHITE
#define COLOR_HIGHLIGHT   ILI9341_CYAN
#define COLOR_BUTTON      ILI9341_DARKGREY
#define COLOR_BUTTON_TEXT ILI9341_WHITE
#define COLOR_ACTIVE      ILI9341_GREEN
#define COLOR_INACTIVE    ILI9341_RED

// RF Colors
#define COLOR_WEAK        ILI9341_BLUE
#define COLOR_MEDIUM      ILI9341_GREEN  
#define COLOR_STRONG      ILI9341_YELLOW
#define COLOR_VERY_STRONG ILI9341_RED

// Menu states
enum MenuState {
    MENU_MAIN,
    MENU_WATERFALL,
    MENU_CHANNEL_SCAN,
    MENU_SETTINGS
};

// Touch button structure
struct TouchButton {
    int x, y, w, h;
    const char* label;
    uint16_t color;
    bool active;
};

// Waterfall configuration
#define CHANNELS 14
#define WATERFALL_HEIGHT 180
#define WATERFALL_WIDTH 280
#define SCAN_TIME_MS 80

// Global variables
static MenuState currentMenu = MENU_MAIN;
static int8_t channelData[CHANNELS];
static uint16_t waterfallBuffer[WATERFALL_HEIGHT][CHANNELS];
static int waterfallLine = 0;
static unsigned long lastScanTime = 0;
static unsigned long lastTouchTime = 0;
static int selectedChannel = 1;
static bool isScanning = false;

// Touch buttons for main menu
TouchButton mainMenuButtons[] = {
    {20, 40, 120, 40, "Waterfall", COLOR_BUTTON, false},
    {160, 40, 120, 40, "Channel", COLOR_BUTTON, false},
    {20, 100, 120, 40, "Settings", COLOR_BUTTON, false},
    {160, 100, 120, 40, "About", COLOR_BUTTON, false}
};

// Control buttons for waterfall
TouchButton waterfallButtons[] = {
    {285, 20, 30, 30, "<", COLOR_BUTTON, false},    // Back
    {285, 60, 30, 30, "||", COLOR_BUTTON, false},   // Pause
    {285, 100, 30, 30, "Ch-", COLOR_BUTTON, false}, // Channel down
    {285, 140, 30, 30, "Ch+", COLOR_BUTTON, false}  // Channel up
};

// Function prototypes
void initDisplay();
void initWiFi();
void drawMainMenu();
void drawWaterfallScreen();
void drawChannelScanScreen();
void handleTouch();
TS_Point getTouchPoint();
bool isTouchInButton(TS_Point point, TouchButton &btn);
void drawButton(TouchButton &btn);
void scanChannel(int channel);
void updateWaterfall();
void drawWaterfall();
uint16_t rssiToColor(int8_t rssi);

void setup() {
    Serial.begin(115200);
    
    initDisplay();
    initWiFi();
    
    // Initialize channel data
    for (int i = 0; i < CHANNELS; i++) {
        channelData[i] = -100;
    }
    
    // Clear waterfall buffer
    for (int y = 0; y < WATERFALL_HEIGHT; y++) {
        for (int x = 0; x < CHANNELS; x++) {
            waterfallBuffer[y][x] = COLOR_BACKGROUND;
        }
    }
    
    drawMainMenu();
    Serial.println("WiFi RF Analyzer Ready");
}

void loop() {
    unsigned long now = millis();
    
    // Handle touch input (debounced)
    if (now - lastTouchTime > 200) {
        handleTouch();
        lastTouchTime = now;
    }
    
    // Handle scanning based on current menu
    if (currentMenu == MENU_WATERFALL && isScanning) {
        if (now - lastScanTime >= SCAN_TIME_MS) {
            static int scanChannel = 1;
            
            scanChannel(scanChannel);
            scanChannel++;
            
            if (scanChannel > CHANNELS) {
                scanChannel = 1;
                updateWaterfall();
                drawWaterfall();
            }
            
            lastScanTime = now;
        }
    }
    
    delay(10);
}

void initDisplay() {
    // Initialize backlight (optional)
    pinMode(TFT_LED, OUTPUT);
    digitalWrite(TFT_LED, HIGH); // Turn on backlight
    
    tft.begin();
    tft.setRotation(3); // Landscape
    tft.fillScreen(COLOR_BACKGROUND);
    
    touch.begin();
    touch.setRotation(1);
    
    // Display startup message
    tft.setTextColor(COLOR_TEXT);
    tft.setTextSize(2);
    tft.setCursor(80, 100);
    tft.print("WiFi RF Analyzer");
    tft.setTextSize(1);
    tft.setCursor(120, 130);
    tft.print("Initializing...");
    
    delay(2000);
    tft.fillScreen(COLOR_BACKGROUND);
}

void initWiFi() {
    WiFi.mode(WIFI_STA);
    WiFi.disconnect();
    esp_wifi_set_promiscuous(false);
}

void drawMainMenu() {
    currentMenu = MENU_MAIN;
    tft.fillScreen(COLOR_BACKGROUND);
    
    // Title
    tft.setTextColor(COLOR_TEXT);
    tft.setTextSize(2);
    tft.setCursor(80, 10);
    tft.print("RF Analyzer Menu");
    
    // Draw menu buttons
    for (int i = 0; i < 4; i++) {
        drawButton(mainMenuButtons[i]);
    }
    
    // Instructions
    tft.setTextSize(1);
    tft.setTextColor(COLOR_TEXT);
    tft.setCursor(20, 200);
    tft.print("Touch buttons to navigate");
}

void drawWaterfallScreen() {
    currentMenu = MENU_WATERFALL;
    isScanning = true;
    tft.fillScreen(COLOR_BACKGROUND);
    
    // Title and info area
    tft.setTextColor(COLOR_TEXT);
    tft.setTextSize(1);
    tft.setCursor(5, 5);
    tft.print("2.4GHz Waterfall - Ch:");
    tft.print(selectedChannel);
    
    // Channel labels
    int channelWidth = WATERFALL_WIDTH / CHANNELS;
    for (int ch = 1; ch <= CHANNELS; ch++) {
        int x = (ch - 1) * channelWidth + 5;
        tft.setCursor(x + channelWidth/2 - 3, 20);
        if (ch == selectedChannel) {
            tft.setTextColor(COLOR_HIGHLIGHT);
        } else {
            tft.setTextColor(COLOR_TEXT);
        }
        tft.print(ch);
    }
    
    // Draw control buttons
    for (int i = 0; i < 4; i++) {
        drawButton(waterfallButtons[i]);
    }
    
    // Status area
    tft.setTextColor(COLOR_TEXT);
    tft.setCursor(5, 220);
    tft.print("Status: Scanning...");
}

void drawChannelScanScreen() {
    currentMenu = MENU_CHANNEL_SCAN;
    tft.fillScreen(COLOR_BACKGROUND);
    
    // Title
    tft.setTextColor(COLOR_TEXT);
    tft.setTextSize(2);
    tft.setCursor(50, 20);
    tft.print("Channel Scanner");
    
    // Show all channels with signal strength
    tft.setTextSize(1);
    for (int ch = 1; ch <= CHANNELS; ch++) {
        int x = 20 + ((ch - 1) % 7) * 40;
        int y = 60 + ((ch - 1) / 7) * 40;
        
        // Channel box
        uint16_t boxColor = rssiToColor(channelData[ch - 1]);
        tft.fillRect(x, y, 35, 30, boxColor);
        tft.drawRect(x, y, 35, 30, COLOR_TEXT);
        
        // Channel number
        tft.setTextColor(COLOR_TEXT);
        tft.setCursor(x + 8, y + 5);
        tft.print(ch);
        
        // RSSI value
        tft.setCursor(x + 2, y + 18);
        if (channelData[ch - 1] > -100) {
            tft.print(channelData[ch - 1]);
        } else {
            tft.print("---");
        }
    }
    
    // Back button
    TouchButton backBtn = {250, 200, 60, 30, "Back", COLOR_BUTTON, false};
    drawButton(backBtn);
    
    // Legend
    tft.setTextColor(COLOR_TEXT);
    tft.setCursor(20, 200);
    tft.print("Colors: Blue=Weak, Green=Med, Yellow=Strong, Red=Very Strong");
}

void handleTouch() {
    if (!touch.touched()) return;
    
    TS_Point point = getTouchPoint();
    if (point.x == 0 && point.y == 0) return; // Invalid touch
    
    switch (currentMenu) {
        case MENU_MAIN:
            if (isTouchInButton(point, mainMenuButtons[0])) {
                drawWaterfallScreen();
            } else if (isTouchInButton(point, mainMenuButtons[1])) {
                // Quick scan all channels first
                for (int ch = 1; ch <= CHANNELS; ch++) {
                    scanChannel(ch);
                    delay(100);
                }
                drawChannelScanScreen();
            } else if (isTouchInButton(point, mainMenuButtons[2])) {
                // Settings - placeholder
                tft.fillScreen(COLOR_BACKGROUND);
                tft.setTextColor(COLOR_TEXT);
                tft.setCursor(100, 100);
                tft.print("Settings Menu");
                tft.setCursor(100, 120);
                tft.print("(Coming Soon)");
                delay(2000);
                drawMainMenu();
            }
            break;
            
        case MENU_WATERFALL:
            if (isTouchInButton(point, waterfallButtons[0])) { // Back
                isScanning = false;
                drawMainMenu();
            } else if (isTouchInButton(point, waterfallButtons[1])) { // Pause
                isScanning = !isScanning;
                waterfallButtons[1].label = isScanning ? "||" : ">";
                drawButton(waterfallButtons[1]);
            } else if (isTouchInButton(point, waterfallButtons[2])) { // Ch-
                if (selectedChannel > 1) {
                    selectedChannel--;
                    drawWaterfallScreen();
                }
            } else if (isTouchInButton(point, waterfallButtons[3])) { // Ch+
                if (selectedChannel < CHANNELS) {
                    selectedChannel++;
                    drawWaterfallScreen();
                }
            }
            break;
            
        case MENU_CHANNEL_SCAN:
            // Check if back button touched
            TouchButton backBtn = {250, 200, 60, 30, "Back", COLOR_BUTTON, false};
            if (isTouchInButton(point, backBtn)) {
                drawMainMenu();
            }
            break;
    }
}

TS_Point getTouchPoint() {
    TS_Point p = touch.getPoint();
    
    // Map touch coordinates to screen coordinates
    int x = map(p.x, TOUCH_MIN_X, TOUCH_MAX_X, 0, SCREEN_WIDTH);
    int y = map(p.y, TOUCH_MIN_Y, TOUCH_MAX_Y, 0, SCREEN_HEIGHT);
    
    // Clamp to screen bounds
    x = constrain(x, 0, SCREEN_WIDTH - 1);
    y = constrain(y, 0, SCREEN_HEIGHT - 1);
    
    return TS_Point(x, y, p.z);
}

bool isTouchInButton(TS_Point point, TouchButton &btn) {
    return (point.x >= btn.x && point.x <= btn.x + btn.w &&
            point.y >= btn.y && point.y <= btn.y + btn.h);
}

void drawButton(TouchButton &btn) {
    // Button background
    tft.fillRect(btn.x, btn.y, btn.w, btn.h, btn.color);
    tft.drawRect(btn.x, btn.y, btn.w, btn.h, COLOR_TEXT);
    
    // Button text
    tft.setTextColor(COLOR_BUTTON_TEXT);
    tft.setTextSize(1);
    
    // Center text in button
    int textWidth = strlen(btn.label) * 6; // Approximate width
    int textX = btn.x + (btn.w - textWidth) / 2;
    int textY = btn.y + (btn.h - 8) / 2;
    
    tft.setCursor(textX, textY);
    tft.print(btn.label);
}

void scanChannel(int channel) {
    esp_wifi_set_channel(channel, WIFI_SECOND_CHAN_NONE);
    channelData[channel - 1] = -100;
    
    delay(20); // Channel settle time
    
    int networkCount = WiFi.scanNetworks(false, false, false, 300, channel);
    for (int i = 0; i < networkCount; i++) {
        int8_t rssi = WiFi.RSSI(i);
        if (rssi > channelData[channel - 1]) {
            channelData[channel - 1] = rssi;
        }
    }
    WiFi.scanDelete();
}

void updateWaterfall() {
    // Add new line to waterfall buffer
    for (int x = 0; x < CHANNELS; x++) {
        waterfallBuffer[waterfallLine][x] = rssiToColor(channelData[x]);
    }
    
    waterfallLine = (waterfallLine + 1) % WATERFALL_HEIGHT;
}

void drawWaterfall() {
    int channelWidth = WATERFALL_WIDTH / CHANNELS;
    int startX = 5;
    int startY = 35;
    
    // Draw waterfall
    for (int y = 0; y < WATERFALL_HEIGHT; y++) {
        int bufferY = (waterfallLine + y) % WATERFALL_HEIGHT;
        
        for (int ch = 0; ch < CHANNELS; ch++) {
            int x = startX + ch * channelWidth;
            uint16_t color = waterfallBuffer[bufferY][ch];
            
            tft.drawFastHLine(x, startY + y, channelWidth, color);
        }
    }
    
    // Update status
    tft.fillRect(80, 220, 100, 10, COLOR_BACKGROUND);
    tft.setTextColor(COLOR_TEXT);
    tft.setTextSize(1);
    tft.setCursor(80, 220);
    tft.printf("Ch%d: %ddBm", selectedChannel, channelData[selectedChannel - 1]);
}

uint16_t rssiToColor(int8_t rssi) {
    if (rssi <= -90) return COLOR_BACKGROUND;
    else if (rssi <= -80) return COLOR_WEAK;
    else if (rssi <= -70) return COLOR_MEDIUM;
    else if (rssi <= -60) return COLOR_STRONG;
    else return COLOR_VERY_STRONG;
}