#include "oled_display.h"

// Global instance
OLEDDisplay oledDisplay;

// Icon definitions (8x8 bitmaps)
const unsigned char PROGMEM icon_wifi_connected[] = {
    0x00, 0x0E, 0x11, 0x04, 0x0A, 0x00, 0x04, 0x00
};

const unsigned char PROGMEM icon_wifi_disconnected[] = {
    0x00, 0x0E, 0x11, 0x04, 0x0A, 0x11, 0x04, 0x11
};

const unsigned char PROGMEM icon_bluetooth[] = {
    0x04, 0x06, 0x15, 0x0E, 0x0E, 0x15, 0x06, 0x04
};

const unsigned char PROGMEM icon_lora[] = {
    0x00, 0x08, 0x14, 0x2A, 0x14, 0x08, 0x00, 0x00
};

const unsigned char PROGMEM icon_zigbee[] = {
    0x00, 0x1C, 0x14, 0x1C, 0x14, 0x1C, 0x00, 0x00
};

const unsigned char PROGMEM icon_attack[] = {
    0x08, 0x1C, 0x2A, 0x49, 0x2A, 0x1C, 0x08, 0x00
};

const unsigned char PROGMEM icon_shield[] = {
    0x08, 0x14, 0x22, 0x41, 0x41, 0x22, 0x14, 0x08
};

const unsigned char PROGMEM icon_warning[] = {
    0x08, 0x14, 0x14, 0x22, 0x22, 0x00, 0x08, 0x00
};

const unsigned char PROGMEM icon_success[] = {
    0x00, 0x01, 0x02, 0x44, 0x28, 0x10, 0x00, 0x00
};

const unsigned char PROGMEM icon_error[] = {
    0x41, 0x22, 0x14, 0x08, 0x14, 0x22, 0x41, 0x00
};

OLEDDisplay::OLEDDisplay() {
    display = nullptr;
    current_page = PAGE_BOOT;
    last_update = 0;
    page_switch_time = 0;
    auto_rotate = true;
    brightness = 255;
    display_enabled = true;
    animation_frame = 0;
    last_animation = 0;
    blink_state = false;
    
    // Initialize status structures
    memset(&system_status, 0, sizeof(system_status));
    memset(&attack_info, 0, sizeof(attack_info));
    memset(&network_info, 0, sizeof(network_info));
    
    attack_info.status = ATTACK_IDLE;
}

OLEDDisplay::~OLEDDisplay() {
    if (display) {
        delete display;
    }
}

bool OLEDDisplay::begin(uint8_t i2c_address) {
    OLED_LOG("Initializing OLED display...");
    
    // Initialize I2C
    Wire.begin();
    
    // Create display object
    display = new Adafruit_SSD1306(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);
    
    if (!display->begin(SSD1306_SWITCHCAPVCC, i2c_address)) {
        OLED_LOG("Failed to initialize SSD1306 display");
        delete display;
        display = nullptr;
        return false;
    }
    
    // Configure display
    display->clearDisplay();
    display->setTextSize(1);
    display->setTextColor(SSD1306_WHITE);
    display->setCursor(0, 0);
    
    setBrightness(brightness);
    
    // Show boot screen
    setPage(PAGE_BOOT);
    update();
    
    OLED_LOG("OLED display initialized successfully");
    return true;
}

void OLEDDisplay::setBrightness(uint8_t new_brightness) {
    brightness = new_brightness;
    if (display) {
        display->ssd1306_command(SSD1306_SETCONTRAST);
        display->ssd1306_command(brightness);
    }
}

void OLEDDisplay::setAutoRotate(bool enable) {
    auto_rotate = enable;
    if (enable) {
        page_switch_time = millis();
    }
}

void OLEDDisplay::enable(bool enabled) {
    display_enabled = enabled;
    if (display) {
        if (enabled) {
            display->ssd1306_command(SSD1306_DISPLAYON);
        } else {
            display->ssd1306_command(SSD1306_DISPLAYOFF);
        }
    }
}

