#include "game.h"
#include "graphics.h"
#include "sound.h"
#include "startscreen.h"
#include "gfx/bg_lab.h"
#include "gfx/endscreen.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_woods.h"
#include "gfx/world_powerp.h"
#include "gfx/world_special1.h"
#include "gfx/world_special2.h"
#include <string.h>
#include <maxmod.h>
#include <stdbool.h>

// Counter for holding L+R buttons to trigger level reset
int triggerHoldCounter = 0;

// Switch mapping for toggle interactions
#define MAX_SWITCHES 50
typedef struct {
    int switchX;
    int switchY;
    int targetX;
    int targetY;
    int switchTargetTile;
    int originalTile;
    int isTargetBlocking;
} SwitchMap;
SwitchMap switchMap[MAX_SWITCHES];
int switchMapCount = 0;

// Update Jigglypuff animation for walls and objects
void updateJigglyAnimation(int* jigglyFrame, int* jigglyDelayCounter, Wall* walls, int wallCount, LevelObject* objects, int objectCount, int* spriteIndex) {
    const int JIGGLY_DELAY = 30;
    (*jigglyDelayCounter)++;
    if (*jigglyDelayCounter >= JIGGLY_DELAY) {
        *jigglyFrame = !*jigglyFrame;
        *jigglyDelayCounter = 0;
    }

    for (int i = 0; i < wallCount && *spriteIndex < 128; i++) {
        if (walls[i].tile == JIGGLY_1 || walls[i].tile == JIGGLY_2) {
            setSprite(*spriteIndex, walls[i].x, walls[i].y, *jigglyFrame ? JIGGLY_2 : JIGGLY_1);
        } else {
            setSprite(*spriteIndex, walls[i].x, walls[i].y, walls[i].tile);
        }
        (*spriteIndex)++;
    }

    for (int i = 0; i < objectCount && *spriteIndex < 128; i++) {
        if (!objects[i].collected) {
            if (objects[i].tile == JIGGLY_1 || objects[i].tile == JIGGLY_2) {
                setSprite(*spriteIndex, objects[i].x, objects[i].y, *jigglyFrame ? JIGGLY_2 : JIGGLY_1);
            } else {
                setSprite(*spriteIndex, objects[i].x, objects[i].y, objects[i].tile);
            }
            (*spriteIndex)++;
        }
    }
}

