/*
 * 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 <malloc.h>

#include "display.h"
#include "commands/echo.h"
#include "keyboard_input.h"
#include "command_processor.h"
#include "epaper_manager.h"
#include "wifi_manager.h"

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

// 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;

static bool execute_command(CommandAction action, const char* input) {
    switch (action) {
        case CMD_REFRESH:
            printf("Command: /refresh\n");
            epaper_force_refresh();
            return true;
        case CMD_CLEAR:
            printf("Command: /clear\n");
            epaper_clear();
            return true;
        case CMD_SCAN:
            printf("Command: /scan\n");
            if (g_display_manager) g_display_manager->refresh("Scanning WiFi...", nullptr);
            epaper_send_update("Scanning WiFi...", true);
            wifi_scan();
            return true;
        case CMD_CONNECT: {
            printf("Command: /connect\n");
            char ssid[33] = {0};
            char password[64] = {0};
            // Skip "/connect " (9 chars)
            const char* args = input + 9;
            
            // Simple parsing: first word is ssid, rest is password
            int parsed = sscanf(args, "%32s %63s", ssid, password);
            if (parsed < 2) {
                printf("Usage: /connect <ssid> <password>\n");
                if (g_display_manager) g_display_manager->refresh("Usage: /connect <ssid> <pass>", nullptr);
                epaper_send_update("Usage: /connect <ssid> <pass>", true);
                return true;
            }
            
            if (g_display_manager) g_display_manager->refresh("Connecting...", ssid);
            // epaper_send_update("Connecting...", true);
            
            if (wifi_connect(ssid, password)) {
                if (g_display_manager) g_display_manager->refresh("Connected!", ssid);
                int rc = wifi_save_credentials(ssid, password);
                if (rc != 0) {
                    char err_msg[32];
                    snprintf(err_msg, sizeof(err_msg), "Flash Error: %d", rc);
                    if (g_display_manager) g_display_manager->refresh(err_msg, nullptr);
                    epaper_send_update(err_msg, true);
                } else {
                    epaper_send_update("Connected!", true);
                }
            } else {
                if (g_display_manager) g_display_manager->refresh("Connection Failed", nullptr);
                epaper_send_update("Connection Failed", true);
            }
            return true;
        }
        case CMD_STATUS: {
            printf("Command: /status\n");
            struct mallinfo m = mallinfo();
            char status_msg[64];
            // fordblks is the free chunk size in the arena. 
            // Note: This might not account for the total available system RAM if the heap hasn't grown to fill it yet.
            // But it gives an idea of fragmentation and available malloc-able memory within the current arena.
            snprintf(status_msg, sizeof(status_msg), "Heap: %d B, IP: %s", m.fordblks, wifi_get_ip());
            printf("%s\n", status_msg);
            if (g_display_manager) g_display_manager->refresh(status_msg, nullptr);
            epaper_send_update(status_msg, true);
            return true;
        }
        default:
            return false;
    }
}

static void process_kbd_report(hid_keyboard_report_t const *report) {
    KeyEvent event;
    if (parse_keyboard_report(report, &event)) {
        if (event.is_backspace) {
            if (g_buffer_index > 0) {
                printf("\b \b");
                g_buffer_index--;
                g_input_buffer[g_buffer_index] = '\0';
                
                // Update OLED
                if (g_display_manager) g_display_manager->refresh(g_input_buffer, g_last_echo);
                
                // Update e-Paper (in-place)
                epaper_send_update(g_input_buffer, false);
            }
        } else if (event.is_enter) {
            printf("\n");
            
            // Always commit the input line to e-Paper so commands are visible
            if (g_buffer_index > 0) {
                epaper_send_update(g_input_buffer, true);
            }
            
            CommandAction action = parse_command(g_input_buffer);
            if (execute_command(action, g_input_buffer)) {
                // Command handled
                if (g_display_manager) {
                    g_display_manager->set_last_echo("Command Executed");
                    g_display_manager->refresh("", "Command Executed");
                }
            } else {
                // Not a command, just text (already committed to e-Paper above)
                
                // Save to last echo
                strncpy(g_last_echo, g_input_buffer, sizeof(g_last_echo) - 1);
                g_last_echo[sizeof(g_last_echo) - 1] = '\0';
                
                // Update OLED
                if (g_display_manager) {
                    g_display_manager->set_last_echo(g_last_echo);
                    g_display_manager->refresh("", g_last_echo);
                }
            }
            
            // Clear buffer
            g_buffer_index = 0;
            g_input_buffer[0] = '\0';
        } else if (event.is_printable) {
            printf("%c", event.ascii);
            if (g_buffer_index < MAX_INPUT_LEN - 1) {
                g_input_buffer[g_buffer_index++] = event.ascii;
                g_input_buffer[g_buffer_index] = '\0';
                
                // Update OLED
                if (g_display_manager) g_display_manager->refresh(g_input_buffer, g_last_echo);
                
                // Update e-Paper on every keystroke
                epaper_send_update(g_input_buffer, false);
            }
        }
    }
}

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

static bool g_keyboard_mounted = false;

// Invoked when device is suspended
void tuh_device_suspend_cb(uint8_t dev_addr) {
    printf("Device address = %d suspended\r\n", dev_addr);
}

// Invoked when device is resumed
void tuh_device_resume_cb(uint8_t dev_addr) {
    printf("Device address = %d resumed\r\n", dev_addr);
}

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");
        g_keyboard_mounted = true;
        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);
    g_keyboard_mounted = false;
    
    // Reset input buffer
    g_buffer_index = 0;
    g_input_buffer[0] = '\0';

    if (g_display_manager) {
        g_display_manager->refresh("Waiting for Keyboard", 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, WIFI and Flash safe execute
    epaper_start_background_thread();
    
    DisplayManager display;
    display.init();
    g_display_manager = &display;
    
    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, Mounted: %d\n", now, g_keyboard_mounted);
            last_print = now;
        }
    }

    return 0;
}