void OLEDDisplay::setPage(DisplayPage page) {
    if (page >= PAGE_COUNT) return;
    
    current_page = page;
    page_switch_time = millis();
    
    if (display) {
        display->clearDisplay();
    }
}

void OLEDDisplay::nextPage() {
    DisplayPage next = (DisplayPage)((current_page + 1) % PAGE_COUNT);
    setPage(next);
}

void OLEDDisplay::previousPage() {
    DisplayPage prev = (DisplayPage)((current_page - 1 + PAGE_COUNT) % PAGE_COUNT);
    setPage(prev);
}

void OLEDDisplay::updateSystemStatus(const SystemStatus& status) {
    system_status = status;
}

void OLEDDisplay::updateAttackInfo(const AttackInfo& info) {
    attack_info = info;
}

void OLEDDisplay::updateNetworkInfo(const NetworkInfo& info) {
    network_info = info;
}

void OLEDDisplay::setAttackStatus(AttackStatus status, const String& target) {
    attack_info.status = status;
    if (!target.isEmpty()) {
        attack_info.target_ssid = target;
    }
}

void OLEDDisplay::setCredentialCount(uint32_t count) {
    attack_info.credentials_captured = count;
}

void OLEDDisplay::setError(const String& error) {
    attack_info.last_error = error;
    setPage(PAGE_ERROR);
}

void OLEDDisplay::clearError() {
    attack_info.last_error = "";
    if (current_page == PAGE_ERROR) {
        setPage(PAGE_MAIN_STATUS);
    }
}

void OLEDDisplay::update() {
    if (!display || !display_enabled) return;
    
    uint32_t now = millis();
    
    // Update animations
    if (now - last_animation >= ANIMATION_UPDATE_INTERVAL) {
        updateAnimations();
        last_animation = now;
    }
    
    // Auto-rotate pages
    if (auto_rotate && now - page_switch_time >= PAGE_AUTO_SWITCH_INTERVAL) {
        nextPage();
    }
    
    // Update display content
    if (now - last_update >= DISPLAY_UPDATE_INTERVAL) {
        display->clearDisplay();
        
        switch (current_page) {
            case PAGE_BOOT:
                drawBootScreen();
                break;
            case PAGE_MAIN_STATUS:
                drawMainStatus();
                break;
            case PAGE_ATTACK_STATUS:
                drawAttackStatus();
                break;
            case PAGE_CREDENTIAL_COUNT:
                drawCredentialCount();
                break;
            case PAGE_NETWORK_INFO:
                drawNetworkInfo();
                break;
            case PAGE_SYSTEM_INFO:
                drawSystemInfo();
                break;
            case PAGE_ERROR:
                drawErrorScreen();
                break;
        }
        
        display->display();
        last_update = now;
    }
}

void OLEDDisplay::clear() {
    if (display) {
        display->clearDisplay();
        display->display();
    }
}

void OLEDDisplay::showMessage(const String& message, uint32_t duration_ms) {
    if (!display) return;
    
    display->clearDisplay();
    drawHeader("MESSAGE");
    
    display->setTextSize(1);
    drawWrappedText(message, 0, 16, SCREEN_WIDTH);
    
    display->display();
    delay(duration_ms);
}

void OLEDDisplay::showProgress(const String& message, float progress) {
    if (!display) return;
    
    display->clearDisplay();
    drawHeader("PROGRESS");
    
    display->setTextSize(1);
    display->setCursor(0, 16);
    display->println(message);
    
    drawProgressBar(0, 32, SCREEN_WIDTH, 8, progress);
    
    display->setCursor(0, 48);
    display->printf("%.1f%%", progress * 100.0f);
    
    display->display();
}

void OLEDDisplay::drawBootScreen() {
    drawHeader("WiFiX Enhanced");
    
    display->setTextSize(1);
    display->setCursor(0, 16);
    display->println("Initializing...");
    
    // Draw loading animation
    drawLoadingAnimation(SCREEN_WIDTH - 16, 16);
    
    display->setCursor(0, 32);
    display->println("ESP32: " + String(system_status.esp32_online ? "OK" : "..."));
    display->println("BW16:  " + String(system_status.bw16_online ? "OK" : "..."));
    display->println("AI:    " + String(system_status.ai_online ? "OK" : "..."));
    
    drawFooter();
}