// Handle switch toggle logic to modify level objects or walls
void handleSwitchToggle(int switchIndex, LevelObject* objects, int* objectCount, Wall* walls, int* wallCount) {
    if (switchIndex == -1) {
        switchMapCount = 0;
        for (int i = 0; i < *objectCount && switchMapCount < MAX_SWITCHES; i++) {
            if (objects[i].tile == SWITCH_ON || objects[i].tile == SWITCH_OFF) {
                switchMap[switchMapCount].switchX = objects[i].x;
                switchMap[switchMapCount].switchY = objects[i].y;
                switchMap[switchMapCount].targetX = (objects[i].switchTargetCol - 1) * SPRITE_SIZE;
                switchMap[switchMapCount].targetY = (objects[i].switchTargetRow - 1) * SPRITE_SIZE;
                switchMap[switchMapCount].switchTargetTile = objects[i].switchTargetTile;
                switchMap[switchMapCount].originalTile = objects[i].originalTile;
                switchMap[switchMapCount].isTargetBlocking = 0;
                for (int k = 0; k < blockingTileCount; k++) {
                    if (objects[i].switchTargetTile == blockingTiles[k].tile) {
                        switchMap[switchMapCount].isTargetBlocking = 1;
                        break;
                    }
                }
                switchMapCount++;
            }
        }
        return;
    }

    SwitchMap* sm = &switchMap[switchIndex];
    int targetX = sm->targetX;
    int targetY = sm->targetY;

    int switchObjectIndex = -1;
    for (int i = 0; i < *objectCount; i++) {
        if (objects[i].x == sm->switchX && objects[i].y == sm->switchY && (objects[i].tile == SWITCH_ON || objects[i].tile == SWITCH_OFF)) {
            switchObjectIndex = i;
            break;
        }
    }

    if (switchObjectIndex == -1) {
        return;
    }

    if (objects[switchObjectIndex].tile == SWITCH_OFF) {
        objects[switchObjectIndex].tile = SWITCH_ON;
        mmEffectEx(&sfx_switch);

        int targetObjectIndex = -1;
        for (int j = 0; j < *objectCount; j++) {
            if (objects[j].x == targetX && objects[j].y == targetY) {
                targetObjectIndex = j;
                break;
            }
        }

        int targetWallIndex = -1;
        for (int j = 0; j < *wallCount; j++) {
            if (walls[j].x == targetX && walls[j].y == targetY) {
                targetWallIndex = j;
                break;
            }
        }

        if (sm->switchTargetTile == 0) {
            if (targetWallIndex != -1) {
                for (int k = targetWallIndex; k < *wallCount - 1; k++) {
                    walls[k] = walls[k + 1];
                }
                (*wallCount)--;
            } else if (targetObjectIndex != -1) {
                for (int k = targetObjectIndex; k < *objectCount - 1; k++) {
                    objects[k] = objects[k + 1];
                }
                (*objectCount)--;
            }
        } else {
            if (targetWallIndex != -1) {
                if (!sm->isTargetBlocking && targetObjectIndex == -1 && *objectCount < 50) {
                    objects[*objectCount].x = targetX;
                    objects[*objectCount].y = targetY;
                    objects[*objectCount].tile = sm->switchTargetTile;
                    objects[*objectCount].collected = 0;
                    objects[*objectCount].soundPlayed = 0;
                    objects[*objectCount].switchTargetRow = 0;
                    objects[*objectCount].switchTargetCol = 0;
                    objects[*objectCount].switchTargetTile = 0;
                    objects[*objectCount].originalTile = 0;
                    objects[*objectCount].isTriggered = 0;
                    (*objectCount)++;
                    for (int k = targetWallIndex; k < *wallCount - 1; k++) {
                        walls[k] = walls[k + 1];
                    }
                    (*wallCount)--;
                } else if (sm->isTargetBlocking) {
                    walls[targetWallIndex].tile = sm->switchTargetTile;
                }
            } else if (targetObjectIndex != -1) {
                objects[targetObjectIndex].tile = sm->switchTargetTile;
            } else if (sm->isTargetBlocking && *wallCount < 100) {
                walls[*wallCount].x = targetX;
                walls[*wallCount].y = targetY;
                walls[*wallCount].tile = sm->switchTargetTile;
                walls[*wallCount].specialTargetLevel = -1;
                (*wallCount)++;
            } else if (!sm->isTargetBlocking && *objectCount < 50) {
                objects[*objectCount].x = targetX;
                objects[*objectCount].y = targetY;
                objects[*objectCount].tile = sm->switchTargetTile;
                objects[*objectCount].collected = 0;
                objects[*objectCount].soundPlayed = 0;
                objects[*objectCount].switchTargetRow = 0;
                objects[*objectCount].switchTargetCol = 0;
                objects[*objectCount].switchTargetTile = 0;
                objects[*objectCount].originalTile = 0;
                objects[*objectCount].isTriggered = 0;
                (*objectCount)++;
            }
        }
    } else {
        objects[switchObjectIndex].tile = SWITCH_OFF;
        mmEffectEx(&sfx_switch);

        int targetObjectIndex = -1;
        for (int j = 0; j < *objectCount; j++) {
            if (objects[j].x == targetX && objects[j].y == targetY) {
                targetObjectIndex = j;
                break;
            }
        }

        int targetWallIndex = -1;
        for (int j = 0; j < *wallCount; j++) {
            if (walls[j].x == targetX && walls[j].y == targetY) {
                targetWallIndex = j;
                break;
            }
        }

        if (sm->originalTile == 0) {
            if (targetWallIndex != -1) {
                for (int k = targetWallIndex; k < *wallCount - 1; k++) {
                    walls[k] = walls[k + 1];
                }
                (*wallCount)--;
            }
            if (targetObjectIndex != -1) {
                for (int k = targetObjectIndex; k < *objectCount - 1; k++) {
                    objects[k] = objects[k + 1];
                }
                (*objectCount)--;
            }
        } else {
            int isOriginalBlocking = 0;
            for (int k = 0; k < blockingTileCount; k++) {
                if (sm->originalTile == blockingTiles[k].tile) {
                    isOriginalBlocking = 1;
                    break;
                }
            }
            if (isOriginalBlocking && *wallCount < 100) {
                if (targetObjectIndex != -1) {
                    walls[*wallCount].x = targetX;
                    walls[*wallCount].y = targetY;
                    walls[*wallCount].tile = sm->originalTile;
                    walls[*wallCount].specialTargetLevel = -1;
                    (*wallCount)++;
                    for (int k = targetObjectIndex; k < *objectCount - 1; k++) {
                        objects[k] = objects[k + 1];
                    }
                    (*objectCount)--;
                } else if (targetWallIndex != -1) {
                    walls[targetWallIndex].tile = sm->originalTile;
                } else {
                    walls[*wallCount].x = targetX;
                    walls[*wallCount].y = targetY;
                    walls[*wallCount].tile = sm->originalTile;
                    walls[*wallCount].specialTargetLevel = -1;
                    (*wallCount)++;
                }
            } else if (*objectCount < 50) {
                if (targetWallIndex != -1) {
                    objects[*objectCount].x = targetX;
                    objects[*objectCount].y = targetY;
                    objects[*objectCount].tile = sm->originalTile;
                    objects[*objectCount].collected = 0;
                    objects[*objectCount].soundPlayed = 0;
                    objects[*objectCount].switchTargetRow = 0;
                    objects[*objectCount].switchTargetCol = 0;
                    objects[*objectCount].switchTargetTile = 0;
                    objects[*objectCount].originalTile = 0;
                    objects[*objectCount].isTriggered = 0;
                    (*objectCount)++;
                    for (int k = targetWallIndex; k < *wallCount - 1; k++) {
                        walls[k] = walls[k + 1];
                    }
                    (*wallCount)--;
                } else if (targetObjectIndex != -1) {
                    objects[targetObjectIndex].tile = sm->originalTile;
                } else {
                    objects[*objectCount].x = targetX;
                    objects[*objectCount].y = targetY;
                    objects[*objectCount].tile = sm->originalTile;
                    objects[*objectCount].collected = 0;
                    objects[*objectCount].soundPlayed = 0;
                    objects[*objectCount].switchTargetRow = 0;
                    objects[*objectCount].switchTargetCol = 0;
                    objects[*objectCount].switchTargetTile = 0;
                    objects[*objectCount].originalTile = 0;
                    objects[*objectCount].isTriggered = 0;
                    (*objectCount)++;
                }
            }
        }
    }
}

