#include <stdio.h>
#include "pico/stdlib.h"
#include "hardware/i2c.h"
#include "ssd1306.h"

// I2C Configuration
#define I2C_PORT i2c0
#define I2C_SDA_PIN 4
#define I2C_SCL_PIN 5

// Rotary Encoder Configuration (KY-040)
#define ENCODER_CLK 14   // Clock pin
#define ENCODER_DT 15    // Data pin
#define ENCODER_SW 16    // Switch/Button pin

// DS1307 Configuration
#define DS1307_ADDR 0x68
#define DS1307_REG_SECONDS 0x00
#define DS1307_REG_MINUTES 0x01
#define DS1307_REG_HOURS 0x02

// Page/Screen enumeration
typedef enum {
    SCREEN_MENU = 0,      // Main menu screen
    SCREEN_TIME,          // Time display
    SCREEN_ALARM,         // Alarm settings
    SCREEN_TIMER          // Countdown timer
} screen_id_t;

// Menu items
typedef enum {
    MENU_TIME = 0,
    MENU_ALARM,
    MENU_TIMER,
    MENU_COUNT
} menu_item_t;

// Timer state
typedef enum {
    TIMER_STOPPED = 0,
    TIMER_RUNNING,
    TIMER_PAUSED,
    TIMER_FINISHED
} timer_state_t;

// Global variables
static ssd1306_t display;
static screen_id_t current_screen = SCREEN_MENU;
static menu_item_t selected_menu_item = MENU_TIME;
static bool encoder_button_pressed = false;
static int8_t encoder_direction = 0;  // -1 = CCW, 0 = none, 1 = CW

// Encoder state tracking
static uint8_t last_clk_state = 1;

// Time structures
typedef struct {
    uint8_t seconds;
    uint8_t minutes;
    uint8_t hours;
} time_24h_t;

typedef struct {
    uint8_t hours;
    uint8_t minutes;
    uint8_t seconds;
    bool is_pm;
} time_12h_t;

// Alarm structure
typedef struct {
    uint8_t hours;
    uint8_t minutes;
    bool enabled;
    bool triggered;
} alarm_t;

static alarm_t alarm_data = {
    .hours = 7,
    .minutes = 30,
    .enabled = false,
    .triggered = false
};

// Timer structure
typedef struct {
    uint32_t duration_seconds;
    uint32_t remaining_seconds;
    uint32_t last_update_time;
    timer_state_t state;
} timer_t;

static timer_t countdown_timer = {
    .duration_seconds = 60,
    .remaining_seconds = 60,
    .last_update_time = 0,
    .state = TIMER_STOPPED
};

// ------------ Helper Functions --------------

uint8_t bcd_to_dec(uint8_t val) {
    return ((val >> 4) * 10) + (val & 0x0F);
}

void convert_to_12h(const time_24h_t *time_24h, time_12h_t *time_12h) {
    time_12h->minutes = time_24h->minutes;
    time_12h->seconds = time_24h->seconds;
    
    if (time_24h->hours == 0) {
        time_12h->hours = 12;
        time_12h->is_pm = false;
    } else if (time_24h->hours < 12) {
        time_12h->hours = time_24h->hours;
        time_12h->is_pm = false;
    } else if (time_24h->hours == 12) {
        time_12h->hours = 12;
        time_12h->is_pm = true;
    } else {
        time_12h->hours = time_24h->hours - 12;
        time_12h->is_pm = true;
    }
}

// ------------ DS1307 Functions --------------

void ds1307_read_register(uint8_t reg, uint8_t *buf, uint8_t len) {
    i2c_write_blocking(I2C_PORT, DS1307_ADDR, &reg, 1, true);
    i2c_read_blocking(I2C_PORT, DS1307_ADDR, buf, len, false);
}

void ds1307_get_time(time_24h_t *time) {
    uint8_t data[3];
    ds1307_read_register(DS1307_REG_SECONDS, data, 3);
    
    time->seconds = bcd_to_dec(data[0] & 0x7F);
    time->minutes = bcd_to_dec(data[1] & 0x7F);
    time->hours = bcd_to_dec(data[2] & 0x3F);
}

