/*
 * 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] = "";

#include "command_processor.h"
#include "epaper_manager.h"
#include "wifi_manager.h"
#include "tcp_debug.h"
#include "keyboard_stdio.h"

// Global display manager instance
DisplayManager* g_display_manager = nullptr;

// Command line buffer
#define MAX_INPUT_LEN 64
static char g_cmd_buffer[MAX_INPUT_LEN];
static int g_cmd_index = 0;

// Image reception state
static bool g_receiving_image = false;
static uint8_t *g_image_buffer = nullptr;
static uint32_t g_image_bytes_received = 0;
static const uint32_t IMAGE_SIZE = 15000; // 400x300 / 8

void on_tcp_connection_change(bool connected) {
    if (!connected) {
        if (g_receiving_image) {
            printf("TCP disconnected, resetting image reception\n");
            g_receiving_image = false;
            if (g_image_buffer) {
                free(g_image_buffer);
                g_image_buffer = nullptr;
            }
        }
        // Also reset command buffer to avoid partial commands
        g_cmd_index = 0;
        g_cmd_buffer[0] = '\0';
    }
}

static bool execute_command(CommandAction action, const char* input) {
    switch (action) {
        case CMD_IMAGE:
            printf("Command: /image - Waiting for %d bytes...\n", IMAGE_SIZE);
            if (g_image_buffer) free(g_image_buffer);
            g_image_buffer = (uint8_t*)malloc(IMAGE_SIZE);
            if (!g_image_buffer) {
                printf("Failed to allocate image buffer!\n");
                return true;
            }
            g_receiving_image = true;
            g_image_bytes_received = 0;
            return true;
        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);
                    tcp_debug_init();
                    tcp_set_connection_callback(on_tcp_connection_change);
                }
            } 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.ascii) {
            keyboard_stdio_push_char(event.ascii);
        }
    }
}

//--------------------------------------------------------------------+
// 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_cmd_index = 0;
    g_cmd_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");
    }
}

static void process_input_char(char c) {
    if (g_receiving_image) {
        if (g_image_buffer) {
            g_image_buffer[g_image_bytes_received++] = (uint8_t)c;
            if (g_image_bytes_received >= IMAGE_SIZE) {
                printf("\nImage received (%d bytes). Displaying...\n", g_image_bytes_received);
                epaper_display_full_image(g_image_buffer, IMAGE_SIZE);
                g_receiving_image = false;
                g_image_buffer = nullptr; // Ownership transferred
            }
        } else {
            g_receiving_image = false; // Error state
        }
        return;
    }

    if (c == 0) return; // Ignore null characters in text mode

    if (c == '\b' || c == 127) {
        if (g_cmd_index > 0) {
            printf("\b \b");
            g_cmd_index--;
            g_cmd_buffer[g_cmd_index] = '\0';
            
            // Update OLED
            if (g_display_manager) g_display_manager->refresh(g_cmd_buffer, g_last_echo);
            
            // Update e-Paper (in-place)
            epaper_send_update(g_cmd_buffer, false);
        }
    } else if (c == '\r' || c == '\n') {
        printf("\n");
        
        // Debug: Print buffer content
        printf("Processing line: '%s' (len=%d)\n", g_cmd_buffer, g_cmd_index);
        printf("Hex: ");
        for(int i=0; i<g_cmd_index; i++) printf("%02X ", (unsigned char)g_cmd_buffer[i]);
        printf("\n");

        // Always commit the input line to e-Paper so commands are visible
        if (g_cmd_index > 0) {
            epaper_send_update(g_cmd_buffer, true);
        }
        
        CommandAction action = parse_command(g_cmd_buffer);
        printf("Parsed Action: %d\n", action); // Debug action

        if (execute_command(action, g_cmd_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_cmd_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_cmd_index = 0;
        g_cmd_buffer[0] = '\0';
    } else {
        printf("%c", c);
        if (g_cmd_index < MAX_INPUT_LEN - 1) {
            g_cmd_buffer[g_cmd_index++] = c;
            g_cmd_buffer[g_cmd_index] = '\0';
            
            // Update OLED
            if (g_display_manager) g_display_manager->refresh(g_cmd_buffer, g_last_echo);
            
            // Update e-Paper on every keystroke
            epaper_send_update(g_cmd_buffer, false);
        }
    }
}

int main() {
    stdio_init_all();
    keyboard_stdio_init();
    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();
        
        int c = getchar_timeout_us(0);
        if (c != PICO_ERROR_TIMEOUT) {
            process_input_char((char)c);
        }

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