/*
 * USB Host Keyboard Example with Display Support
 */

#include "pico/stdlib.h"
#include "pico/time.h"
#include "pico/multicore.h"
#include "tusb.h"
#include <stdio.h>
#include <ctype.h>
#include <cstring>
#include <stdlib.h>

#include "display.h"
#include "commands/echo.h"

// e-Paper library includes
extern "C" {
    #include "DEV_Config.h"
    #include "EPD_7in5b_V2.h"
    #include "GUI_Paint.h"
}

// Holds last echoed line for display
static char g_last_echo[128] = "";

// e-Paper display buffers
static UBYTE *g_epd_image = NULL;  // Black image buffer
static UBYTE *g_epd_red = NULL;    // Red image buffer

// Entry list for display
#define MAX_ENTRIES 15
#define ENTRY_LENGTH 64
typedef struct {
    char entries[MAX_ENTRIES][ENTRY_LENGTH];
    int count;
} EntryList;
static EntryList g_entry_list = {{}, 0};
static bool g_force_full_refresh = false;  // Flag to force full refresh when list is full

// Synchronization flag for multicore initialization
static volatile bool g_display_ready = false;

// Display update message structure for inter-core communication
typedef struct {
    EntryList entries;
    bool use_partial;  // Use partial refresh instead of full refresh
    UWORD xstart;      // Partial refresh region
    UWORD ystart;
    UWORD xend;
    UWORD yend;
} DisplayMessage;

void init_epaper_display() {
    printf("Initializing 7.5\" e-Paper display (B V2)...\r\n");
    
    // Initialize the hardware
    if (DEV_Module_Init() != 0) {
        printf("Failed to initialize e-Paper hardware!\r\n");
        return;
    }
    
    // Initialize the display
    printf("EPD_7IN5B_V2_Init()\r\n");
    EPD_7IN5B_V2_Init_Fast();
    printf("EPD_7IN5B_V2_Clear()\r\n");
    // EPD_7IN5B_V2_Clear();
    // DEV_Delay_ms(500);
    
    // Create image buffers for black and red content
    // 7.5" display: 800x480, 8 pixels per byte, so (800/8) * 480 = 48000 bytes each
    UWORD imagesize = (EPD_7IN5B_V2_WIDTH / 8) * EPD_7IN5B_V2_HEIGHT;
    
    g_epd_image = (UBYTE *)malloc(imagesize);
    g_epd_red = (UBYTE *)malloc(imagesize);
    
    if (g_epd_image == NULL || g_epd_red == NULL) {
        printf("Failed to allocate memory for e-Paper image buffers!\r\n");
        return;
    }
    
    // Initialize red buffer to all white (0xFF = no red pixels)
    for (UWORD i = 0; i < imagesize; i++) {
        g_epd_red[i] = 0xFF;
    }
    
    // Setup paint buffer for black content
    printf("White background\r\n");
    Paint_NewImage(g_epd_image, EPD_7IN5B_V2_WIDTH, EPD_7IN5B_V2_HEIGHT, 0, WHITE);
    Paint_SelectImage(g_epd_image);
    Paint_Clear(WHITE);
    
    // Draw header
    EPD_7IN5B_V2_Init_Part();
    EPD_7IN5B_V2_Display_Partial(g_epd_image, 0, 0, EPD_7IN5B_V2_WIDTH, EPD_7IN5B_V2_HEIGHT);
    EPD_7IN5B_V2_Display_Partial(g_epd_image, 0, 0, EPD_7IN5B_V2_WIDTH, EPD_7IN5B_V2_HEIGHT);
    EPD_7IN5B_V2_Display_Partial(g_epd_image, 0, 0, EPD_7IN5B_V2_WIDTH, EPD_7IN5B_V2_HEIGHT);

    Paint_SelectImage(g_epd_red);
    printf("Drawing header\r\n");
    Paint_DrawString_EN(10, 10, "What's new today:", &Font24, WHITE, RED);
    
    // Display the image with both black and red buffers
    EPD_7IN5B_V2_Display_Partial(g_epd_red, 0, 0, EPD_7IN5B_V2_WIDTH, EPD_7IN5B_V2_HEIGHT);
    
    printf("e-Paper display ready!\r\n");
}

