#include "startscreen.h"
#include "game.h"
#include "graphics.h"
#include "sound.h"
#include <string.h>
#include "gfx/world_lab.h"
#include "gfx/world_ice.h"
#include "gfx/world_cave.h"
#include "gfx/world_city.h"
#include "gfx/world_mansion.h"
#include "gfx/world_powerp.h"
#include "gfx/world_woods.h"

// Structure for mapping letter tiles to their widths
typedef struct {
    int tile;
    int width;
} LetterWidth;

// Array of letter widths for text rendering
static const LetterWidth letterWidths[] = {
    {LETTER_A, 7}, {LETTER_B, 7}, {LETTER_C, 7}, {LETTER_D, 7},
    {LETTER_E, 7}, {LETTER_F, 7}, {LETTER_G, 7}, {LETTER_H, 7},
    {LETTER_I, 6}, {LETTER_J, 7}, {LETTER_K, 7}, {LETTER_L, 7},
    {LETTER_M, 9}, {LETTER_N, 8}, {LETTER_O, 7}, {LETTER_P, 7},
    {LETTER_Q, 8}, {LETTER_R, 7}, {LETTER_S, 7}, {LETTER_T, 8},
    {LETTER_U, 7}, {LETTER_V, 7}, {LETTER_W, 9}, {LETTER_X, 7},
    {LETTER_Y, 7}, {LETTER_Z, 7},
    {SYMBOL_EX, 4}, {SYMBOL_QUESTION, 7}, {SYMBOL_COLON, 4}, {SYMBOL_PERIOD, 4}
};
static const int letterWidthCount = sizeof(letterWidths) / sizeof(letterWidths[0]);

// Get width of a letter tile
static int getLetterWidth(int tile) {
    for (int i = 0; i < letterWidthCount; i++) {
        if (letterWidths[i].tile == tile) return letterWidths[i].width;
    }
    return 9;
}

// Structure for menu option rendering
typedef struct {
    const int* tiles;
    int len;
    int x;
    int y;
    int spriteStartIndex;
} MenuOption;

// Render a menu option with centered text
static int renderMenuOption(MenuOption* option, int startX, int startY) {
    int width = 0, xOffset = 0, SPACE_WIDTH = 7;
    const int SPACE_TILE = -1;
    for (int i = 0; i < option->len; i++) {
        width += (option->tiles[i] != SPACE_TILE) ? getLetterWidth(option->tiles[i]) : SPACE_WIDTH;
        if (i < option->len - 1) width += 2;
    }
    option->x = startX ? startX : (SCREEN_WIDTH - width) / 2;
    option->y = startY;
    for (int i = 0; i < option->len; i++) {
        if (option->tiles[i] != SPACE_TILE) {
            setSprite(option->spriteStartIndex + i, option->x + xOffset, option->y, option->tiles[i]);
            xOffset += getLetterWidth(option->tiles[i]) + 2;
        } else {
            xOffset += SPACE_WIDTH + 2;
        }
    }
    return width;
}

// Character-to-tile mapping for text rendering
static const int charToTile[128] = {
    [' '] = -1,
    ['A'] = LETTER_A, ['B'] = LETTER_B, ['C'] = LETTER_C, ['D'] = LETTER_D,
    ['E'] = LETTER_E, ['F'] = LETTER_F, ['G'] = LETTER_G, ['H'] = LETTER_H,
    ['I'] = LETTER_I, ['J'] = LETTER_J, ['K'] = LETTER_K, ['L'] = LETTER_L,
    ['M'] = LETTER_M, ['N'] = LETTER_N, ['O'] = LETTER_O, ['P'] = LETTER_P,
    ['Q'] = LETTER_Q, ['R'] = LETTER_R, ['S'] = LETTER_S, ['T'] = LETTER_T,
    ['U'] = LETTER_U, ['V'] = LETTER_V, ['W'] = LETTER_W, ['X'] = LETTER_X,
    ['Y'] = LETTER_Y, ['Z'] = LETTER_Z,
    ['!'] = SYMBOL_EX, ['?'] = SYMBOL_QUESTION, [':'] = SYMBOL_COLON, ['.'] = SYMBOL_PERIOD,
    ['\n'] = -2
};