// Main game loop and initialization
int main(void) {
    REG_DISPCNT = MODE0 | BG0_ENABLE | OBJ_ENABLE | OBJ_1D_MAP;
    initSound();
    initGraphics();

    static int musicInitialized = 0;
    static WorldType currentMusicWorldType = -1;

game_start:
    showStartScreen();

    Wall walls[100];
    LevelObject objects[50];
    int wallCount = 0, objectCount = 0, playerX = 80, playerY = 80, coinCount = 0;
    int guiX = 212, guiY = 147;
    WorldType worldType;
    int x, y, dx, dy, baseTile, playerVisible, isElectro, jigglyFrame, jigglyDelayCounter;
    int spriteIndex;
    u16 prevKeys;
    mm_sfxhand pokeballHandle;
    int spawnTiles[] = {PLAYER_SPAWN_1, PLAYER_SPAWN_2, PLAYER_SPAWN_3, PLAYER_SPAWN_4, PLAYER_SPAWN_5, PLAYER_SPAWN_6};

    loadLevel(currentLevel, walls, &wallCount, objects, &objectCount, &playerX, &playerY, &coinCount, &worldType);
    resetGraphicsForLevel(worldType);
    resetSprites(walls, wallCount, objects, objectCount, guiX, guiY, coinCount);
    saveLevel(currentLevel);

    handleSwitchToggle(-1, objects, &objectCount, walls, &wallCount);

    if (!musicInitialized || currentMusicWorldType != worldType) {
        initSound();
        playMusic("", worldType);
        musicInitialized = 1;
        currentMusicWorldType = worldType;
    }

    x = playerX; y = playerY; dx = 0; dy = 0; baseTile = PLAYER_DOWN;
    prevKeys = 0; playerVisible = 1; pokeballHandle = 0; isElectro = 0;
    jigglyFrame = 0; jigglyDelayCounter = 0;
    spriteIndex = 3;

    for (int i = 0; i < objectCount; i++) {
        objects[i].isTriggered = 0;
    }

    for (int i = 0; i < wallCount && spriteIndex < 128; i++) {
        setSprite(spriteIndex++, walls[i].x, walls[i].y, walls[i].tile);
    }
    for (int i = 0; i < objectCount && spriteIndex < 128; i++) {
        if (!objects[i].collected) {
            setSprite(spriteIndex++, objects[i].x, objects[i].y, objects[i].tile);
        }
    }
    updateCounter(coinCount, guiX, guiY);
    updateOAM();

    VBlankIntrWait();
    REG_DISPCNT = MODE0 | BG0_ENABLE | OBJ_ENABLE | OBJ_1D_MAP;
    mmEffectEx(&sfx_spawn);
    playSpawnAnimation(x, y, spawnTiles);
    mmEffectEx(&sfx_player);

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

        if ((keys & KEY_L) && (keys & KEY_R)) {
            triggerHoldCounter++;
        } else {
            triggerHoldCounter = 0;
        }

        if (dx == 0 && dy == 0 && playerVisible) {
            if ((keys & KEY_LEFT) && !(prevKeys & KEY_LEFT)) { 
                dx = -SPRITE_SPEED; dy = 0; 
                baseTile = isElectro ? PLAYER_LEFT_ELECTRO : PLAYER_LEFT; 
            }
            else if ((keys & KEY_RIGHT) && !(prevKeys & KEY_RIGHT)) { 
                dx = SPRITE_SPEED; dy = 0; 
                baseTile = isElectro ? PLAYER_RIGHT_ELECTRO : PLAYER_RIGHT; 
            }
            else if ((keys & KEY_UP) && !(prevKeys & KEY_UP)) { 
                dx = 0; dy = -SPRITE_SPEED; 
                baseTile = isElectro ? PLAYER_UP_ELECTRO : PLAYER_UP; 
            }
            else if ((keys & KEY_DOWN) && !(prevKeys & KEY_DOWN)) { 
                dx = 0; dy = SPRITE_SPEED; 
                baseTile = isElectro ? PLAYER_DOWN_ELECTRO : PLAYER_DOWN; 
            }
        }

        int newX = x + dx, newY = y + dy;
        int levelChanged = 0;

        if ((dx != 0 || dy != 0) && playerVisible) {
            if (checkSpecialTileCollision(newX, newY, walls, wallCount, coinCount, &levelChanged)) {
                dx = 0; dy = 0;
                mmEffectEx(&sfx_exit);
                for (int j = 5; j >= 0; j--) {
                    setSprite(0, x, y, spawnTiles[j]);
                    updateOAM();
                    delay(5);
                }
                playerVisible = 0;
                hideSprite(0);
                updateOAM();

                if (levelChanged == 9) {
                    mmEffectEx(&sfx_special1);
                    delay(30);
                } else if (levelChanged == 54) {
                    mmEffectEx(&sfx_special2);
                    delay(30);
                }

                for (int j = 30; j >= 0; j--) {
                    mmSetModuleVolume((j * 512) / 30);
                    VBlankIntrWait();
                    mmFrame();
                }

                loadLevel(levelChanged, walls, &wallCount, objects, &objectCount, &playerX, &playerY, &coinCount, &worldType);
                resetGraphicsForLevel(worldType);
                resetSprites(walls, wallCount, objects, objectCount, guiX, guiY, coinCount);
                handleSwitchToggle(-1, objects, &objectCount, walls, &wallCount);

                mmSetModuleVolume(512);
                if (currentMusicWorldType != worldType) {
                    playMusic("", worldType);
                    currentMusicWorldType = worldType;
                }

                x = playerX; y = playerY; dx = 0; dy = 0; baseTile = PLAYER_DOWN;
                playerVisible = 1; isElectro = 0; jigglyFrame = 0; jigglyDelayCounter = 0;
                spriteIndex = 3;

                for (int j = 0; j < objectCount; j++) {
                    objects[j].isTriggered = 0;
                }

                for (int j = 0; j < wallCount && spriteIndex < 128; j++) {
                    setSprite(spriteIndex++, walls[j].x, walls[j].y, walls[j].tile);
                }
                for (int j = 0; j < objectCount && spriteIndex < 128; j++) {
                    if (!objects[j].collected) {
                        setSprite(spriteIndex++, objects[j].x, objects[j].y, objects[j].tile);
                    }
                }
                updateCounter(coinCount, guiX, guiY);
                updateOAM();
                mmEffectEx(&sfx_spawn);
                playSpawnAnimation(x, y, spawnTiles);
                mmEffectEx(&sfx_player);
                continue;
            }
        }

        if (!checkCollisionWithWrap(newX, newY, walls, wallCount, objects, objectCount, coinCount, currentLevel) && playerVisible) {
            x = (newX + SCREEN_WIDTH) % SCREEN_WIDTH;
            y = (newY + SCREEN_HEIGHT) % SCREEN_HEIGHT;

            bool warpedThisFrame = false;
            bool rampedThisFrame = false;
            bool stoppedThisFrame = false;

            for (int i = 0; i < objectCount; i++) {
                if (objects[i].tile == STOP && (x != objects[i].x || y != objects[i].y)) {
                    objects[i].soundPlayed = 0;
                }
                if ((objects[i].tile == SWITCH_OFF || objects[i].tile == SWITCH_ON) && (x != objects[i].x || y != objects[i].y)) {
                    objects[i].isTriggered = 0;
                }

                if (x == objects[i].x && y == objects[i].y) {
                    if (!objects[i].collected && (objects[i].tile == COIN || (objects[i].tile >= BERRY_1 && objects[i].tile <= BERRY_4))) {
                        objects[i].collected = 1;
                        coinCount--;
                        mmEffectEx(objects[i].tile == COIN ? &sfx_coin : &sfx_berry);
                    }

                    if (objects[i].tile == STOP && !stoppedThisFrame) {
                        if (!objects[i].soundPlayed) {
                            mmEffectEx(&sfx_stop);
                            objects[i].soundPlayed = 1;
                        }
                        dx = 0;
                        dy = 0;
                        stoppedThisFrame = true;
                    }

                    if (objects[i].tile == WARP && !warpedThisFrame) {
                        for (int j = 0; j < objectCount; j++) {
                            if (objects[j].tile == WARP && i != j) {
                                mmEffectEx(&sfx_warp);
                                x = objects[j].x;
                                y = objects[j].y;
                                warpedThisFrame = true;
                                i = -1;
                                break;
                            }
                        }
                        if (warpedThisFrame) continue;
                    }

                    if (!rampedThisFrame && (objects[i].tile == RAMP_LEFT || objects[i].tile == RAMP_RIGHT)) {
                        mmEffectEx(&sfx_ramp);
                        if (objects[i].tile == RAMP_LEFT) {
                            if (dx > 0) { dx = 0; dy = -SPRITE_SPEED; baseTile = isElectro ? PLAYER_UP_ELECTRO : PLAYER_UP; }
                            else if (dx < 0) { dx = 0; dy = SPRITE_SPEED; baseTile = isElectro ? PLAYER_DOWN_ELECTRO : PLAYER_DOWN; }
                            else if (dy > 0) { dx = -SPRITE_SPEED; dy = 0; baseTile = isElectro ? PLAYER_LEFT_ELECTRO : PLAYER_LEFT; }
                            else if (dy < 0) { dx = SPRITE_SPEED; dy = 0; baseTile = isElectro ? PLAYER_RIGHT_ELECTRO : PLAYER_RIGHT; }
                        } else {
                            if (dx > 0) { dx = 0; dy = SPRITE_SPEED; baseTile = isElectro ? PLAYER_DOWN_ELECTRO : PLAYER_DOWN; }
                            else if (dx < 0) { dx = 0; dy = -SPRITE_SPEED; baseTile = isElectro ? PLAYER_UP_ELECTRO : PLAYER_UP; }
                            else if (dy > 0) { dx = SPRITE_SPEED; dy = 0; baseTile = isElectro ? PLAYER_RIGHT_ELECTRO : PLAYER_RIGHT; }
                            else if (dy < 0) { dx = -SPRITE_SPEED; dy = 0; baseTile = isElectro ? PLAYER_LEFT_ELECTRO : PLAYER_LEFT; }
                        }
                        rampedThisFrame = true;
                    }

                    if (objects[i].tile == SWITCH_OFF || objects[i].tile == SWITCH_ON) {
                        if (!objects[i].isTriggered) {
                            for (int k = 0; k < switchMapCount; k++) {
                                if (switchMap[k].switchX == objects[i].x && switchMap[k].switchY == objects[i].y) {
                                    handleSwitchToggle(k, objects, &objectCount, walls, &wallCount);
                                    updateSpatialGrid(walls, wallCount, objects, objectCount);
                                    break;
                                }
                            }
                            objects[i].isTriggered = 1;
                        }
                    }

                    if (objects[i].tile == EXIT && coinCount == 0) {
                        dx = 0; dy = 0;
                        mmEffectEx(&sfx_exit);
                        for (int j = 5; j >= 0; j--) {
                            setSprite(0, x, y, spawnTiles[j]);
                            updateOAM();
                            delay(5);
                        }
                        playerVisible = 0;
                        hideSprite(0);
                        updateOAM();

                        if (currentLevel != 10 && currentLevel != 20 && currentLevel != 30 && currentLevel != 40 && currentLevel != 50 && currentLevel != 60 && currentLevel != 70 && (currentLevel + 1) < maxLevels) {
                            for (int j = 30; j >= 0; j--) {
                                mmSetModuleVolume((j * 512) / 30);
                                VBlankIntrWait();
                                mmFrame();
                            }
                        }

                        delay(5);
                        currentLevel++;

                        if (currentLevel == 9 || currentLevel == 54) {
                            currentLevel++;
                        }

                        if (currentLevel >= maxLevels) {
                            showEndScreen();
                            mmSetModuleVolume(0);
                            initGraphics();                            
                            musicInitialized = 0;
                            currentMusicWorldType = -1;
                            saveLevel(0);
                            goto game_start;
                        }

                        if (currentLevel == 11 || currentLevel == 21 || currentLevel == 31 || currentLevel == 41 || currentLevel == 51 || currentLevel == 62) {
                            memset((void*)0x06000000, 0, 0xA000);
                            for (int k = 0; k < 128; k++) shadowOAM[k].attr0 = 0x0200;
                            updateOAM();
                            memset((void*)0x06010000, 0, 0x4000);
                            REG_DISPCNT = 0x0403;
                            volatile u16* vram = (volatile u16*)0x06000000;

                            if(currentLevel == 11) memcpy((void*)vram, worldice, WORLDICE_SIZE);
                            if(currentLevel == 21) memcpy((void*)vram, worldcave, WORLDCAVE_SIZE);
                            if(currentLevel == 31) memcpy((void*)vram, worldwoods, WORLDWOODS_SIZE);
                            if(currentLevel == 41) memcpy((void*)vram, worldcity, WORLDCITY_SIZE);
                            if(currentLevel == 51) memcpy((void*)vram, worldmansion, WORLDMANSION_SIZE);
                            if(currentLevel == 62) memcpy((void*)vram, worldpowerp, WORLDPOWERP_SIZE);

                            delay(180);

                            for (int j = 120; j >= 0; j--) {
                                mmSetModuleVolume((j * 512) / 120);
                                VBlankIntrWait();
                                mmFrame();
                            }

                            while (1) {
                                u16 current_keys = ~REG_KEYINPUT & 0x03FF;
                                if (current_keys & (KEY_A | KEY_START)) {
                                    mmSetModuleVolume(512);
                                    VBlankIntrWait();
                                    for (int j = 0; j < 16; j++) BG_PAL[j] = 0x0000;
                                    for (int j = 0; j < 256; j++) PAL_OBJ[j] = 0x0000;
                                    memset((void*)0x06000000, 0, 0xA000);
                                    REG_DISPCNT = 0;
                                    break;
                                 }
                                VBlankIntrWait();
                                mmFrame();
                            }
                        }

                        loadLevel(currentLevel, walls, &wallCount, objects, &objectCount, &playerX, &playerY, &coinCount, &worldType);
                        saveLevel(currentLevel);
                        resetGraphicsForLevel(worldType);
                        resetSprites(walls, wallCount, objects, objectCount, guiX, guiY, coinCount);
                        handleSwitchToggle(-1, objects, &objectCount, walls, &wallCount);

                        mmSetModuleVolume(512);
                        if (currentMusicWorldType != worldType) {
                            playMusic("", worldType);
                            currentMusicWorldType = worldType;
                        }

                        x = playerX; y = playerY; dx = 0; dy = 0; baseTile = PLAYER_DOWN; 
                        playerVisible = 1; isElectro = 0; jigglyFrame = 0; jigglyDelayCounter = 0;
                        spriteIndex = 3;

                        for (int j = 0; j < objectCount; j++) objects[j].isTriggered = 0;

                        for (int j = 0; j < wallCount && spriteIndex < 128; j++) setSprite(spriteIndex++, walls[j].x, walls[j].y, walls[j].tile);
                        for (int j = 0; j < objectCount && spriteIndex < 128; j++) {
                            if (!objects[j].collected) setSprite(spriteIndex++, objects[j].x, objects[j].y, objects[j].tile);
                        }
                        updateCounter(coinCount, guiX, guiY);
                        updateOAM();
                        mmEffectEx(&sfx_spawn);
                        playSpawnAnimation(x, y, spawnTiles);
                        mmEffectEx(&sfx_player);
                        continue;
                    }
                }
            }
        } else if (playerVisible) {
            dx = 0; dy = 0;
            for (int i = 0; i < wallCount; i++) {
                if (checkCollision(newX, newY, walls[i].x, walls[i].y, 0)) {
                    for (int j = 0; j < blockingTileCount; j++) {
                        if (walls[i].tile == blockingTiles[j].tile) {
                            if (strstr(blockingTiles[j].name, "_SPECIAL_") || strstr(blockingTiles[j].name, "BOX") || strstr(blockingTiles[j].name, "QBLOCK")) {
                                mmEffectEx(&sfx_pipe); 
                                break;
                            } else if (strstr(blockingTiles[j].name, "WALL")) {
                                mmEffectEx(&sfx_wall); 
                                break;
                            } else if (strstr(blockingTiles[j].name, "JIGGLY") && !isElectro) {
                                mmEffectEx(&sfx_jiggly);
                                mmEffectEx(&sfx_kiss);

                                int jigglySpriteIndex = -1;
                                for (int k = 3; k < spriteIndex; k++) {
                                    if (shadowOAM[k].attr0 != 0x0200 &&
                                        (shadowOAM[k].attr1 & 0x01FF) == walls[i].x && 
                                        (shadowOAM[k].attr0 & 0x00FF) == walls[i].y && 
                                        ((shadowOAM[k].attr2 & 0x03FF) == JIGGLY_1 || (shadowOAM[k].attr2 & 0x03FF) == JIGGLY_2)) {
                                        jigglySpriteIndex = k;
                                        break;
                                    }
                                }

                                if (jigglySpriteIndex != -1) {
                                    setSprite(jigglySpriteIndex, walls[i].x, walls[i].y, JIGGLY_KISS);
                                    updateOAM();
                                    delay(60);
                                    setSprite(jigglySpriteIndex, walls[i].x, walls[i].y, jigglyFrame ? JIGGLY_2 : JIGGLY_1);
                                    updateOAM();
                                }

                                mmEffectEx(&sfx_happy);
                                mmEffectEx(&sfx_electro);
                                int electroTile = (baseTile == PLAYER_LEFT) ? PLAYER_LEFT_ELECTRO :
                                                 (baseTile == PLAYER_RIGHT) ? PLAYER_RIGHT_ELECTRO :
                                                 (baseTile == PLAYER_UP) ? PLAYER_UP_ELECTRO : PLAYER_DOWN_ELECTRO;
                                int normalTile = baseTile;
                                for (int k = 0; k < 5; k++) {
                                    setSprite(0, x, y, electroTile);
                                    updateOAM();
                                    delay(5);
                                    setSprite(0, x, y, normalTile);
                                    updateOAM();
                                    delay(5);
                                }
                                isElectro = 1;
                                baseTile = electroTile;

                                if (jigglySpriteIndex != -1) {
                                    setSprite(jigglySpriteIndex, walls[i].x, walls[i].y, jigglyFrame ? JIGGLY_2 : JIGGLY_1);
                                    updateOAM();
                                }
                                break;
                            }
                        }
                    }
                }
            }

            for (int i = 0; i < objectCount; i++) {
                if (!objects[i].collected && checkCollision(newX, newY, objects[i].x, objects[i].y, 0) && 
                    (objects[i].tile == JIGGLY_1 || objects[i].tile == JIGGLY_2) && !isElectro) {
                    mmEffectEx(&sfx_jiggly);
                    mmEffectEx(&sfx_kiss);

                    int jigglySpriteIndex = -1;
                    for (int k = 3; k < spriteIndex; k++) {
                        if (shadowOAM[k].attr0 != 0x0200 &&
                            (shadowOAM[k].attr1 & 0x01FF) == objects[i].x && 
                            (shadowOAM[k].attr0 & 0x00FF) == objects[i].y && 
                            ((shadowOAM[k].attr2 & 0x03FF) == JIGGLY_1 || (shadowOAM[k].attr2 & 0x03FF) == JIGGLY_2)) {
                            jigglySpriteIndex = k;
                            break;
                        }
                    }

                    if (jigglySpriteIndex != -1) {
                        setSprite(jigglySpriteIndex, objects[i].x, objects[i].y, JIGGLY_KISS);
                        updateOAM();
                        delay(60);
                        setSprite(jigglySpriteIndex, objects[i].x, objects[i].y, jigglyFrame ? JIGGLY_2 : JIGGLY_1);
                        updateOAM();
                    }

                    mmEffectEx(&sfx_happy);
                    mmEffectEx(&sfx_electro);
                    int electroTile = (baseTile == PLAYER_LEFT) ? PLAYER_LEFT_ELECTRO :
                                     (baseTile == PLAYER_RIGHT) ? PLAYER_RIGHT_ELECTRO :
                                     (baseTile == PLAYER_UP) ? PLAYER_UP_ELECTRO : PLAYER_DOWN_ELECTRO;
                    int normalTile = baseTile;
                    for (int k = 0; k < 5; k++) {
                        setSprite(0, x, y, electroTile);
                        updateOAM();
                        delay(5);
                        setSprite(0, x, y, normalTile);
                        updateOAM();
                        delay(5);
                    }
                    isElectro = 1;
                    baseTile = electroTile;

                    if (jigglySpriteIndex != -1) {
                        setSprite(jigglySpriteIndex, objects[i].x, objects[i].y, jigglyFrame ? JIGGLY_2 : JIGGLY_1);
                        updateOAM();
                    }
                    break;
                }
            }
        }
        
        bool resetLevel = false;
        for (int i = 0; i < objectCount; i++) {
            if (objects[i].tile == POKEBALL && !objects[i].collected && checkCollision(x, y, objects[i].x, objects[i].y, 0)) {
                if (isElectro) {
                    int pokeballSpriteIndex = -1;
                    for (int j = 3; j < spriteIndex; j++) {
                        if ((shadowOAM[j].attr1 & 0x01FF) == objects[i].x && 
                            (shadowOAM[j].attr0 & 0x00FF) == objects[i].y && 
                            (shadowOAM[j].attr2 & 0x03FF) == POKEBALL) {
                            pokeballSpriteIndex = j;
                            break;
                        }
                    }
                    if (pokeballSpriteIndex != -1) {
                        int explodeTiles[] = {POKEBALL_EXPLODE_1, POKEBALL_EXPLODE_2, POKEBALL_EXPLODE_3, POKEBALL_EXPLODE_4};
                        mmEffectEx(&sfx_explode);
                        for (int j = 0; j < 4; j++) {
                            setSprite(pokeballSpriteIndex, objects[i].x, objects[i].y, explodeTiles[j]);
                            updateOAM();
                            delay(5);
                        }
                        hideSprite(pokeballSpriteIndex);
                        updateOAM();
                        objects[i].collected = 1;
                        isElectro = 0;
                        baseTile = (baseTile == PLAYER_LEFT_ELECTRO) ? PLAYER_LEFT :
                                   (baseTile == PLAYER_RIGHT_ELECTRO) ? PLAYER_RIGHT :
                                   (baseTile == PLAYER_UP_ELECTRO) ? PLAYER_UP : PLAYER_DOWN;
                    }
                } else {
                    resetLevel = true;
                    break;
                }
            }
        }

        if (triggerHoldCounter >= 60) {
            resetLevel = true;
            triggerHoldCounter = 0;
        }

        if (resetLevel) {
            pokeballHandle = mmEffectEx(&sfx_pokeball);
            for (int j = 5; j >= 0; j--) {
                setSprite(0, x, y, spawnTiles[j]);
                updateOAM();
                delay(5);
            }
            playerVisible = 0;
            hideSprite(0);
            updateOAM();
            delay(30);
            if (pokeballHandle) { mmEffectCancel(pokeballHandle); pokeballHandle = 0; }
            restoreLevel(currentLevel, walls, &wallCount, objects, &objectCount, &playerX, &playerY, &coinCount, &worldType);
            handleSwitchToggle(-1, objects, &objectCount, walls, &wallCount);
            resetGraphicsForLevel(worldType);
            resetSprites(walls, wallCount, objects, objectCount, guiX, guiY, coinCount);
            x = playerX; y = playerY; dx = 0; dy = 0; baseTile = PLAYER_DOWN;
            playerVisible = 1; isElectro = 0; jigglyFrame = 0; jigglyDelayCounter = 0;
            spriteIndex = 3;
            for (int j = 0; j < objectCount; j++) objects[j].isTriggered = 0;
            for (int j = 0; j < wallCount && spriteIndex < 128; j++) setSprite(spriteIndex++, walls[j].x, walls[j].y, walls[j].tile);
            for (int j = 0; j < objectCount && spriteIndex < 128; j++) {
                if (!objects[j].collected) setSprite(spriteIndex++, objects[j].x, objects[j].y, objects[j].tile);
            }
            updateCounter(coinCount, guiX, guiY);
            updateOAM();
            mmEffectEx(&sfx_spawn);
            playSpawnAnimation(x, y, spawnTiles);
            mmEffectEx(&sfx_player);
        }

        for (int i = 0; i < 128; i++) hideSprite(i);
        if (playerVisible) setSprite(0, x, y, baseTile);
        updateCounter(coinCount, guiX, guiY);
        spriteIndex = 3;
        updateJigglyAnimation(&jigglyFrame, &jigglyDelayCounter, walls, wallCount, objects, objectCount, &spriteIndex);
        if (playerVisible) {
            if (x + SPRITE_SIZE > SCREEN_WIDTH && spriteIndex < 128) setSprite(spriteIndex++, x - SCREEN_WIDTH, y, baseTile);
            if (y + SPRITE_SIZE > SCREEN_HEIGHT && spriteIndex < 128) setSprite(spriteIndex++, x, y - SCREEN_HEIGHT, baseTile);
            if (x + SPRITE_SIZE > SCREEN_WIDTH && y + SPRITE_SIZE > SCREEN_HEIGHT && spriteIndex < 128) {
                setSprite(spriteIndex++, x - SCREEN_WIDTH, y - SCREEN_HEIGHT, baseTile);
            }
        }
        updateOAM();

        prevKeys = keys;
    }
    return 0;
}