#include <stdio.h>
#include <string.h>
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "freertos/semphr.h"
#include "esp_log.h"
#include "esp_check.h"

#include "driver/i2c_master.h"

#include "ssd1306.h"
#include "font8x8_basic.h"
#include "display_task.h" // Now includes game_framework.h indirectly

#define TAG "DISPLAY_TASK"

#define I2C_SDA_GPIO 19
#define I2C_SCL_GPIO 18
#define OLED_I2C_ADDRESS 0x3C
#define I2C_PORT I2C_NUM_1
#define I2C_SPEED_HZ 400000

// OLED_WIDTH, OLED_HEIGHT, SCORE_DISPLAY_HEIGHT etc. are now from game_framework.h

#define N_DISPLAY_LINES (OLED_HEIGHT / 8)
#define MAX_LINE_LENGTH (OLED_WIDTH / 8) + 1

SSD1306_t dev;
SemaphoreHandle_t xDisplayMutex = NULL;
static char display_lines[N_DISPLAY_LINES][MAX_LINE_LENGTH];
static bool dirty_lines[N_DISPLAY_LINES] = { true }; // Track which lines changed

// Define the display event queue (its definition)
QueueHandle_t display_event_queue = NULL;
// Store the global button input queue for use in game over screen
static QueueHandle_t s_button_input_queue = NULL;


typedef enum {
    DISPLAY_STATE_MENU, // State for displaying the game selection menu
    DISPLAY_STATE_GAME_RUNNING,
    DISPLAY_STATE_GAME_OVER,
    DISPLAY_STATE_OTHER_INFO
} display_state_t;

// Initialize to MENU state so it starts by displaying the menu
static display_state_t s_current_display_state = DISPLAY_STATE_MENU;
static int s_game_over_score = 0;
static int s_displayed_score = -1; // Tracks the score currently drawn on the display

// Flag to indicate a state change, triggering a full screen clear and redraw
static bool s_state_changed = true;

// Pointer to the currently active game instance
static const game_t *s_current_game = NULL;

static esp_err_t display_i2c_ng_init(void) {
    i2c_master_bus_config_t bus_config = {
        .clk_source = I2C_CLK_SRC_DEFAULT,
        .i2c_port = I2C_PORT,
        .scl_io_num = I2C_SCL_GPIO,
        .sda_io_num = I2C_SDA_GPIO,
        .flags = {
            .enable_internal_pullup = true,
        }
    };

    ESP_RETURN_ON_ERROR(i2c_new_master_bus(&bus_config, &dev._i2c_bus_handle), TAG, "Failed to create I2C bus");

    i2c_device_config_t dev_cfg = {
        .dev_addr_length = I2C_ADDR_BIT_LEN_7,
        .device_address = OLED_I2C_ADDRESS,
        .scl_speed_hz = I2C_SPEED_HZ,
    };

    ESP_RETURN_ON_ERROR(i2c_master_bus_add_device(dev._i2c_bus_handle, &dev_cfg, &dev._i2c_dev_handle), TAG, "Failed to add I2C OLED device");

    return ESP_OK;
}