void OLEDDisplay::drawMainStatus() {
    drawHeader("MAIN STATUS");
    
    // System status icons
    int icon_y = 16;
    drawStatusIcon(0, icon_y, system_status.esp32_online, "E");
    drawStatusIcon(16, icon_y, system_status.bw16_online, "B");
    drawStatusIcon(32, icon_y, system_status.ai_online, "A");
    drawWiFiIcon(48, icon_y, system_status.wifi_connected);
    
    if (system_status.wifi_connected) {
        drawSignalBars(64, icon_y, system_status.wifi_signal);
    }
    
    // Attack status
    display->setTextSize(1);
    display->setCursor(0, 32);
    
    switch (attack_info.status) {
        case ATTACK_IDLE:
            display->println("Status: IDLE");
            break;
        case ATTACK_SCANNING:
            display->println("Status: SCANNING");
            drawLoadingAnimation(SCREEN_WIDTH - 16, 32);
            break;
        case ATTACK_DEAUTH_ACTIVE:
            display->println("Status: DEAUTH");
            drawPulseAnimation(SCREEN_WIDTH - 16, 32, 6);
            break;
        case ATTACK_PORTAL_ACTIVE:
            display->println("Status: PORTAL");
            break;
        case ATTACK_SUCCESS:
            display->println("Status: SUCCESS");
            break;
        case ATTACK_FAILED:
            display->println("Status: FAILED");
            break;
    }
    
    display->printf("Creds: %d", attack_info.credentials_captured);
    
    drawFooter();
}

void OLEDDisplay::drawAttackStatus() {
    drawHeader("ATTACK STATUS");
    
    display->setTextSize(1);
    display->setCursor(0, 16);
    
    if (!attack_info.target_ssid.isEmpty()) {
        display->println("Target:");
        display->println(truncateString(attack_info.target_ssid, 16));
        display->println();
    }
    
    display->printf("Deauth: %d\n", attack_info.deauth_packets_sent);
    display->printf("Connects: %d\n", attack_info.portal_connections);
    display->printf("Duration: %s\n", formatDuration(attack_info.attack_duration).c_str());
    
    drawFooter();
}

void OLEDDisplay::drawCredentialCount() {
    drawHeader("CREDENTIALS");
    
    // Large credential count
    display->setTextSize(3);
    String count_str = String(attack_info.credentials_captured);
    int text_width = count_str.length() * 18; // Approximate width
    int x = (SCREEN_WIDTH - text_width) / 2;
    display->setCursor(x, 20);
    display->println(count_str);
    
    // Additional info
    display->setTextSize(1);
    display->setCursor(0, 48);
    display->printf("Success: %.1f%%", attack_info.success_rate);
    
    drawFooter();
}

void OLEDDisplay::drawNetworkInfo() {
    drawHeader("NETWORK INFO");
    
    display->setTextSize(1);
    display->setCursor(0, 16);
    
    if (!network_info.ap_ssid.isEmpty()) {
        display->println("AP: " + truncateString(network_info.ap_ssid, 12));
        display->println("IP: " + network_info.ap_ip.toString());
        display->printf("Clients: %d\n", network_info.connected_clients);
    }
    
    if (!network_info.sta_ssid.isEmpty()) {
        display->println("STA: " + truncateString(network_info.sta_ssid, 12));
    }
    
    display->printf("Nearby: %d", network_info.nearby_networks);
    
    drawFooter();
}

void OLEDDisplay::drawSystemInfo() {
    drawHeader("SYSTEM INFO");
    
    display->setTextSize(1);
    display->setCursor(0, 16);
    
    display->println("Uptime: " + formatUptime(system_status.uptime));
    display->println("Memory: " + formatMemory(system_status.free_memory));
    display->printf("CPU: %.1f%%\n", system_status.cpu_usage);
    display->printf("Temp: %dC", system_status.temperature);
    
    drawFooter();
}