/**
 * Add an entry to the display list and send update to core 1
 */
void send_display_update(const char *entry) {
    if (!g_display_ready) return;  // Don't send if display isn't ready
    
    // Check if list is at capacity
    if (g_entry_list.count >= MAX_ENTRIES) {
        // List is full - clear all entries and start fresh
        printf("[Core 0] List is FULL - clearing all entries and starting fresh\n");
        g_entry_list.count = 0;
        g_force_full_refresh = true;
    }
    
    // Add entry to global list
    strncpy(g_entry_list.entries[g_entry_list.count], entry, ENTRY_LENGTH - 1);
    g_entry_list.entries[g_entry_list.count][ENTRY_LENGTH - 1] = '\0';
    g_entry_list.count++;
    
    // Allocate message structure
    DisplayMessage *msg = (DisplayMessage *)malloc(sizeof(DisplayMessage));
    if (msg == NULL) {
        printf("Failed to allocate display message!\n");
        return;
    }
    
    // Copy entry list to message
    msg->entries = g_entry_list;
    
    // Decide whether to use partial or full refresh
    msg->use_partial = !g_force_full_refresh;
    
    if (g_force_full_refresh) {
        printf("[Core 0] Forcing full refresh this update\n");
        g_force_full_refresh = false;  // Reset flag after setting it in message
    }
    
    // Setup partial refresh region (text area at top of display)
    msg->xstart = 0;
    msg->ystart = 50;      // Start below the header
    msg->xend = 800;       // Full width
    msg->yend = 480;       // Full height from header onwards
    
    // Send message pointer to core 1
    multicore_fifo_push_blocking((uint32_t)msg);
    
    // Wait for core 1 to complete the update
    multicore_fifo_pop_blocking();
}

/**
 * Core 1 function: Runs display initialization and update handling
 */
void core1_display_init() {
    init_epaper_display();
    g_display_ready = true;
    
    printf("[Core 1] Display ready, waiting for update messages...\n");
    
    // Initialize for partial refresh on second call
    EPD_7IN5B_V2_Init_Part();
    
    // Core 1 main loop: handle display updates via FIFO
    while (true) {
        // Check if there's a message from core 0
        if (multicore_fifo_rvalid()) {
            // Read the message pointer
            uint32_t msg_addr = multicore_fifo_pop_blocking();
            DisplayMessage *msg = (DisplayMessage *)msg_addr;
            
            printf("[Core 1] Updating display with %d entries\n", msg->entries.count);
            
            // Update the e-Paper display
            if (g_epd_image != NULL && g_epd_red != NULL) {
                UWORD imagesize = (EPD_7IN5B_V2_WIDTH / 8) * EPD_7IN5B_V2_HEIGHT;
                
                // For partial refresh, only clear the text area
                UWORD width = (EPD_7IN5B_V2_WIDTH / 8);
                UWORD y_start = msg->ystart;
                UWORD y_end = msg->yend;
                
                Paint_SelectImage(g_epd_image);
                Paint_Clear(WHITE);
                if(!msg->use_partial){
                    // for loop 5 times display white partial
                    int i = 0;
                    for(i = 0; i < 1; i++){
                        EPD_7IN5B_V2_Display_Partial(g_epd_image, msg->xstart, msg->ystart, msg->xend, msg->yend);
                    }
                }
                
                // Draw all entries starting below header
                UWORD y_pos = 0;
                // paint only the last 3 entries
                int start_index = msg->entries.count > 2 ? msg->entries.count - 2 : 0;
                for (int i = start_index; i < msg->entries.count; i++) {
                    if (y_pos + 25 < 480) {  //Don't draw beyond screen
                        Paint_DrawString_EN(20, i*25, msg->entries.entries[i], &Font16, WHITE, BLACK);
                        y_pos += 25;  // Space between entries
                    } else {
                        printf("Skipping entry to avoid overflow\n");
                    }
                }
                // Use partial or full refresh
                printf("[Core 1] Using partial refresh\n");
                EPD_7IN5B_V2_Display_Partial(g_epd_image, msg->xstart, msg->ystart, msg->xend, msg->yend);
            }
            
            // Free the message (it was allocated by core 0)
            free(msg);
            
            // Signal back to core 0 that update is complete
            multicore_fifo_push_blocking(1);
        }
        
        sleep_ms(10);
    }
}

