#include <Wire.h>
#include <U8g2lib.h>
#include <EEPROM.h>

// Display type selection - uncomment ONE of these:
#define SH1106
//#define SSD1306

#ifdef SH1106
    U8G2_SH1106_128X64_NONAME_F_HW_I2C u8g2(U8G2_R0, U8X8_PIN_NONE);
#elif defined(SSD1306)
    U8G2_SSD1306_128X64_NONAME_F_HW_I2C u8g2(U8G2_R0, U8X8_PIN_NONE);
#else
    #error "Please define either SH1106 or SSD1306"
#endif

// Display constants
#define DISPLAY_WIDTH 128
#define DISPLAY_HEIGHT 64

// Pin definitions
#define BTN_UP     12
#define BTN_DOWN   14
#define BTN_LEFT   27
#define BTN_RIGHT  26
#define BTN_MENU   25  // Was jump button
#define BTN_ACTION 16  // New button
#define BTN_PICK   17  // New button
#define BTN_DROP   18  // New button

// Map constants - removed semicolons
#define MAP_WIDTH 16
#define MAP_HEIGHT 16

// Game states
enum GameState {
    GAME_TITLE,
    GAME_PLAYING,
    GAME_MENU,
    GAME_INVENTORY
};

// Menu options
enum MenuOption {
    MENU_RESUME,
    MENU_SAVE,
    MENU_OPTIONS,
    MENU_QUIT,
    MENU_COUNT
};

// Item structure
struct Item {
    uint8_t id;
    char name[12];
    bool active;
};

// Player structure
struct Player {
    int16_t x;
    int16_t y;
    uint8_t health;
    uint8_t facing; // 0=up, 1=right, 2=down, 3=left
} player;

// Game variables
GameState currentState = GAME_TITLE;
MenuOption currentMenuOption = MENU_RESUME;
Item inventory[10];  // Max 10 items
uint8_t inventoryCount = 0;
uint8_t inventorySelection = 0;
bool buttonStates[8] = {false};
bool lastButtonStates[8] = {false};

// Tile map declaration
uint8_t gameMap[MAP_HEIGHT][MAP_WIDTH];

void setup() {
    // Initialize display
    u8g2.begin();
    u8g2.setFont(u8g2_font_6x10_tf);
    
    #ifdef SSD1306
        u8g2.setContrast(128);
    #endif

    // Initialize buttons
    pinMode(BTN_UP, INPUT_PULLUP);
    pinMode(BTN_DOWN, INPUT_PULLUP);
    pinMode(BTN_LEFT, INPUT_PULLUP);
    pinMode(BTN_RIGHT, INPUT_PULLUP);
    pinMode(BTN_MENU, INPUT_PULLUP);
    pinMode(BTN_ACTION, INPUT_PULLUP);
    pinMode(BTN_PICK, INPUT_PULLUP);
    pinMode(BTN_DROP, INPUT_PULLUP);

    // Initialize EEPROM for save functionality
    EEPROM.begin(512);

    initializeGame();
}

void initializeGame() {
    // Initialize player
    player.x = DISPLAY_WIDTH / 2;
    player.y = DISPLAY_HEIGHT / 2;
    player.health = 100;
    player.facing = 2; // facing down

    // Clear inventory
    inventoryCount = 0;
    for (int i = 0; i < 10; i++) {
        inventory[i].active = false;
    }

    // Initialize simple map (0=floor, 1=wall)
    for (int y = 0; y < MAP_HEIGHT; y++) {
        for (int x = 0; x < MAP_WIDTH; x++) {
            gameMap[y][x] = (x == 0 || x == MAP_WIDTH-1 || 
                            y == 0 || y == MAP_HEIGHT-1) ? 1 : 0;
        }
    }
}

bool isButtonPressed(uint8_t pin) {
    return !digitalRead(pin);
}

bool isButtonJustPressed(uint8_t pin, uint8_t index) {
    bool currentState = isButtonPressed(pin);
    bool justPressed = currentState && !lastButtonStates[index];
    lastButtonStates[index] = currentState;
    return justPressed;
}

void updatePlayer() {
    int16_t newX = player.x;
    int16_t newY = player.y;
    
    if (isButtonPressed(BTN_LEFT))  { newX -= 2; player.facing = 3; }
    if (isButtonPressed(BTN_RIGHT)) { newX += 2; player.facing = 1; }
    if (isButtonPressed(BTN_UP))    { newY -= 2; player.facing = 0; }
    if (isButtonPressed(BTN_DOWN))  { newY += 2; player.facing = 2; }

    // Check map boundaries and collisions
    int mapX = newX / 8;
    int mapY = newY / 8;
    
    if (mapX >= 0 && mapX < MAP_WIDTH && 
        mapY >= 0 && mapY < MAP_HEIGHT && 
        gameMap[mapY][mapX] == 0) {
        player.x = newX;
        player.y = newY;
    }
}