// ------------ Rotary Encoder Functions --------------

void encoder_read() {
    uint8_t clk_state = gpio_get(ENCODER_CLK);
    uint8_t dt_state = gpio_get(ENCODER_DT);
    
    // Detect rotation on CLK falling edge
    if (clk_state != last_clk_state && clk_state == 0) {
        if (dt_state != clk_state) {
            encoder_direction = 1;   // Clockwise
        } else {
            encoder_direction = -1;  // Counter-clockwise
        }
    }
    
    last_clk_state = clk_state;
}

void gpio_callback(uint gpio, uint32_t events) {
    if (gpio == ENCODER_SW && (events & GPIO_IRQ_EDGE_FALL)) {
        encoder_button_pressed = true;
    }
    
    if ((gpio == ENCODER_CLK || gpio == ENCODER_DT) && (events & GPIO_IRQ_EDGE_FALL)) {
        encoder_read();
    }
}

void encoder_init() {
    // Initialize CLK pin
    gpio_init(ENCODER_CLK);
    gpio_set_dir(ENCODER_CLK, GPIO_IN);
    gpio_pull_up(ENCODER_CLK);
    
    // Initialize DT pin
    gpio_init(ENCODER_DT);
    gpio_set_dir(ENCODER_DT, GPIO_IN);
    gpio_pull_up(ENCODER_DT);
    
    // Initialize SW (button) pin
    gpio_init(ENCODER_SW);
    gpio_set_dir(ENCODER_SW, GPIO_IN);
    gpio_pull_up(ENCODER_SW);
    
    // Read initial state
    last_clk_state = gpio_get(ENCODER_CLK);
    
    // Enable interrupts
    gpio_set_irq_enabled_with_callback(ENCODER_CLK, GPIO_IRQ_EDGE_FALL | GPIO_IRQ_EDGE_RISE, true, &gpio_callback);
    gpio_set_irq_enabled(ENCODER_DT, GPIO_IRQ_EDGE_FALL | GPIO_IRQ_EDGE_RISE, true);
    gpio_set_irq_enabled(ENCODER_SW, GPIO_IRQ_EDGE_FALL, true);
}

// ------------ Alarm Functions --------------

void check_alarm(const time_24h_t *current_time) {
    if (!alarm_data.enabled) return;
    
    if (current_time->hours == alarm_data.hours && 
        current_time->minutes == alarm_data.minutes &&
        current_time->seconds == 0 &&
        !alarm_data.triggered) {
        
        alarm_data.triggered = true;
        printf("*** ALARM TRIGGERED! ***\n");
        
        for (int i = 0; i < 3; i++) {
            ssd1306_invert(&display, 1);
            ssd1306_show(&display);
            sleep_ms(200);
            ssd1306_invert(&display, 0);
            ssd1306_show(&display);
            sleep_ms(200);
        }
    }
    
    if (current_time->minutes != alarm_data.minutes) {
        alarm_data.triggered = false;
    }
}

void alarm_toggle() {
    alarm_data.enabled = !alarm_data.enabled;
    printf("Alarm %s\n", alarm_data.enabled ? "ENABLED" : "DISABLED");
}

// ------------ Timer Functions --------------

void timer_start() {
    if (countdown_timer.state == TIMER_STOPPED) {
        countdown_timer.remaining_seconds = countdown_timer.duration_seconds;
        countdown_timer.last_update_time = to_ms_since_boot(get_absolute_time());
        countdown_timer.state = TIMER_RUNNING;
    } else if (countdown_timer.state == TIMER_PAUSED) {
        countdown_timer.last_update_time = to_ms_since_boot(get_absolute_time());
        countdown_timer.state = TIMER_RUNNING;
    }
}

void timer_stop() {
    if (countdown_timer.state == TIMER_RUNNING) {
        countdown_timer.state = TIMER_PAUSED;
    }
}

void timer_reset() {
    countdown_timer.remaining_seconds = countdown_timer.duration_seconds;
    countdown_timer.state = TIMER_STOPPED;
}