// Global DisplayManager pointer
static DisplayManager* g_display_manager = nullptr;

// Keyboard buffer
#define MAX_INPUT_LEN 64
static char g_input_buffer[MAX_INPUT_LEN];
static int g_buffer_index = 0;

// ASCII code for Backspace (often sent as 0x08)
#define ASCII_BACKSPACE 8

static const uint8_t keycode2ascii[128][2] =  {
    {0, 0}, /* 0x00 */
    {0, 0}, /* 0x01 */
    {0, 0}, /* 0x02 */
    {0, 0}, /* 0x03 */
    {'a', 'A'}, /* 0x04 */
    {'b', 'B'}, /* 0x05 */
    {'c', 'C'}, /* 0x06 */
    {'d', 'D'}, /* 0x07 */
    {'e', 'E'}, /* 0x08 */
    {'f', 'F'}, /* 0x09 */
    {'g', 'G'}, /* 0x0a */
    {'h', 'H'}, /* 0x0b */
    {'i', 'I'}, /* 0x0c */
    {'j', 'J'}, /* 0x0d */
    {'k', 'K'}, /* 0x0e */
    {'l', 'L'}, /* 0x0f */
    {'m', 'M'}, /* 0x10 */
    {'n', 'N'}, /* 0x11 */
    {'o', 'O'}, /* 0x12 */
    {'p', 'P'}, /* 0x13 */
    {'q', 'Q'}, /* 0x14 */
    {'r', 'R'}, /* 0x15 */
    {'s', 'S'}, /* 0x16 */
    {'t', 'T'}, /* 0x17 */
    {'u', 'U'}, /* 0x18 */
    {'v', 'V'}, /* 0x19 */
    {'w', 'W'}, /* 0x1a */
    {'x', 'X'}, /* 0x1b */
    {'y', 'Y'}, /* 0x1c */
    {'z', 'Z'}, /* 0x1d */
    {'1', '!'}, /* 0x1e */
    {'2', '@'}, /* 0x1f */
    {'3', '#'}, /* 0x20 */
    {'4', '$'}, /* 0x21 */
    {'5', '%'}, /* 0x22 */
    {'6', '^'}, /* 0x23 */
    {'7', '&'}, /* 0x24 */
    {'8', '*'}, /* 0x25 */
    {'9', '('}, /* 0x26 */
    {'0', ')'}, /* 0x27 */
    {'\r', '\r'}, /* 0x28 ENTER */
    {'\x1b', '\x1b'}, /* 0x29 ESCAPE */
    {'\b', '\b'}, /* 0x2a BACKSPACE */
    {'\t', '\t'}, /* 0x2b TAB */
    {' ', ' '}, /* 0x2c SPACE */
    {'-', '_'}, /* 0x2d MINUS */
    {'=', '+'}, /* 0x2e EQUAL */
    {'[', '{'}, /* 0x2f BRACKET_LEFT */
    {']', '}'}, /* 0x30 BRACKET_RIGHT */
    {'\\', '|'}, /* 0x31 BACKSLASH */
    {'#', '~'}, /* 0x32 EUROPE_1 */
    {';', ':'}, /* 0x33 SEMICOLON */
    {'\'', '\"'}, /* 0x34 APOSTROPHE */
    {'`', '~'}, /* 0x35 GRAVE */
    {',', '<'}, /* 0x36 COMMA */
    {'.', '>'}, /* 0x37 PERIOD */
    {'/', '?'}, /* 0x38 SLASH */
    {0, 0}, /* 0x39 CAPS_LOCK */
    {0, 0}, /* 0x3a F1 */
    {0, 0}, /* 0x3b F2 */
    {0, 0}, /* 0x3c F3 */
    {0, 0}, /* 0x3d F4 */
    {0, 0}, /* 0x3e F5 */
    {0, 0}, /* 0x3f F6 */
    {0, 0}, /* 0x40 F7 */
    {0, 0}, /* 0x41 F8 */
    {0, 0}, /* 0x42 F9 */
    {0, 0}, /* 0x43 F10 */
    {0, 0}, /* 0x44 F11 */
    {0, 0}, /* 0x45 F12 */
    {0, 0}, /* 0x46 PRINT_SCREEN */
    {0, 0}, /* 0x47 SCROLL_LOCK */
    {0, 0}, /* 0x48 PAUSE */
    {0, 0}, /* 0x49 INSERT */
    {0, 0}, /* 0x4a HOME */
    {0, 0}, /* 0x4b PAGE_UP */
    {0, 0}, /* 0x4c DELETE */
    {0, 0}, /* 0x4d END */
    {0, 0}, /* 0x4e PAGE_DOWN */
    {0, 0}, /* 0x4f RIGHT_ARROW */
    {0, 0}, /* 0x50 LEFT_ARROW */
    {0, 0}, /* 0x51 DOWN_ARROW */
    {0, 0}, /* 0x52 UP_ARROW */
    {0, 0}, /* 0x53 NUM_LOCK */
    {'/', '/'}, /* 0x54 KP_DIVIDE */
    {'*', '*'}, /* 0x55 KP_MULTIPLY */
    {'-', '-'}, /* 0x56 KP_MINUS */
    {'+', '+'}, /* 0x57 KP_PLUS */
    {'\r', '\r'}, /* 0x58 KP_ENTER */
    {'1', '1'}, /* 0x59 KP_1 */
    {'2', '2'}, /* 0x5a KP_2 */
    {'3', '3'}, /* 0x5b KP_3 */
    {'4', '4'}, /* 0x5c KP_4 */
    {'5', '5'}, /* 0x5d KP_5 */
    {'6', '6'}, /* 0x5e KP_6 */
    {'7', '7'}, /* 0x5f KP_7 */
    {'8', '8'}, /* 0x60 KP_8 */
    {'9', '9'}, /* 0x61 KP_9 */
    {'0', '0'}, /* 0x62 KP_0 */
    {'.', '.'}, /* 0x63 KP_DECIMAL */
};

