#ifndef OLED_DISPLAY_H
#define OLED_DISPLAY_H

#include <Arduino.h>
#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>
#include <WiFi.h>

// Display configuration
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
#define OLED_RESET -1
#define SCREEN_ADDRESS 0x3C

// Display pages
enum DisplayPage {
    PAGE_BOOT = 0,
    PAGE_MAIN_STATUS,
    PAGE_ATTACK_STATUS,
    PAGE_CREDENTIAL_COUNT,
    PAGE_NETWORK_INFO,
    PAGE_SYSTEM_INFO,
    PAGE_ERROR,
    PAGE_COUNT
};

// Attack status types
enum AttackStatus {
    ATTACK_IDLE = 0,
    ATTACK_SCANNING,
    ATTACK_DEAUTH_ACTIVE,
    ATTACK_PORTAL_ACTIVE,
    ATTACK_SUCCESS,
    ATTACK_FAILED
};

// System status
struct SystemStatus {
    bool esp32_online;
    bool bw16_online;
    bool ai_online;
    bool wifi_connected;
    uint8_t wifi_signal;
    uint32_t uptime;
    uint32_t free_memory;
    float cpu_usage;
    uint8_t temperature;
};

// Attack information
struct AttackInfo {
    AttackStatus status;
    String target_ssid;
    String target_bssid;
    uint8_t target_channel;
    uint32_t deauth_packets_sent;
    uint32_t portal_connections;
    uint32_t credentials_captured;
    uint32_t attack_duration;
    float success_rate;
    String last_error;
};

// Network information
struct NetworkInfo {
    String ap_ssid;
    String ap_password;
    IPAddress ap_ip;
    uint8_t connected_clients;
    String sta_ssid;
    IPAddress sta_ip;
    uint8_t nearby_networks;
};

class OLEDDisplay {
private:
    Adafruit_SSD1306* display;
    DisplayPage current_page;
    uint32_t last_update;
    uint32_t page_switch_time;
    bool auto_rotate;
    uint8_t brightness;
    bool display_enabled;
    
    // Status data
    SystemStatus system_status;
    AttackInfo attack_info;
    NetworkInfo network_info;
    
    // Animation variables
    uint8_t animation_frame;
    uint32_t last_animation;
    bool blink_state;
    
    // Display methods
    void drawBootScreen();
    void drawMainStatus();
    void drawAttackStatus();
    void drawCredentialCount();
    void drawNetworkInfo();
    void drawSystemInfo();
    void drawErrorScreen();
    
    // Helper methods
    void drawHeader(const char* title);
    void drawFooter();
    void drawProgressBar(int x, int y, int width, int height, float progress);
    void drawSignalBars(int x, int y, uint8_t signal_strength);
    void drawBattery(int x, int y, uint8_t level);
    void drawWiFiIcon(int x, int y, bool connected);
    void drawStatusIcon(int x, int y, bool status, const char* icon);
    void drawScrollingText(int x, int y, const String& text, int max_width);
    
    // Animation methods
    void updateAnimations();
    void drawLoadingAnimation(int x, int y);
    void drawPulseAnimation(int x, int y, int radius);
    
    // Text formatting
    String formatUptime(uint32_t seconds);
    String formatMemory(uint32_t bytes);
    String formatDuration(uint32_t seconds);
    String truncateString(const String& str, int max_chars);

public:
    OLEDDisplay();
    ~OLEDDisplay();
    
    // Initialization
    bool begin(uint8_t i2c_address = SCREEN_ADDRESS);
    void setBrightness(uint8_t brightness);
    void setAutoRotate(bool enable);
    void enable(bool enabled);
    
    // Page management
    void setPage(DisplayPage page);
    void nextPage();
    void previousPage();
    DisplayPage getCurrentPage() const { return current_page; }
    
    // Status updates
    void updateSystemStatus(const SystemStatus& status);
    void updateAttackInfo(const AttackInfo& info);
    void updateNetworkInfo(const NetworkInfo& info);
    
    // Quick status updates
    void setAttackStatus(AttackStatus status, const String& target = "");
    void setCredentialCount(uint32_t count);
    void setError(const String& error);
    void clearError();
    
    // Display control
    void update();
    void clear();
    void showMessage(const String& message, uint32_t duration_ms = 2000);
    void showProgress(const String& message, float progress);
    
    // Utility methods
    void drawCenteredText(const String& text, int y);
    void drawRightAlignedText(const String& text, int x, int y);
    void drawWrappedText(const String& text, int x, int y, int max_width);
    
    // Debug methods
    void printStatus();
    void testDisplay();
};

// Global instance
extern OLEDDisplay oledDisplay;

// Convenience macros
#define OLED_LOG(msg) Serial.println("[OLED] " + String(msg))
#define OLED_LOG_F(fmt, ...) Serial.printf("[OLED] " fmt "\n", ##__VA_ARGS__)

// Display update intervals (ms)
#define DISPLAY_UPDATE_INTERVAL 100
#define PAGE_AUTO_SWITCH_INTERVAL 5000
#define ANIMATION_UPDATE_INTERVAL 200

// Icons (8x8 bitmaps)
extern const unsigned char PROGMEM icon_wifi_connected[];
extern const unsigned char PROGMEM icon_wifi_disconnected[];
extern const unsigned char PROGMEM icon_bluetooth[];
extern const unsigned char PROGMEM icon_lora[];
extern const unsigned char PROGMEM icon_zigbee[];
extern const unsigned char PROGMEM icon_attack[];
extern const unsigned char PROGMEM icon_shield[];
extern const unsigned char PROGMEM icon_warning[];
extern const unsigned char PROGMEM icon_success[];
extern const unsigned char PROGMEM icon_error[];

#endif // OLED_DISPLAY_H