#include "serial_uploader.h"
#include "game_launcher.h"
#include "lua_game_loader.h"
#include "sd_card.h"
#include "pico/stdlib.h"
#include <cstring>
#include <cstdio>
#include <cctype>

// Maximum file size: 64KB (should be plenty for Lua games)
#define MAX_FILE_SIZE (64 * 1024)

SerialUploader::SerialUploader(GameLauncher* launcher)
    : state(IDLE)
    , game_launcher(launcher)
    , file_size(0)
    , bytes_received(0)
    , file_buffer(nullptr)
    , base64_index(0)
{
    filename[0] = '\0';
    last_uploaded_name[0] = '\0';
}

void SerialUploader::reset() {
    state = IDLE;
    file_size = 0;
    bytes_received = 0;
    base64_index = 0;
    filename[0] = '\0';

    if (file_buffer) {
        delete[] file_buffer;
        file_buffer = nullptr;
    }
}

uint8_t SerialUploader::decode_base64_char(char c) {
    if (c >= 'A' && c <= 'Z') return c - 'A';
    if (c >= 'a' && c <= 'z') return c - 'a' + 26;
    if (c >= '0' && c <= '9') return c - '0' + 52;
    if (c == '+') return 62;
    if (c == '/') return 63;
    return 0;
}

void SerialUploader::decode_base64_block(const char* input, uint8_t* output) {
    uint32_t combined = (decode_base64_char(input[0]) << 18) |
                        (decode_base64_char(input[1]) << 12) |
                        (decode_base64_char(input[2]) << 6) |
                        decode_base64_char(input[3]);

    output[0] = (combined >> 16) & 0xFF;
    output[1] = (combined >> 8) & 0xFF;
    output[2] = combined & 0xFF;
}

bool SerialUploader::write_file_to_sd() {
    if (bytes_received == 0) {
        printf("ERROR No data received!\n");
        return false;
    }

    if (!file_buffer) {
        printf("ERROR File buffer is NULL!\n");
        return false;
    }

    // Set SPI speed to SD card speed (12.5 MHz) and save previous speed
    uint prev_speed = sd_card_set_spi_speed();

    // Ensure /games directory exists
    FRESULT fr = f_mkdir("/games");
    if (fr != FR_OK && fr != FR_EXIST) {
        printf("ERROR Failed to create /games directory: %d\n", fr);
        // Try to continue anyway in case it already exists
    }

    // Build file path - overwrite if file exists
    char filepath[128];
    snprintf(filepath, sizeof(filepath), "/games/%s", filename);

    FIL fil;
    fr = f_open(&fil, filepath, FA_CREATE_ALWAYS | FA_WRITE);

    if (fr != FR_OK) {
        printf("ERROR Failed to open file for writing: %d\n", fr);
        sd_card_restore_spi_speed(prev_speed);
        return false;
    }

    // Write data in chunks (more reliable for large files)
    const uint32_t CHUNK_SIZE = 512;
    uint32_t total_written = 0;

    while (total_written < bytes_received) {
        uint32_t chunk = (bytes_received - total_written > CHUNK_SIZE) ? CHUNK_SIZE : (bytes_received - total_written);

        UINT bytes_written = 0;
        fr = f_write(&fil, file_buffer + total_written, chunk, &bytes_written);

        if (fr != FR_OK) {
            printf("ERROR f_write failed at byte %u: error %d\n", total_written, fr);
            f_close(&fil);
            sd_card_restore_spi_speed(prev_speed);
            return false;
        }

        if (bytes_written != chunk) {
            printf("ERROR Partial write: wrote %u/%u bytes at offset %u\n", bytes_written, chunk, total_written);
            f_close(&fil);
            sd_card_restore_spi_speed(prev_speed);
            return false;
        }

        total_written += bytes_written;
    }

    // Sync to ensure data is written to SD card
    fr = f_sync(&fil);
    if (fr != FR_OK) {
        printf("ERROR f_sync failed: %d\n", fr);
    }

    f_close(&fil);

    printf("✓ Wrote %u bytes to %s\n", total_written, filepath);

    // Restore display SPI speed
    sd_card_restore_spi_speed(prev_speed);

    return true;
}

void SerialUploader::launch_game() {
    // Extract base game name (remove .lua extension)
    strncpy(last_uploaded_name, filename, sizeof(last_uploaded_name) - 1);
    last_uploaded_name[sizeof(last_uploaded_name) - 1] = '\0';

    // Remove .lua extension
    char* ext = strstr(last_uploaded_name, ".lua");
    if (ext) {
        *ext = '\0';
    }

    printf("Re-scanning Lua games to pick up new file...\n");

    // Clear existing games before re-scanning (prevents duplicates and memory leaks)
    game_launcher->clear_games();
    LuaGameLoader::clear_factory_data();

    // Re-scan Lua games to pick up the newly uploaded file
    // Note: LuaGameLoader::register_all_games handles SPI speed internally
    int new_games = LuaGameLoader::register_all_games(game_launcher);
    printf("Found %d Lua games after re-scan\n", new_games);
}