static inline bool find_key_in_report(hid_keyboard_report_t const *report, uint8_t keycode) {
    for(uint8_t i=0; i<6; i++) {
        if (report->keycode[i] == keycode) return true;
    }
    return false;
}

static void process_kbd_report(hid_keyboard_report_t const *report) {
    static hid_keyboard_report_t prev_report = { 0, 0, {0} };

    for(uint8_t i=0; i<6; i++) {
        if (report->keycode[i]) {
            if (find_key_in_report(&prev_report, report->keycode[i])) {
                // exist in previous report means the current key is holding
            } else {
                // not existed in previous report means the current key is pressed
                bool const is_shift = report->modifier & (KEYBOARD_MODIFIER_LEFTSHIFT | KEYBOARD_MODIFIER_RIGHTSHIFT);
                uint8_t ch = keycode2ascii[report->keycode[i]][is_shift ? 1 : 0];
                
                if (ch) {
                    printf("%c", ch);
                    
                    if (ch == '\r' || ch == '\n') {
                        printf("\n");
                        // Process command
                        g_input_buffer[g_buffer_index] = '\0';
                        
                        // Echo logic
                        echo_and_reset(g_input_buffer, &g_buffer_index);
                        
                        // Save last echoed (all caps) for display
                        size_t k = 0;
                        for (; g_input_buffer[k] != '\0' && k < sizeof(g_last_echo) - 1; ++k) {
                            g_last_echo[k] = g_input_buffer[k];
                        }
                        g_last_echo[k] = '\0';
                        
                        // Update displays
                        if (g_display_manager) {
                            g_display_manager->set_last_echo(g_last_echo);
                            g_display_manager->refresh("", g_last_echo);
                        }
                        send_display_update(g_last_echo);
                        
                        // Reset buffer
                        g_buffer_index = 0;
                        g_input_buffer[0] = '\0';
                        
                    } else if (ch == '\b' || ch == 127) {
                        if (g_buffer_index > 0) {
                            g_buffer_index--;
                            g_input_buffer[g_buffer_index] = '\0';
                            if (g_display_manager) {
                                g_display_manager->refresh(g_input_buffer, g_last_echo);
                            }
                        }
                    } else if (g_buffer_index < MAX_INPUT_LEN - 1) {
                        g_input_buffer[g_buffer_index] = ch;
                        g_buffer_index++;
                        g_input_buffer[g_buffer_index] = '\0';
                        if (g_display_manager) {
                            g_display_manager->refresh(g_input_buffer, g_last_echo);
                        }
                    }
                }
            }
        }
    }
    prev_report = *report;
}