static void display_manager_task(void *pvParameters) {
    ESP_LOGI(TAG, "Display Manager Task started.");

    if (display_i2c_ng_init() != ESP_OK) {
        ESP_LOGE(TAG, "I2C init failed, halting task.");
        vTaskDelete(NULL);
    }

    ssd1306_init(&dev, OLED_WIDTH, OLED_HEIGHT);
    ssd1306_clear_screen(&dev, false);
    ssd1306_show_buffer(&dev);

    // Static variables to track previous button states for edge detection in game over screen
    static bool left_button_pressed_prev_go = false;
    static bool right_button_pressed_prev_go = false;

    while (1) {
        if (xSemaphoreTake(xDisplayMutex, portMAX_DELAY) == pdTRUE) {
            // Check for state change and clear screen if necessary
            if (s_state_changed) {
                ssd1306_clear_screen(&dev, false); // Clear the entire display buffer
                // Mark all lines dirty to ensure they are redrawn with new content
                for (int i = 0; i < N_DISPLAY_LINES; i++) {
                    dirty_lines[i] = true;
                }
                s_state_changed = false; // Reset flag after clearing
            }

            switch (s_current_display_state) {
            case DISPLAY_STATE_MENU: {
                // In menu state, simply display the text set by display_update_text
                for (int i = 0; i < N_DISPLAY_LINES; i++) {
                    if (dirty_lines[i] && strlen(display_lines[i]) > 0) {
                        ssd1306_display_text(&dev, i, display_lines[i], strlen(display_lines[i]), false);
                        dirty_lines[i] = false;
                    }
                }
                break;
            }
            case DISPLAY_STATE_GAME_RUNNING: {
                if (s_current_game && s_current_game->is_active && s_current_game->is_active()) {
                    if (s_current_game->handle_input) {
                        s_current_game->handle_input();
                    }
                    bool game_over = false;
                    if (s_current_game->update_game_state) {
                        game_over = s_current_game->update_game_state();
                    }
                    if (s_current_game->draw_game) {
                        s_current_game->draw_game(&dev); // This now only draws game elements
                    }

                    // Update score display only if it has changed
                    int current_game_score = 0;
                    if (s_current_game && s_current_game->get_score) {
                        current_game_score = s_current_game->get_score();
                    }

                    if (current_game_score != s_displayed_score) {
                        char score_str[MAX_LINE_LENGTH];
                        snprintf(score_str, sizeof(score_str), "Score: %d", current_game_score);
                        // Clear the entire line 0 before drawing the new score to avoid artifacts
                        ssd1306_display_text(&dev, 0, "                ", strlen("                "), false); // Clear line 0
                        ssd1306_display_text(&dev, 0, score_str, strlen(score_str), false);
                        s_displayed_score = current_game_score;
                    }

                    if (game_over) {
                        s_game_over_score = current_game_score; // Use the final score
                        s_current_display_state = DISPLAY_STATE_GAME_OVER;
                        if (s_current_game && s_current_game->set_active) {
                            s_current_game->set_active(false);
                        }
                        s_state_changed = true; // Mark state changed for next iteration (game over screen)
                        s_displayed_score = -1; // Reset score display tracker
                        // Reset button states for game over screen input
                        left_button_pressed_prev_go = false;
                        right_button_pressed_prev_go = false;
                    }
                } else {
                    // This block handles cases where the game might become inactive externally
                    // or if is_active() becomes false for other reasons.
                    if (s_current_game && s_current_game->get_score) {
                        s_game_over_score = s_current_game->get_score(); // Get final score
                    } else {
                        s_game_over_score = 0; // Default if no game or get_score not implemented
                    }
                    s_current_display_state = DISPLAY_STATE_GAME_OVER;
                    s_state_changed = true; // Mark state changed for next iteration (game over screen)
                    s_displayed_score = -1; // Reset score display tracker
                    // Reset button states for game over screen input
                    left_button_pressed_prev_go = false;
                    right_button_pressed_prev_go = false;
                }
                break;
            }
            case DISPLAY_STATE_GAME_OVER: {
                char score_display[MAX_LINE_LENGTH];
                snprintf(score_display, sizeof(score_display), "Score: %d", s_game_over_score);
                const char *game_over_messages[] = {
                    "    GAME OVER!    ",
                    "                  ",
                    score_display,
                    "                  ",
                    "  Left: RESTART   ", // Instructions for restart
                    " Right: MENU      ", // Instructions for menu
                    NULL, NULL
                };

                for (int i = 0; i < N_DISPLAY_LINES; i++) {
                    if (game_over_messages[i] && strlen(game_over_messages[i]) > 0) { // Only draw if message exists
                        ssd1306_display_text(&dev, i, game_over_messages[i], strlen(game_over_messages[i]), false);
                        dirty_lines[i] = false; // Mark as drawn
                    } else if (dirty_lines[i]) { // Clear line if it was dirty but no message
                        ssd1306_display_text(&dev, i, "                ", strlen("                "), false);
                        dirty_lines[i] = false;
                    }
                }

                // Handle input on game over screen using the global button_input_queue
                int input_signal = 0;
                if (s_button_input_queue != NULL && xQueueReceive(s_button_input_queue, &input_signal, 0) == pdTRUE) {
                    
                    bool left_button_now_pressed = (input_signal == -1);
                    bool right_button_now_pressed = (input_signal == 1);

                    bool left_press_event = left_button_now_pressed && !left_button_pressed_prev_go;
                    bool right_press_event = right_button_now_pressed && !right_button_pressed_prev_go;

                    if (left_press_event) { // Left button pressed: Restart current game
                        ESP_LOGI(TAG, "Restarting current game...");
                        s_current_display_state = DISPLAY_STATE_GAME_RUNNING;
                        if (s_current_game && s_current_game->init_game) {
                            s_current_game->init_game();
                        }
                        if (s_current_game && s_current_game->set_active) {
                            s_current_game->set_active(true);
                        }
                        s_state_changed = true; // Mark state changed for next iteration (game running)
                        s_displayed_score = -1;
                    } else if (right_press_event) { // Right button pressed: Return to main menu
                        ESP_LOGI(TAG, "Returning to main menu...");
                        s_current_display_state = DISPLAY_STATE_MENU; // Transition to menu state
                        if (s_current_game && s_current_game->set_active) {
                            s_current_game->set_active(false); // Deactivate current game
                        }
                        s_state_changed = true; // Mark state changed for next iteration (menu)
                        // Signal main.c to re-run the menu logic
                        display_event_t event = DISPLAY_EVENT_RETURN_TO_MENU;
                        if (display_event_queue != NULL) {
                            xQueueSend(display_event_queue, &event, 0);
                        } else {
                            ESP_LOGE(TAG, "Display event queue is NULL! Cannot signal main.");
                        }
                    }
                    // Update previous button states for next iteration in game over screen
                    left_button_pressed_prev_go = left_button_now_pressed;
                    right_button_pressed_prev_go = right_button_now_pressed;
                }
                break;
            }
            case DISPLAY_STATE_OTHER_INFO: {
                for (int i = 0; i < N_DISPLAY_LINES; i++) {
                    if (dirty_lines[i] && strlen(display_lines[i]) > 0) {
                        ssd1306_display_text(&dev, i, display_lines[i], strlen(display_lines[i]), false);
                        dirty_lines[i] = false;
                    }
                }
                break;
            }
            }

            ssd1306_show_buffer(&dev); // Push the entire buffer to the display
            xSemaphoreGive(xDisplayMutex);
        }
        vTaskDelay(pdMS_TO_TICKS(5));
    }
}