bool SerialUploader::complete_launch() {
    // This should only be called when it's safe (no display refresh in progress)
    // Now try to launch the newly uploaded game by name
    printf("Attempting to launch game: %s\n", last_uploaded_name);
    bool launched = game_launcher->select_game_by_name(last_uploaded_name);

    if (launched) {
        printf("LAUNCHED %s\n", last_uploaded_name);
    } else {
        printf("ERROR Failed to launch game: %s\n", last_uploaded_name);
    }

    // Reset state back to IDLE
    reset();

    return launched;
}

bool SerialUploader::process() {
    if (state == IDLE) {
        // Check for "UPLOAD" command
        int c = getchar_timeout_us(0);
        if (c == PICO_ERROR_TIMEOUT) return false;

        if (c == 'U') {
            // Check for "UPLOAD "
            static char cmd_buffer[8];
            cmd_buffer[0] = c;
            int idx = 1;

            // Read "PLOAD "
            for (int i = 0; i < 6; i++) {
                c = getchar_timeout_us(100000); // 100ms timeout
                if (c == PICO_ERROR_TIMEOUT) return false;
                cmd_buffer[idx++] = c;
            }
            cmd_buffer[idx] = '\0';

            if (strcmp(cmd_buffer, "UPLOAD ") == 0) {
                state = RECEIVING_FILENAME;
                filename[0] = '\0';
                return false;
            }
        }
    }

    if (state == RECEIVING_FILENAME) {
        // Read filename until space
        int idx = strlen(filename);
        while (idx < sizeof(filename) - 1) {
            int c = getchar_timeout_us(100000);
            if (c == PICO_ERROR_TIMEOUT) break;

            if (c == ' ') {
                filename[idx] = '\0';
                state = RECEIVING_SIZE;
                file_size = 0;
                return false;
            }

            filename[idx++] = (char)c;
        }
    }

    if (state == RECEIVING_SIZE) {
        // Read file size until newline
        char size_buffer[16];
        int idx = 0;

        while (idx < sizeof(size_buffer) - 1) {
            int c = getchar_timeout_us(100000);
            if (c == PICO_ERROR_TIMEOUT) break;

            if (c == '\n' || c == '\r') {
                size_buffer[idx] = '\0';
                file_size = atoi(size_buffer);

                if (file_size == 0 || file_size > MAX_FILE_SIZE) {
                    printf("ERROR Invalid file size: %u\n", file_size);
                    reset();
                    return false;
                }

                // Allocate buffer for decoded data
                // file_size is the ORIGINAL (decoded) size, so allocate exactly that
                file_buffer = new uint8_t[file_size];

                if (!file_buffer) {
                    printf("ERROR Failed to allocate %u bytes for file buffer\n", file_size);
                    reset();
                    return false;
                }

                bytes_received = 0;
                base64_index = 0;

                state = RECEIVING_DATA;
                printf("Receiving %u bytes for %s...\n", file_size, filename);
                return false;
            }

            size_buffer[idx++] = (char)c;
        }
    }

    if (state == RECEIVING_DATA) {
        // Read base64 data until "END"
        while (true) {
            int c = getchar_timeout_us(1000);
            if (c == PICO_ERROR_TIMEOUT) break;

            // Check for "END" marker
            static char end_check[4] = {0, 0, 0, 0};
            end_check[0] = end_check[1];
            end_check[1] = end_check[2];
            end_check[2] = end_check[3];
            end_check[3] = (char)c;

            if (end_check[0] == '\n' && end_check[1] == 'E' && end_check[2] == 'N' && end_check[3] == 'D') {
                state = WRITING_FILE;
                return false;
            }

            // Skip whitespace
            if (isspace(c)) continue;

            // Accumulate base64 characters
            if ((c >= 'A' && c <= 'Z') || (c >= 'a' && c <= 'z') ||
                (c >= '0' && c <= '9') || c == '+' || c == '/' || c == '=') {

                base64_buffer[base64_index++] = (char)c;

                // Decode every 4 characters
                if (base64_index == 4) {
                    uint8_t decoded[3];
                    decode_base64_block(base64_buffer, decoded);

                    // Handle padding
                    int decoded_count = 3;
                    if (base64_buffer[2] == '=') decoded_count = 1;
                    else if (base64_buffer[3] == '=') decoded_count = 2;

                    // Copy decoded bytes
                    for (int i = 0; i < decoded_count; i++) {
                        if (bytes_received < file_size) {
                            file_buffer[bytes_received++] = decoded[i];
                        }
                    }

                    base64_index = 0;
                }
            }
        }
    }

    if (state == WRITING_FILE) {
        if (write_file_to_sd()) {
            state = LAUNCHING_GAME;
            // Prepare for launch by scanning games, but don't actually launch yet
            launch_game();
        } else {
            reset();
        }
        return false;
    }

    if (state == LAUNCHING_GAME) {
        // Stay in this state until main loop calls complete_launch() when safe
        return false;
    }

    return false;
}