//--------------------------------------------------------------------+
// TinyUSB Callbacks
//--------------------------------------------------------------------+

void tuh_hid_mount_cb(uint8_t dev_addr, uint8_t instance, uint8_t const* desc_report, uint16_t desc_len) {
    printf("HID device address = %d, instance = %d is mounted\r\n", dev_addr, instance);

    uint8_t const itf_protocol = tuh_hid_interface_protocol(dev_addr, instance);
    if (itf_protocol == HID_ITF_PROTOCOL_KEYBOARD) {
        printf("Keyboard mounted\r\n");
        if (g_display_manager) {
            g_display_manager->refresh("Keyboard Connected", nullptr);
        }
        if (!tuh_hid_receive_report(dev_addr, instance)) {
            printf("Error: cannot request to receive report\r\n");
        }
    }
}

void tuh_hid_umount_cb(uint8_t dev_addr, uint8_t instance) {
    printf("HID device address = %d, instance = %d is unmounted\r\n", dev_addr, instance);
    if (g_display_manager) {
        g_display_manager->refresh("Keyboard Disconnected", nullptr);
    }
}

void tuh_hid_report_received_cb(uint8_t dev_addr, uint8_t instance, uint8_t const* report, uint16_t len) {
    uint8_t const itf_protocol = tuh_hid_interface_protocol(dev_addr, instance);

    if (itf_protocol == HID_ITF_PROTOCOL_KEYBOARD) {
        process_kbd_report((hid_keyboard_report_t const*) report);
    }

    // continue to request to receive report
    if (!tuh_hid_receive_report(dev_addr, instance)) {
        printf("Error: cannot request to receive report\r\n");
    }
}

int main() {
    stdio_init_all();
    sleep_ms(3000); // Give time for power to settle and serial to connect
    printf("System Booting...\n");
    
    // Launch display initialization on core 1
    printf("Launching e-Paper display init on core 1...\n");
    multicore_launch_core1(core1_display_init);
    
    DisplayManager display;
    display.init();
    g_display_manager = &display;
    
    display.refresh("Waiting for Keyboard", nullptr);
    
    printf("Initializing TinyUSB Host...\n");
    tuh_init(BOARD_TUH_RHPORT);
    printf("TinyUSB Host Initialized.\n");

    uint32_t last_print = 0;
    while (true) {
        tuh_task();
        
        uint32_t now = to_ms_since_boot(get_absolute_time());
        if (now - last_print > 5000) {
            printf("Heartbeat: %u\n", now);
            last_print = now;
        }
    }

    return 0;
}