void display_update_text(const char *lines[], size_t num_lines) {
    if (xDisplayMutex == NULL) return;
    if (xSemaphoreTake(xDisplayMutex, portMAX_DELAY) == pdTRUE) {
        for (size_t i = 0; i < N_DISPLAY_LINES; i++) {
            const char *new_line = (i < num_lines && lines[i]) ? lines[i] : "";
            // Only update if content changed or if it's a line that needs to be cleared
            if (strncmp(display_lines[i], new_line, MAX_LINE_LENGTH) != 0 || strlen(new_line) == 0) {
                strncpy(display_lines[i], new_line, MAX_LINE_LENGTH - 1);
                display_lines[i][MAX_LINE_LENGTH - 1] = '\0';
                dirty_lines[i] = true;
            }
        }
        xSemaphoreGive(xDisplayMutex);
    }
}

// Updated function signature: now takes the button_input_queue
esp_err_t start_display_task(QueueHandle_t button_input_queue) {
    xDisplayMutex = xSemaphoreCreateMutex();
    if (xDisplayMutex == NULL) return ESP_FAIL;
    xSemaphoreGive(xDisplayMutex);

    // Store the button input queue
    s_button_input_queue = button_input_queue;

    // Create the display event queue
    display_event_queue = xQueueCreate(1, sizeof(display_event_t));
    if (display_event_queue == NULL) {
        ESP_LOGE(TAG, "Failed to create display event queue!");
        vSemaphoreDelete(xDisplayMutex);
        return ESP_FAIL;
    }

    // The display_manager_task will start in DISPLAY_STATE_MENU and display whatever
    // display_update_text provides (which will be the menu from game_menu).
    // The s_state_changed flag is already true by default, ensuring a clear on first run.

    if (xTaskCreate(display_manager_task, "Display_Manager_Task", 4096, NULL, 5, NULL) != pdPASS) {
        vSemaphoreDelete(xDisplayMutex);
        vQueueDelete(display_event_queue); // Clean up queue on task creation failure
        return ESP_FAIL;
    }
    return ESP_OK;
}

// New function to set the active game and transition the display state
void display_set_game_active(const game_t *game_instance) {
    if (game_instance == NULL) {
        ESP_LOGE(TAG, "Attempted to set NULL game instance!");
        return;
    }
    if (xDisplayMutex == NULL || xSemaphoreTake(xDisplayMutex, portMAX_DELAY) != pdTRUE) {
        ESP_LOGE(TAG, "Failed to take display mutex for setting game active.");
        return;
    }

    s_current_game = game_instance; // Store the game instance
    s_current_display_state = DISPLAY_STATE_GAME_RUNNING; // Transition to game running state
    s_state_changed = true; // Mark state changed to trigger a full clear on next iteration

    // Initialize the game and set it active
    if (s_current_game->init_game) {
        s_current_game->init_game();
    }
    if (s_current_game->set_active) {
        s_current_game->set_active(true);
    }
    // Pass the global button input queue to the game
    if (s_current_game->set_input_queue) {
        s_current_game->set_input_queue(s_button_input_queue);
    } else {
        ESP_LOGW(TAG, "Current game does not have set_input_queue function.");
    }

    s_displayed_score = -1; // Reset score display tracker for new game
    // ssd1306_clear_screen(&dev, false); // Removed: Clearing now handled by s_state_changed flag
    // Mark all lines dirty to force a full redraw by the game (handled by s_state_changed)
    // for (int i = 0; i < N_DISPLAY_LINES; i++) {
    //     dirty_lines[i] = true;
    // }

    xSemaphoreGive(xDisplayMutex);
    ESP_LOGI(TAG, "Display transitioned to GAME_RUNNING state for new game.");
}