// Display text screen with centered text
static void showTextScreen(const char* text) {
    for (int i = 0; i < 128; i++) hideSprite(i);

    int textLen = strlen(text), textTiles[textLen], lineCount = 1;
    for (int i = 0; i < textLen; i++) {
        textTiles[i] = charToTile[(unsigned char)text[i]];
        if (textTiles[i] == -2) lineCount++;
    }

    int maxLineWidth = 0, currentLineWidth = 0, lineWidths[lineCount], lineIndex = 0;
    const int SPACE_WIDTH = 7;
    for (int i = 0; i < textLen; i++) {
        if (textTiles[i] == -2) {
            lineWidths[lineIndex++] = currentLineWidth;
            if (currentLineWidth > maxLineWidth) maxLineWidth = currentLineWidth;
            currentLineWidth = 0;
        } else if (textTiles[i] != -1) {
            currentLineWidth += getLetterWidth(textTiles[i]) + 2;
        } else {
            currentLineWidth += SPACE_WIDTH + 2;
        }
    }
    lineWidths[lineIndex] = currentLineWidth;
    if (currentLineWidth > maxLineWidth) maxLineWidth = currentLineWidth;

    int totalHeight = lineCount * SPRITE_SIZE + (lineCount - 1) * 4;
    int textX = (SCREEN_WIDTH - maxLineWidth) / 2;
    int textY = (SCREEN_HEIGHT - totalHeight) / 2;

    int xOffset = 0, spriteIndex = 0;
    int currentY = textY;
    for (int i = 0; i < textLen; i++) {
        if (textTiles[i] == -2) {
            xOffset = 0;
            currentY += SPRITE_SIZE + 4;
        } else if (textTiles[i] != -1) {
            setSprite(spriteIndex++, textX + xOffset, currentY, textTiles[i]);
            xOffset += getLetterWidth(textTiles[i]) + 2;
        } else {
            xOffset += SPACE_WIDTH + 2;
        }
    }
    updateOAM();

    u16 prevKeys = 0;
    while (1) {
        u16 keys = ~REG_KEYINPUT & 0x03FF;
        VBlankIntrWait();
        mmFrame();
        if (keys & KEY_B && !(prevKeys & KEY_B)) {
            break;
        }
        prevKeys = keys;
    }
}

// Restore menu options and pokeball sprite
static void restoreMenu(MenuOption* options, int optionCount, int pokeballIndex, int pokeballX, int pokeballY) {
    for (int i = 0; i < 128; i++) hideSprite(i);
    for (int i = 0; i < optionCount; i++) {
        renderMenuOption(&options[i], options[i].x, options[i].y);
    }
    setSprite(pokeballIndex, pokeballX, pokeballY, POKEBALL);
    updateOAM();
}

// Clear VRAM and hide all sprites
static void clearVRAMandSprites(void) {
    memset((void*)0x06000000, 0, 0xA000);
    for (int i = 0; i < 128; i++) {
        shadowOAM[i].attr0 = 0x0200;
    }
    updateOAM();
    memset((void*)0x06010000, 0, 0x4000);
}

// Shut down display and sound
static void shutdownDisplay(void) {
    mmSetModuleVolume(0);
    mmStop();
    mmEffectCancelAll();
    mmFrame();
    REG_SOUNDCNT_H = 0;
    REG_SOUNDCNT_X = 0;
    VBlankIntrWait();
    for (int i = 0; i < 16; i++) BG_PAL[i] = 0x0000;
    for (int i = 0; i < 256; i++) PAL_OBJ[i] = 0x0000;
    memset((void*)0x06000000, 0, 0xA000);
    REG_DISPCNT = 0;
}

// Load and display level-specific screen
static void loadLevelScreen(const void* worldData, int worldSize, const char* musicName, int musicId) {
    REG_DISPCNT = 0x0403;
    volatile u16* vram = (volatile u16*)0x06000000;
    memcpy((void*)vram, worldData, worldSize);
    playMusic(musicName, musicId);
    delay(180);

    while (1) {
        u16 keys = ~REG_KEYINPUT & 0x03FF;
        if (keys & (KEY_A | KEY_START)) {
            shutdownDisplay();
            break;
        }
        VBlankIntrWait();
        mmFrame();
    }
}

// Configuration for level-specific screens
typedef struct {
    int level;
    const void* worldData;
    int worldSize;
    const char* musicName;
    int musicId;
} LevelConfig;

static const LevelConfig levelConfigs[] = {
    {0, worldlab, WORLDLAB_SIZE, "LAB", WORLD_LAB},
    {11, worldice, WORLDICE_SIZE, "ICE", WORLD_ICE},
    {21, worldcave, WORLDCAVE_SIZE, "CAVE", WORLD_CAVE},
    {31, worldwoods, WORLDWOODS_SIZE, "WOODS", WORLD_WOODS},
    {41, worldcity, WORLDCITY_SIZE, "CITY", WORLD_CITY},
    {51, worldmansion, WORLDMANSION_SIZE, "MANSION", WORLD_MANSION},
    {62, worldpowerp, WORLDPOWERP_SIZE, "POWERP", WORLD_POWERP}
};
static const int levelConfigCount = sizeof(levelConfigs) / sizeof(levelConfigs[0]);

