#include "UIManager.h"
#include "config.h"
#include "SharedState.h"

// --- HARDWARE CONFIGURATION ---
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
#define OLED_RESET    -1 
#define SCREEN_ADDRESS 0x3C

UIManager ui;

UIManager::UIManager() 
    : display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET)
{
}

void UIManager::begin() {
    // Setup Display
    Wire.begin();
    if(!display.begin(SSD1306_SWITCHCAPVCC, SCREEN_ADDRESS)) {
        Serial.println(F("SSD1306 allocation failed"));
        for(;;);
    }
    display.clearDisplay();
    display.display();

    // Setup NeoPixel Matrix
    ledMatrix.begin();
}

void UIManager::showMessage(const char* msg) {
    display.clearDisplay();
    display.setCursor(10, 25);
    display.setTextColor(SSD1306_WHITE);
    display.setTextSize(2);
    display.print(msg);
    display.display();
    delay(500);
    display.setTextSize(1);
}
 
void UIManager::draw(UIState currentState, int menuSelection, 
          int midiChannel, int tempo, MelodyStrategy* currentStrategy, 
          int queuedTheme, int currentThemeIndex,
          int numScaleNotes, const int* scaleNotes, int melodySeed, int currentTrackNumSteps,
          bool mutationEnabled, bool songModeEnabled, 
          const Step sequence[][NUM_STEPS], int playbackStep, bool isPlaying,
          int randomizeTrack, const bool* trackMute, const int* trackIntensities) {
 
    display.clearDisplay();
    display.setTextSize(1);
    display.setTextColor(SSD1306_WHITE);
    display.setCursor(0, 0);

    switch(currentState) {
        case UI_MENU_MAIN:
            drawMenu(menuSelection, currentState, midiChannel, tempo, currentStrategy->getName(), queuedTheme, currentThemeIndex, numScaleNotes, scaleNotes, melodySeed, currentTrackNumSteps, mutationEnabled, songModeEnabled, isPlaying, randomizeTrack, trackMute, trackIntensities);
            break;
        case UI_SETUP_CHANNEL_EDIT:
            drawEditScreen("SET MIDI CHANNEL", "CH: ", midiChannel);
            break;
        case UI_EDIT_TEMPO:
            drawEditScreen("SET TEMPO", "BPM: ", tempo);
            break;
        case UI_EDIT_STEPS:
            drawEditScreen("SET STEPS", "LEN: ", currentTrackNumSteps);
            break;
        case UI_EDIT_FLAVOUR:
            drawEditScreen("SET FLAVOUR", "", currentStrategy->getName());
            break;
        case UI_EDIT_ROOT:
            {
                const char* noteNames[] = {"C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B"};
                drawEditScreen("SET KEY", "", noteNames[globalRoot % 12]);
            }
            break;
        case UI_EDIT_PROG_ORDER:
            drawEditScreen("PROG ORDER", "", (progressionOrder == PROG_ORDER_SEQUENCE) ? "Sequence" : "Random");
            break;
        case UI_EDIT_PROG_TEMPLATE:
            drawEditScreen("APPLY TEMPLATE", "", progressions[templateSelectionIndex].name);
            break;
        case UI_EDIT_PROG_LENGTH:
            drawEditScreen("PROG LENGTH", "Steps: ", currentProgression.length);
            break;
        case UI_EDIT_PROG_STEP_ROOT:
            drawEditScreen("STEP MODULATION", "Semi: ", currentProgression.steps[stepEditIndex/3].rootOffset);
            break;
        case UI_EDIT_PROG_STEP_TYPE:
            {
                const char* typeNames[] = {"Chrom", "Major", "Minor", "H.Min", "P.Maj", "P.Min", "ChdMaj", "ChdMin", "ChdDim", "Chd7"};
                int type = currentProgression.steps[stepEditIndex/3].scaleType;
                if (type >= 0 && type < 10)
                    drawEditScreen("STEP CHORD", "", typeNames[type]);
                else
                    drawEditScreen("STEP CHORD", "", "?");
            }
            break;
        case UI_EDIT_PROG_STEP_REPS:
            drawEditScreen("STEP REPEATS", "Reps: ", currentProgression.steps[stepEditIndex/3].repeats);
            break;
        case UI_EDIT_INTENSITY:
            drawNumberEditor("SET INTENSITY", trackIntensities[randomizeTrack], 1, 10);
            break;
        case UI_RANDOMIZE_TRACK_EDIT:
            drawEditScreen("SET TRACK", "TRK: ", randomizeTrack + 1);
            break;
    }
    display.display();
}