void timer_update() {
    if (countdown_timer.state == TIMER_RUNNING) {
        uint32_t now = to_ms_since_boot(get_absolute_time());
        uint32_t elapsed_ms = now - countdown_timer.last_update_time;
        
        if (elapsed_ms >= 1000) {
            uint32_t elapsed_seconds = elapsed_ms / 1000;
            countdown_timer.last_update_time = now;
            
            if (countdown_timer.remaining_seconds > elapsed_seconds) {
                countdown_timer.remaining_seconds -= elapsed_seconds;
            } else {
                countdown_timer.remaining_seconds = 0;
                countdown_timer.state = TIMER_FINISHED;
                printf("*** TIMER FINISHED! ***\n");
                
                for (int i = 0; i < 5; i++) {
                    ssd1306_invert(&display, 1);
                    ssd1306_show(&display);
                    sleep_ms(200);
                    ssd1306_invert(&display, 0);
                    ssd1306_show(&display);
                    sleep_ms(200);
                }
            }
        }
    }
}

// ------------ Screen Drawing Functions --------------

// MAIN MENU SCREEN
void draw_screen_menu() {
    ssd1306_clear(&display);
    
    // Title
    ssd1306_draw_string(&display, 35, 2, 1, "MAIN MENU");
    ssd1306_draw_line(&display, 0, 12, 127, 12);
    
    // Menu items with selection indicator
    const char* menu_items[] = {"Time", "Alarm", "Timer"};
    
    for (int i = 0; i < MENU_COUNT; i++) {
        uint8_t y_pos = 20 + (i * 12);
        
        // Draw selection arrow
        if (i == selected_menu_item) {
            ssd1306_draw_string(&display, 10, y_pos, 1, ">");
        }
        
        // Draw menu item text
        ssd1306_draw_string(&display, 25, y_pos, 1, menu_items[i]);
    }
    
    // Instructions at bottom
    ssd1306_draw_line(&display, 0, 52, 127, 52);
    ssd1306_draw_string(&display, 8, 54, 1, "Rotate:Select");
    ssd1306_draw_string(&display, 8, 62, 1, "Click:Enter");  // Won't show on 64px, but added
    
    ssd1306_show(&display);
}

// TIME SCREEN
void draw_screen_time() {
    time_24h_t time_24h;
    ds1307_get_time(&time_24h);
    
    time_12h_t time_12h;
    convert_to_12h(&time_24h, &time_12h);
    
    ssd1306_clear(&display);
    
    // Title
    ssd1306_draw_string(&display, 25, 2, 1, "12-HOUR TIME");
    ssd1306_draw_line(&display, 0, 12, 127, 12);
    
    // Large time display
    char time_str[16];
    snprintf(time_str, sizeof(time_str), "%02d:%02d:%02d", 
             time_12h.hours, time_12h.minutes, time_12h.seconds);
    ssd1306_draw_string(&display, 10, 24, 2, time_str);
    
    // AM/PM indicator
    ssd1306_draw_string(&display, 50, 44, 2, time_12h.is_pm ? "PM" : "AM");
    
    // Back instruction
    ssd1306_draw_string(&display, 20, 54, 1, "Click: BACK");
    
    ssd1306_show(&display);
    
    check_alarm(&time_24h);
}