void OLEDDisplay::drawErrorScreen() {
    drawHeader("ERROR");
    
    // Error icon
    display->drawBitmap(SCREEN_WIDTH/2 - 4, 16, icon_error, 8, 8, SSD1306_WHITE);
    
    display->setTextSize(1);
    drawWrappedText(attack_info.last_error, 0, 32, SCREEN_WIDTH);
    
    drawFooter();
}

void OLEDDisplay::drawHeader(const char* title) {
    display->setTextSize(1);
    display->setCursor(0, 0);
    display->println(title);
    display->drawLine(0, 8, SCREEN_WIDTH, 8, SSD1306_WHITE);
}

void OLEDDisplay::drawFooter() {
    display->drawLine(0, SCREEN_HEIGHT - 10, SCREEN_WIDTH, SCREEN_HEIGHT - 10, SSD1306_WHITE);
    
    display->setTextSize(1);
    display->setCursor(0, SCREEN_HEIGHT - 8);
    
    // Page indicator
    display->printf("%d/%d", current_page + 1, PAGE_COUNT);
    
    // Time or other info
    String uptime = formatUptime(system_status.uptime);
    int text_width = uptime.length() * 6;
    display->setCursor(SCREEN_WIDTH - text_width, SCREEN_HEIGHT - 8);
    display->print(uptime);
}

void OLEDDisplay::drawProgressBar(int x, int y, int width, int height, float progress) {
    // Border
    display->drawRect(x, y, width, height, SSD1306_WHITE);
    
    // Fill
    int fill_width = (int)(progress * (width - 2));
    if (fill_width > 0) {
        display->fillRect(x + 1, y + 1, fill_width, height - 2, SSD1306_WHITE);
    }
}

void OLEDDisplay::drawSignalBars(int x, int y, uint8_t signal_strength) {
    int bars = map(signal_strength, 0, 100, 0, 4);
    
    for (int i = 0; i < 4; i++) {
        int bar_height = (i + 1) * 2;
        if (i < bars) {
            display->fillRect(x + i * 3, y + 8 - bar_height, 2, bar_height, SSD1306_WHITE);
        } else {
            display->drawRect(x + i * 3, y + 8 - bar_height, 2, bar_height, SSD1306_WHITE);
        }
    }
}

void OLEDDisplay::drawBattery(int x, int y, uint8_t level) {
    // Battery outline
    display->drawRect(x, y, 12, 6, SSD1306_WHITE);
    display->drawRect(x + 12, y + 1, 2, 4, SSD1306_WHITE);
    
    // Battery fill
    int fill_width = map(level, 0, 100, 0, 10);
    if (fill_width > 0) {
        display->fillRect(x + 1, y + 1, fill_width, 4, SSD1306_WHITE);
    }
}

void OLEDDisplay::drawWiFiIcon(int x, int y, bool connected) {
    if (connected) {
        display->drawBitmap(x, y, icon_wifi_connected, 8, 8, SSD1306_WHITE);
    } else {
        display->drawBitmap(x, y, icon_wifi_disconnected, 8, 8, SSD1306_WHITE);
    }
}

void OLEDDisplay::drawStatusIcon(int x, int y, bool status, const char* icon) {
    display->setTextSize(1);
    display->setCursor(x, y);
    
    if (status) {
        display->setTextColor(SSD1306_WHITE);
    } else {
        // Draw inverted for offline status
        display->fillRect(x, y, 8, 8, SSD1306_WHITE);
        display->setTextColor(SSD1306_BLACK);
    }
    
    display->print(icon);
    display->setTextColor(SSD1306_WHITE); // Reset color
}

void OLEDDisplay::drawScrollingText(int x, int y, const String& text, int max_width) {
    int text_width = text.length() * 6;
    
    if (text_width <= max_width) {
        display->setCursor(x, y);
        display->print(text);
    } else {
        // Implement scrolling logic
        static int scroll_offset = 0;
        static uint32_t last_scroll = 0;
        
        if (millis() - last_scroll > 200) {
            scroll_offset = (scroll_offset + 1) % (text_width + max_width);
            last_scroll = millis();
        }
        
        display->setCursor(x - scroll_offset, y);
        display->print(text);
    }
}