void UIManager::drawNumberEditor(const char* title, int value, int minVal, int maxVal) {
    display.println(title);
    display.drawLine(0, 8, 128, 8, SSD1306_WHITE);
    
    // Display value
    display.setCursor(20, 20);
    display.setTextSize(2);
    display.print(value);

    // Graphical bar
    int barWidth = 100;
    int barX = (SCREEN_WIDTH - barWidth) / 2;
    int barY = 40;
    int barHeight = 10;
    float percentage = (float)(value - minVal) / (maxVal - minVal);
    int fillWidth = (int)(percentage * barWidth);

    display.drawRect(barX, barY, barWidth, barHeight, SSD1306_WHITE);
    display.fillRect(barX, barY, fillWidth, barHeight, SSD1306_WHITE);

    display.setTextSize(1);
    display.setCursor(0, 54);
    display.println(F(" (Press to confirm)"));
}

void UIManager::drawEditScreen(const char* title, const char* label, const char* valueStr) {
    display.println(title);
    display.drawLine(0, 8, 128, 8, SSD1306_WHITE);
    display.setCursor(20, 25);
    display.setTextSize(2);
    if (label && *label) display.print(label);
    display.print(valueStr);
    display.setTextSize(1);
    display.setCursor(0, 50);
    display.println(F(" (Press to confirm)"));
}

void UIManager::drawEditScreen(const char* title, const char* label, int value) {
    char buf[16];
    itoa(value, buf, 10);
    drawEditScreen(title, label, buf);
}