// ALARM SCREEN
void draw_screen_alarm() {
    ssd1306_clear(&display);
    
    // Title
    ssd1306_draw_string(&display, 40, 2, 1, "ALARM");
    ssd1306_draw_line(&display, 0, 12, 127, 12);
    
    // Alarm time in 12h format
    time_24h_t alarm_time_24h = {
        .hours = alarm_data.hours,
        .minutes = alarm_data.minutes,
        .seconds = 0
    };
    time_12h_t alarm_time_12h;
    convert_to_12h(&alarm_time_24h, &alarm_time_12h);
    
    char alarm_str[16];
    snprintf(alarm_str, sizeof(alarm_str), "%02d:%02d %s", 
             alarm_time_12h.hours, alarm_time_12h.minutes,
             alarm_time_12h.is_pm ? "PM" : "AM");
    ssd1306_draw_string(&display, 20, 22, 2, alarm_str);
    
    // Status with visual indicator
    ssd1306_draw_string(&display, 25, 40, 1, "Status:");
    if (alarm_data.enabled) {
        ssd1306_draw_square(&display, 70, 38, 10, 10);  // Filled
        ssd1306_draw_string(&display, 85, 40, 1, "ON");
    } else {
        ssd1306_draw_empty_square(&display, 70, 38, 10, 10);  // Empty
        ssd1306_draw_string(&display, 85, 40, 1, "OFF");
    }
    
    // Instructions
    ssd1306_draw_string(&display, 8, 54, 1, "Rotate:Toggle");
    ssd1306_draw_string(&display, 80, 54, 1, "Clk:BACK");
    
    ssd1306_show(&display);
}

// TIMER SCREEN
void draw_screen_timer() {
    timer_update();
    
    ssd1306_clear(&display);
    
    // Title
    ssd1306_draw_string(&display, 35, 2, 1, "TIMER");
    ssd1306_draw_line(&display, 0, 12, 127, 12);
    
    // Timer display
    uint32_t minutes = countdown_timer.remaining_seconds / 60;
    uint32_t seconds = countdown_timer.remaining_seconds % 60;
    
    char timer_str[16];
    snprintf(timer_str, sizeof(timer_str), "%02lu:%02lu", minutes, seconds);
    ssd1306_draw_string(&display, 25, 22, 3, timer_str);
    
    // State indicator
    const char* state_str = "";
    switch (countdown_timer.state) {
        case TIMER_STOPPED:
            state_str = "READY";
            break;
        case TIMER_RUNNING:
            state_str = "RUNNING";
            // Progress bar
            uint32_t progress = (countdown_timer.remaining_seconds * 100) / countdown_timer.duration_seconds;
            uint8_t bar_width = (progress * 118) / 100;
            ssd1306_draw_empty_square(&display, 5, 46, 118, 6);
            if (bar_width > 0) {
                ssd1306_draw_square(&display, 6, 47, bar_width, 4);
            }
            break;
        case TIMER_PAUSED:
            state_str = "PAUSED";
            break;
        case TIMER_FINISHED:
            state_str = "DONE!";
            if ((to_ms_since_boot(get_absolute_time()) / 500) % 2) {
                ssd1306_invert(&display, 1);
            }
            break;
    }
    ssd1306_draw_string(&display, 40, 40, 1, state_str);
    
    // Instructions
    if (countdown_timer.state == TIMER_STOPPED || countdown_timer.state == TIMER_PAUSED) {
        ssd1306_draw_string(&display, 15, 54, 1, "Rot:Start");
    } else if (countdown_timer.state == TIMER_RUNNING) {
        ssd1306_draw_string(&display, 15, 54, 1, "Rot:Pause");
    } else if (countdown_timer.state == TIMER_FINISHED) {
        ssd1306_draw_string(&display, 15, 54, 1, "Rot:Reset");
    }
    ssd1306_draw_string(&display, 80, 54, 1, "Clk:BACK");
    
    ssd1306_show(&display);
}

// ------------ Screen Router --------------

void draw_current_screen() {
    switch (current_screen) {
        case SCREEN_MENU:
            draw_screen_menu();
            break;
        case SCREEN_TIME:
            draw_screen_time();
            break;
        case SCREEN_ALARM:
            draw_screen_alarm();
            break;
        case SCREEN_TIMER:
            draw_screen_timer();
            break;
    }
}

// ------------ Input Handlers --------------

void handle_menu_encoder_rotation(int8_t direction) {
    if (direction > 0) {
        // Clockwise - move down
        selected_menu_item = (selected_menu_item + 1) % MENU_COUNT;
    } else if (direction < 0) {
        // Counter-clockwise - move up
        if (selected_menu_item == 0) {
            selected_menu_item = MENU_COUNT - 1;
        } else {
            selected_menu_item--;
        }
    }
    printf("Menu selection: %d\n", selected_menu_item);
}