// Load and start a specific level
void loadAndStartLevel(int initialLevel) {
    currentLevel = initialLevel;

    for (int i = 0; i < levelConfigCount; i++) {
        if (levelConfigs[i].level == currentLevel) {
            clearVRAMandSprites();
            loadLevelScreen(levelConfigs[i].worldData, levelConfigs[i].worldSize,
                           levelConfigs[i].musicName, levelConfigs[i].musicId);
            return;
        }
    }
    shutdownDisplay();
}

// Display the start screen with menu options
void showStartScreen(void) {
    for (int i = 0; i < 16; i++) {
        BG_PAL[i] = (i < bg_menuPalLen / 2) ? bg_menuPal[i] : originalPal[i];
    }

    volatile u16* vramTile = VRAM_BG;
    for (int i = 0; i < bg_menuTilesLen / 2; i++) {
        vramTile[i] = bg_menuTiles[i];
    }

    fillBackgroundWithTileset(BG_MENU_TILE_1, BG_MENU_TILE_2, BG_MENU_TILE_3, BG_MENU_TILE_4);

    static const int newGameTiles[] = {LETTER_N, LETTER_E, LETTER_W, -1, LETTER_G, LETTER_A, LETTER_M, LETTER_E};
    static const int continueTiles[] = {LETTER_C, LETTER_O, LETTER_N, LETTER_T, LETTER_I, LETTER_N, LETTER_U, LETTER_E};
    static const int aboutTiles[] = {LETTER_A, LETTER_B, LETTER_O, LETTER_U, LETTER_T};
    static const int creditsTiles[] = {LETTER_C, LETTER_R, LETTER_E, LETTER_D, LETTER_I, LETTER_T, LETTER_S};

    MenuOption options[] = {
        {newGameTiles, 8, 0, 0, 0},
        {continueTiles, 8, 0, 0, 8},
        {aboutTiles, 5, 0, 0, 16},
        {creditsTiles, 7, 0, 0, 21}
    };
    const int optionCount = sizeof(options) / sizeof(options[0]);
    const int verticalSpacing = 10;

    int totalHeight = optionCount * SPRITE_SIZE + (optionCount - 1) * verticalSpacing;
    int startY = (SCREEN_HEIGHT - totalHeight) / 2;
    for (int i = 0; i < optionCount; i++) {
        renderMenuOption(&options[i], 0, startY + i * (SPRITE_SIZE + verticalSpacing));
    }

    int selectedOption = 0, pokeballIndex = 28, pokeballHeight = 24;
    int pokeballX = options[0].x - SPRITE_SIZE - 4, pokeballY = options[0].y + (SPRITE_SIZE - pokeballHeight) / 2;
    setSprite(pokeballIndex, pokeballX, pokeballY, POKEBALL);
    u16 prevKeys = 0;
    updateOAM();

menu_loop:
    while (1) {
        u16 keys = ~REG_KEYINPUT & 0x03FF;
        VBlankIntrWait();
        mmFrame();

        if ((keys & KEY_UP) && !(prevKeys & KEY_UP)) {
            selectedOption--;
            if (selectedOption < 0) selectedOption = optionCount - 1;
            mmEffectEx(&sfx_berry);
        }
        if ((keys & KEY_DOWN) && !(prevKeys & KEY_DOWN)) {
            selectedOption++;
            if (selectedOption >= optionCount) selectedOption = 0;
            mmEffectEx(&sfx_berry);
        }

        pokeballX = options[selectedOption].x - SPRITE_SIZE - 4;
        pokeballY = options[selectedOption].y + (SPRITE_SIZE - pokeballHeight) / 2;
        setSprite(pokeballIndex, pokeballX, pokeballY, POKEBALL);
        updateOAM();

        if ((keys & (KEY_START | KEY_A)) && !(prevKeys & (KEY_START | KEY_A))) {
            mmEffectEx(&sfx_start);
            delay(60);
            if (selectedOption == 0) {
                loadAndStartLevel(0);
                break;
            } else if (selectedOption == 1) {
                loadAndStartLevel(loadSavedLevel());
                break;
            } else if (selectedOption == 2) {
                showTextScreen("HELP KIDNAPPED VOLTORB\nRETURN TO POWER PLANT!\nCOLLECT ALL BERRIES AND\nSTONES. DODGE POKEBALLS\nOR YOU WILL BE\nCAPTURED AGAIN!");
                restoreMenu(options, optionCount, pokeballIndex, pokeballX, pokeballY);
                goto menu_loop;
            } else if (selectedOption == 3) {
                showTextScreen("GAME BY: ACEMAN\nMSX: ACEMAN\nGFX: MIXED SOURCES\nCODE: GROK AND CLAUDE\nLANGUAGE MODELS\nENDSCREEN: PHAERE");
                restoreMenu(options, optionCount, pokeballIndex, pokeballX, pokeballY);
                goto menu_loop;
            }
        }
        prevKeys = keys;
    }

    seedRand((u16)(REG_TM0CNT_L));
    for (int i = 0; i < 128; i++) hideSprite(i);
    updateOAM();
}