void UIManager::drawMenu(int selection, UIState currentState, int midiChannel, int tempo, const char* flavourName, 
              int queuedTheme, int currentThemeIndex, int numScaleNotes,
              const int* scaleNotes, int melodySeed, int currentTrackNumSteps, bool mutationEnabled,
              bool songModeEnabled, bool isPlaying, int randomizeTrack, const bool* trackMute, const int* trackIntensities) {

    // Calculate visual cursor position and scroll offset
    int visualCursor = 0;
    for(int i=0; i<selection; i++) {
        if(isItemVisible(i)) visualCursor++;
    }
    
    const int MAX_LINES = 7; // No title, so we have more space
    int startVisualIndex = 0;
    if (visualCursor >= MAX_LINES) {
        startVisualIndex = visualCursor - (MAX_LINES - 1);
    }
    
    int currentVisualIndex = 0;
    int y = 0;

    for (int i = 0; i < menuItemsCount; i++) {
        if (!isItemVisible(i)) continue;
        
        if (currentVisualIndex >= startVisualIndex) {
            if (y > 55) break;

        MenuItemID id = menuItems[i].id;

        if (i == selection) {
            display.fillRect(0, y, 128, 9, SSD1306_WHITE);
            display.setTextColor(SSD1306_BLACK, SSD1306_WHITE);
        } else {
            display.setTextColor(SSD1306_WHITE);
        }
        
        int x = 2 + (menuItems[i].indentLevel * 6);
        display.setCursor(x, y + 1);
        
        if (menuItems[i].isGroup) {
            display.print(menuItems[i].expanded ? F("v ") : F("> "));
        }
        
        if (id >= MENU_ID_PROG_STEP_START && id < MENU_ID_PROG_STEP_END) {
            int stepIdx = (id - MENU_ID_PROG_STEP_START) / 3;
            if (stepIdx == visualProgressionStep) {
                display.setCursor(x - 6 * 2, y + 1);
                display.print(F("| "));
            } else if (stepIdx == queuedProgressionStep) {
                display.setCursor(x - 6 * 2, y + 1);
                display.print(F(", "));
            }
        }
        display.print(menuItems[i].label);

        if (id == MENU_ID_CHANNEL) {
            display.print(F(": ")); display.print(midiChannel);
        }
        
        // Dynamic values
            if (id == MENU_ID_PLAYBACK) { display.print(F(": ")); display.print(isPlaying ? F("ON") : F("OFF")); }
            else if (id == MENU_ID_MELODY) {
                display.print(F(": ")); display.print(melodySeed);
            } else if (id == MENU_ID_TEMPO) { display.print(F(": ")); display.print(tempo); }
            else if (id == MENU_ID_ROOT) { 
                const char* noteNames[] = {"C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B"};
                display.print(F(": ")); display.print(noteNames[globalRoot % 12]); 
            }
            else if (id == MENU_ID_STEPS) { display.print(F(": ")); display.print(currentTrackNumSteps); }
            else if (id == MENU_ID_SONG_MODE) { display.print(F(": ")); display.print(songModeEnabled ? F("ON") : F("OFF")); }
            else if (id == MENU_ID_TRACK_SELECT) { display.print(F(": ")); display.print(randomizeTrack + 1); }
            else if (id == MENU_ID_MUTE) { display.print(F(": ")); display.print(trackMute[randomizeTrack] ? F("YES") : F("NO")); }
            else if (id == MENU_ID_FLAVOUR) { display.print(F(": ")); display.print(flavourName); }
            else if (id == MENU_ID_PROG_ORDER) { display.print(F(": ")); display.print((progressionOrder == PROG_ORDER_SEQUENCE) ? F("Seq") : F("Rnd")); }
            else if (id == MENU_ID_PROG_LENGTH) { display.print(F(": ")); display.print(currentProgression.length); }
            else if (id >= MENU_ID_PROG_STEP_START && id < MENU_ID_PROG_STEP_END) {
                // Only show steps that are within the current length
                int stepIdx = (id - MENU_ID_PROG_STEP_START) / 3;
                if (stepIdx >= currentProgression.length) {
                    continue; // Skip rendering this item
                }
                int idx = id - MENU_ID_PROG_STEP_START;
                int step = idx / 3;
                int type = idx % 3; // 0=Root, 1=Type, 2=Reps
                display.print(F(": "));
                if (type == 1) {
                    const char* typeNames[] = {"Chrom", "Maj", "Min", "HMin", "PMaj", "PMin", "CMaj", "CMin", "CDim", "C7"};
                    int t = currentProgression.steps[step].scaleType;
                    if (t >= 0 && t < 10) display.print(typeNames[t]);
                } else if (type == 2) {
                    display.print(currentProgression.steps[step].repeats);
                } else {
                    int root = currentProgression.steps[step].rootOffset;
                    if (root > 0) display.print(F("+"));
                    display.print(root);
                }
            }
            else if (id == MENU_ID_INTENSITY) {
                display.print(F(": "));
                display.print(trackIntensities[randomizeTrack]);
                int val = trackIntensities[randomizeTrack];
                int barX = display.getCursorX() + 3;
                int barY = y + 2;
                int maxW = 20;
                int h = 5;
                uint16_t color = (i == selection) ? SSD1306_BLACK : SSD1306_WHITE;
                display.drawRect(barX, barY, maxW + 2, h, color);
                display.fillRect(barX + 1, barY + 1, (val * maxW) / 10, h - 2, color);
            }
            else if (id == MENU_ID_MUTATION) { display.print(F(": ")); display.print(mutationEnabled ? F("ON") : F("OFF")); }
            else if (id == MENU_ID_PROTECTED_MODE) { display.print(F(": ")); display.print(protectedMode ? F("ON") : F("OFF")); }
            
            if (id >= MENU_ID_THEME_1 && id <= MENU_ID_THEME_7) {
                int themeIdx = id - MENU_ID_THEME_1 + 1;
                if (queuedTheme == themeIdx) display.print(F(" [NEXT]"));
                if (currentThemeIndex == themeIdx) display.print(F(" *"));
            }

        y += 9;
        }
        
        currentVisualIndex++;
    }
}

void UIManager::updateLeds(const Step sequence[][NUM_STEPS], int playbackStep, bool isPlaying, 
                UIState currentState, bool songModeEnabled, 
                int songRepeatsRemaining, bool sequenceChangeScheduled, PlayMode playMode,
                int selectedTrack, const int* numSteps, int numScaleNotes, const int* scaleNotes, const bool* trackMute) {
    ledMatrix.update(sequence, playbackStep, isPlaying, currentState, songModeEnabled, songRepeatsRemaining, sequenceChangeScheduled, playMode, selectedTrack, numSteps, trackMute);
}