void handle_menu_button_press() {
    // Enter selected menu item
    switch (selected_menu_item) {
        case MENU_TIME:
            current_screen = SCREEN_TIME;
            printf("Entered TIME screen\n");
            break;
        case MENU_ALARM:
            current_screen = SCREEN_ALARM;
            printf("Entered ALARM screen\n");
            break;
        case MENU_TIMER:
            current_screen = SCREEN_TIMER;
            printf("Entered TIMER screen\n");
            break;
    }
}

void handle_alarm_encoder_rotation(int8_t direction) {
    if (direction != 0) {
        alarm_toggle();
    }
}

void handle_timer_encoder_rotation(int8_t direction) {
    if (direction != 0) {
        if (countdown_timer.state == TIMER_STOPPED || countdown_timer.state == TIMER_PAUSED) {
            timer_start();
        } else if (countdown_timer.state == TIMER_RUNNING) {
            timer_stop();
        } else if (countdown_timer.state == TIMER_FINISHED) {
            timer_reset();
        }
    }
}

void handle_back_button() {
    current_screen = SCREEN_MENU;
    printf("Returned to MENU\n");
}

// ------------ Main --------------

int main() {
    stdio_init_all();
    
    // Initialize I2C
    i2c_init(I2C_PORT, 400000);
    gpio_set_function(I2C_SDA_PIN, GPIO_FUNC_I2C);
    gpio_set_function(I2C_SCL_PIN, GPIO_FUNC_I2C);
    gpio_pull_up(I2C_SDA_PIN);
    gpio_pull_up(I2C_SCL_PIN);
    
    sleep_ms(250);
    
    // Initialize OLED
    display.external_vcc = false;
    ssd1306_init(&display, 128, 64, 0x3C, I2C_PORT);
    
    // Initialize rotary encoder
    encoder_init();
    
    printf("\n=== Menu-Based OLED System ===\n");
    printf("Rotary Encoder Navigation:\n");
    printf("  GPIO %d: CLK\n", ENCODER_CLK);
    printf("  GPIO %d: DT\n", ENCODER_DT);
    printf("  GPIO %d: SW (Button)\n\n", ENCODER_SW);
    
    printf("Controls:\n");
    printf("  Rotate: Navigate/Control\n");
    printf("  Click: Select/Back\n\n");
    
    // Display menu
    draw_current_screen();
    
    uint32_t last_update = 0;
    
    while (1) {
        // Handle encoder rotation
        if (encoder_direction != 0) {
            int8_t direction = encoder_direction;
            encoder_direction = 0;  // Reset
            
            switch (current_screen) {
                case SCREEN_MENU:
                    handle_menu_encoder_rotation(direction);
                    break;
                case SCREEN_ALARM:
                    handle_alarm_encoder_rotation(direction);
                    break;
                case SCREEN_TIMER:
                    handle_timer_encoder_rotation(direction);
                    break;
                default:
                    break;
            }
            
            draw_current_screen();
            sleep_ms(50);  // Debounce
        }
        
        // Handle encoder button press
        if (encoder_button_pressed) {
            encoder_button_pressed = false;
            
            if (current_screen == SCREEN_MENU) {
                handle_menu_button_press();
            } else {
                handle_back_button();
            }
            
            draw_current_screen();
            sleep_ms(200);  // Debounce
        }
        
        // Auto-refresh for time and running timer
        uint32_t now = to_ms_since_boot(get_absolute_time());
        uint32_t update_interval = 1000;
        
        if (countdown_timer.state == TIMER_RUNNING && current_screen == SCREEN_TIMER) {
            update_interval = 500;
        }
        
        if (now - last_update >= update_interval) {
            last_update = now;
            if (current_screen == SCREEN_TIME || 
                (current_screen == SCREEN_TIMER && countdown_timer.state == TIMER_RUNNING)) {
                draw_current_screen();
            }
        }
        
        sleep_ms(10);
    }
    
    return 0;
}