void OLEDDisplay::updateAnimations() {
    animation_frame = (animation_frame + 1) % 8;
    blink_state = !blink_state;
}

void OLEDDisplay::drawLoadingAnimation(int x, int y) {
    const char* frames[] = {"|", "/", "-", "\\", "|", "/", "-", "\\"};
    display->setCursor(x, y);
    display->print(frames[animation_frame]);
}

void OLEDDisplay::drawPulseAnimation(int x, int y, int radius) {
    int pulse_radius = radius + (animation_frame % 4);
    display->drawCircle(x, y, pulse_radius, SSD1306_WHITE);
}

String OLEDDisplay::formatUptime(uint32_t seconds) {
    uint32_t days = seconds / 86400;
    uint32_t hours = (seconds % 86400) / 3600;
    uint32_t minutes = (seconds % 3600) / 60;
    
    if (days > 0) {
        return String(days) + "d " + String(hours) + "h";
    } else if (hours > 0) {
        return String(hours) + "h " + String(minutes) + "m";
    } else {
        return String(minutes) + "m " + String(seconds % 60) + "s";
    }
}

String OLEDDisplay::formatMemory(uint32_t bytes) {
    if (bytes >= 1024 * 1024) {
        return String(bytes / (1024 * 1024)) + "MB";
    } else if (bytes >= 1024) {
        return String(bytes / 1024) + "KB";
    } else {
        return String(bytes) + "B";
    }
}

String OLEDDisplay::formatDuration(uint32_t seconds) {
    if (seconds >= 3600) {
        return String(seconds / 3600) + "h " + String((seconds % 3600) / 60) + "m";
    } else if (seconds >= 60) {
        return String(seconds / 60) + "m " + String(seconds % 60) + "s";
    } else {
        return String(seconds) + "s";
    }
}

String OLEDDisplay::truncateString(const String& str, int max_chars) {
    if (str.length() <= max_chars) {
        return str;
    } else {
        return str.substring(0, max_chars - 3) + "...";
    }
}

void OLEDDisplay::drawCenteredText(const String& text, int y) {
    int text_width = text.length() * 6; // Approximate width for size 1
    int x = (SCREEN_WIDTH - text_width) / 2;
    display->setCursor(x, y);
    display->print(text);
}

void OLEDDisplay::drawRightAlignedText(const String& text, int x, int y) {
    int text_width = text.length() * 6;
    display->setCursor(x - text_width, y);
    display->print(text);
}

void OLEDDisplay::drawWrappedText(const String& text, int x, int y, int max_width) {
    int chars_per_line = max_width / 6; // Approximate characters per line
    int current_y = y;
    
    for (int i = 0; i < text.length(); i += chars_per_line) {
        String line = text.substring(i, min((int)text.length(), i + chars_per_line));
        display->setCursor(x, current_y);
        display->println(line);
        current_y += 8;
        
        if (current_y >= SCREEN_HEIGHT - 8) break; // Don't overflow screen
    }
}

void OLEDDisplay::printStatus() {
    OLED_LOG_F("Current page: %d", current_page);
    OLED_LOG_F("Auto rotate: %s", auto_rotate ? "enabled" : "disabled");
    OLED_LOG_F("Display enabled: %s", display_enabled ? "yes" : "no");
    OLED_LOG_F("Brightness: %d", brightness);
    OLED_LOG_F("Attack status: %d", attack_info.status);
    OLED_LOG_F("Credentials captured: %d", attack_info.credentials_captured);
}

void OLEDDisplay::testDisplay() {
    OLED_LOG("Testing display...");
    
    for (int page = 0; page < PAGE_COUNT; page++) {
        setPage((DisplayPage)page);
        update();
        delay(2000);
    }
    
    OLED_LOG("Display test completed");
}