void drawPlayer() {
    // Draw player as triangle pointing in facing direction
    const uint8_t size = 3;
    switch (player.facing) {
        case 0: // up
            u8g2.drawTriangle(
                player.x, player.y - size,
                player.x - size, player.y + size,
                player.x + size, player.y + size
            );
            break;
        case 1: // right
            u8g2.drawTriangle(
                player.x + size, player.y,
                player.x - size, player.y - size,
                player.x - size, player.y + size
            );
            break;
        case 2: // down
            u8g2.drawTriangle(
                player.x, player.y + size,
                player.x - size, player.y - size,
                player.x + size, player.y - size
            );
            break;
        case 3: // left
            u8g2.drawTriangle(
                player.x - size, player.y,
                player.x + size, player.y - size,
                player.x + size, player.y + size
            );
            break;
    }
}

void drawMap() {
    // Calculate visible portion of map
    int startX = (player.x / 8) - 8;
    int startY = (player.y / 8) - 4;
    int endX = startX + 16;
    int endY = startY + 8;

    // Clamp to map boundaries
    startX = max(0, min((int)MAP_WIDTH, startX));
    startY = max(0, min((int)MAP_HEIGHT, startY));
    endX = max(0, min((int)MAP_WIDTH, endX));
    endY = max(0, min((int)MAP_HEIGHT, endY));

    // Draw visible tiles
    for (int y = startY; y < endY; y++) {
        for (int x = startX; x < endX; x++) {
            if (gameMap[y][x] == 1) {
                u8g2.drawBox(x*8 - (player.x/8)*8 + 64, 
                            y*8 - (player.y/8)*8 + 32, 
                            8, 8);
            }
        }
    }
}

void drawMenu() {
    const char* menuItems[] = {
        "Resume",
        "Save Game",
        "Options",
        "Quit to Title"
    };

    u8g2.drawBox(20, 10, 88, 44);
    u8g2.setDrawColor(0);
    
    for (int i = 0; i < MENU_COUNT; i++) {
        if (i == currentMenuOption) {
            u8g2.drawBox(22, 12 + i*10, 84, 10);
            u8g2.setDrawColor(0);
            u8g2.drawStr(24, 20 + i*10, menuItems[i]);
            u8g2.setDrawColor(1);
        } else {
            u8g2.drawStr(24, 20 + i*10, menuItems[i]);
        }
    }
    u8g2.setDrawColor(1);
}

void drawInventory() {
    u8g2.drawBox(20, 10, 88, 44);
    u8g2.setDrawColor(0);
    
    u8g2.drawStr(24, 20, "Inventory:");
    
    for (int i = 0; i < inventoryCount; i++) {
        if (i == inventorySelection) {
            u8g2.drawBox(22, 22 + i*10, 84, 10);
            u8g2.setDrawColor(0);
            u8g2.drawStr(24, 30 + i*10, inventory[i].name);
            u8g2.setDrawColor(1);
        } else {
            u8g2.drawStr(24, 30 + i*10, inventory[i].name);
        }
    }
    u8g2.setDrawColor(1);
}

void drawTitleScreen() {
    u8g2.setFont(u8g2_font_8x13B_tf);
    u8g2.drawStr(30, 30, "Mini RPG");
    u8g2.setFont(u8g2_font_6x10_tf);
    u8g2.drawStr(20, 50, "Press ACTION to start");
}

void handleMenu() {
    if (isButtonJustPressed(BTN_UP, 0)) {
        currentMenuOption = (MenuOption)((currentMenuOption - 1 + MENU_COUNT) % MENU_COUNT);
    }
    if (isButtonJustPressed(BTN_DOWN, 1)) {
        currentMenuOption = (MenuOption)((currentMenuOption + 1) % MENU_COUNT);
    }
    if (isButtonJustPressed(BTN_ACTION, 5)) {
        switch (currentMenuOption) {
            case MENU_RESUME:
                currentState = GAME_PLAYING;
                break;
            case MENU_SAVE:
                saveGame();
                break;
            case MENU_OPTIONS:
                // Implement options menu
                break;
            case MENU_QUIT:
                currentState = GAME_TITLE;
                initializeGame();
                break;
        }
    }
}

void saveGame() {
    // Simple save implementation
    EEPROM.put(0, player);
    EEPROM.put(sizeof(player), inventory);
    EEPROM.commit();
}

void loadGame() {
    EEPROM.get(0, player);
    EEPROM.get(sizeof(player), inventory);
}

void loop() {
    u8g2.clearBuffer();

    switch (currentState) {
        case GAME_TITLE:
            drawTitleScreen();
            if (isButtonJustPressed(BTN_ACTION, 5)) {
                currentState = GAME_PLAYING;
            }
            break;

        case GAME_PLAYING:
            updatePlayer();
            drawMap();
            drawPlayer();
            
            if (isButtonJustPressed(BTN_MENU, 4)) {
                currentState = GAME_MENU;
            }
            if (isButtonJustPressed(BTN_PICK, 6)) {
                currentState = GAME_INVENTORY;
            }
            break;

        case GAME_MENU:
            drawMap();
            drawPlayer();
            drawMenu();
            handleMenu();
            break;

        case GAME_INVENTORY:
            drawMap();
            drawPlayer();
            drawInventory();
            if (isButtonJustPressed(BTN_MENU, 4)) {
                currentState = GAME_PLAYING;
            }
            break;
    }

    u8g2.sendBuffer();
    delay(16); // ~60 FPS
}