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

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

UIManager ui;

UIManager::UIManager() 
    : display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET),
      pixels(NUM_PIXELS, PIN_NEOPIXEL, NEO_GRB + NEO_KHZ800)
{
}

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
    pixels.begin();
    pixels.setBrightness(40);
    pixels.clear();
    pixels.show();
}

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 navSelection, EditMode editMode, 
          int midiChannel, int tempo, MelodyStrategy* currentStrategy, 
          int queuedTheme, int currentThemeIndex, 
          int numScaleNotes, const int* scaleNotes, int melodySeed, 
          bool mutationEnabled, bool songModeEnabled, 
          const Step sequence[][NUM_STEPS], int scrollOffset, int playbackStep, bool isPlaying,
          const char* mainMenu[], int mainMenuCount,
          const char* randomizeMenu[], int randomizeMenuCount,
          const char* setupMenu[], int setupMenuCount, int theme1Index, 
          PlayMode playMode, int currentTrack, int randomizeTrack, const bool* trackMute) {

    display.clearDisplay();
    display.setTextSize(1);
    display.setTextColor(SSD1306_WHITE);
    display.setCursor(0, 0);

    switch(currentState) {
        case UI_TRACKER:
            drawTracker(navSelection, editMode, midiChannel, sequence, scrollOffset, playbackStep, isPlaying, playMode, currentTrack);
            break;
        case UI_MENU_MAIN:
            drawMenu("MAIN MENU", mainMenu, mainMenuCount, menuSelection, currentState, midiChannel, tempo, currentStrategy->getName(), queuedTheme, currentThemeIndex, numScaleNotes, scaleNotes, melodySeed, mutationEnabled, songModeEnabled, theme1Index, playMode, randomizeTrack, trackMute);
            break;
        case UI_MENU_RANDOMIZE:
            drawMenu("RANDOMIZE", randomizeMenu, randomizeMenuCount, menuSelection, currentState, midiChannel, tempo, currentStrategy->getName(), queuedTheme, currentThemeIndex, numScaleNotes, scaleNotes, melodySeed, mutationEnabled, songModeEnabled, theme1Index, playMode, randomizeTrack, trackMute);
            break;
        case UI_MENU_SETUP:
            drawMenu("SETUP", setupMenu, setupMenuCount, menuSelection, currentState, midiChannel, tempo, currentStrategy->getName(), queuedTheme, currentThemeIndex, numScaleNotes, scaleNotes, melodySeed, mutationEnabled, songModeEnabled, theme1Index, playMode, randomizeTrack, trackMute);
            break;
        case UI_SETUP_CHANNEL_EDIT:
            display.println(F("SET MIDI CHANNEL"));
            display.drawLine(0, 8, 128, 8, SSD1306_WHITE);
            display.setCursor(20, 25);
            display.setTextSize(2);
            display.print(F("CH: "));
            if (midiChannel < 10) display.print(F(" "));
            display.print(midiChannel);
            display.setTextSize(1);
            display.setCursor(0, 50);
            display.println(F(" (Press to confirm)"));
            break;
        case UI_EDIT_TEMPO:
            display.println(F("SET TEMPO"));
            display.drawLine(0, 8, 128, 8, SSD1306_WHITE);
            display.setCursor(20, 25);
            display.setTextSize(2);
            display.print(F("BPM: "));
            display.print(tempo);
            display.setTextSize(1);
            display.setCursor(0, 50);
            display.println(F(" (Press to confirm)"));
            break;
        case UI_EDIT_FLAVOUR:
            display.println(F("SET FLAVOUR"));
            display.drawLine(0, 8, 128, 8, SSD1306_WHITE);
            display.setCursor(20, 25);
            display.setTextSize(2);
            display.print(currentStrategy->getName());
            display.setTextSize(1);
            display.setCursor(0, 50);
            display.println(F(" (Press to confirm)"));
            break;
        case UI_SETUP_PLAYMODE_EDIT:
            display.println(F("SET PLAY MODE"));
            display.drawLine(0, 8, 128, 8, SSD1306_WHITE);
            display.setCursor(20, 25);
            display.setTextSize(2);
            display.print(playMode == MODE_MONO ? "Mono" : "Poly");
            display.setTextSize(1);
            display.setCursor(0, 50);
            display.println(F(" (Press to confirm)"));
            break;
        case UI_RANDOMIZE_TRACK_EDIT:
            display.println(F("SET TRACK"));
            display.drawLine(0, 8, 128, 8, SSD1306_WHITE);
            display.setCursor(20, 25);
            display.setTextSize(2);
            display.print(F("TRK: "));
            display.print(randomizeTrack + 1);
            display.setTextSize(1);
            display.setCursor(0, 50);
            display.println(F(" (Press to confirm)"));
            break;
    }
    display.display();
}

void UIManager::drawMenu(const char* title, const char* items[], int count, int selection, 
              UIState currentState, int midiChannel, int tempo, const char* flavourName, 
              int queuedTheme, int currentThemeIndex, int numScaleNotes, 
              const int* scaleNotes, int melodySeed, bool mutationEnabled, 
              bool songModeEnabled, int theme1Index, PlayMode playMode, int randomizeTrack, const bool* trackMute) {
    display.println(title);
    display.drawLine(0, 8, 128, 8, SSD1306_WHITE);

    int start = 0;
    if (selection >= 5) start = selection - 4;

    int y = 10;
    for (int i = start; i < count; i++) {
        if (y > 55) break;

        if (i == selection) {
            display.fillRect(0, y, 128, 8, SSD1306_WHITE);
            display.setTextColor(SSD1306_BLACK, SSD1306_WHITE);
        } else {
            display.setTextColor(SSD1306_WHITE);
        }
        display.setCursor(2, y);
        display.print(items[i]);

        if (currentState == UI_MENU_SETUP && i == 1) {
            display.print(F(": ")); display.print(playMode == MODE_MONO ? "Mono" : "Poly");
        }
        if (currentState == UI_MENU_SETUP && i == 2) {
            display.print(F(": ")); display.print(midiChannel);
        }
        if (currentState == UI_MENU_RANDOMIZE && playMode == MODE_POLY && i == 1) {
            display.print(F(": ")); display.print(randomizeTrack + 1);
        }
        if (currentState == UI_MENU_RANDOMIZE && playMode == MODE_POLY && i == 2) {
            display.print(F(": ")); display.print(trackMute[randomizeTrack] ? F("YES") : F("NO"));
        }
        if (currentState == UI_MENU_RANDOMIZE && i >= theme1Index && queuedTheme == (i - theme1Index + 1)) {
            display.print(F(" [NEXT]"));
        }
        if (currentState == UI_MENU_RANDOMIZE && i >= theme1Index && currentThemeIndex == (i - theme1Index + 1)) {
            display.print(F(" *"));
        }
        if (currentState == UI_MENU_RANDOMIZE) {
            int track_offset = (playMode == MODE_POLY) ? 2 : 0;
            if (i == 1 + track_offset) { // Melody
                display.print(F(": ")); display.print(melodySeed);
            } else if (i == 2 + track_offset) { // Flavour
                display.print(F(": ")); display.print(flavourName);
            } else if (i == 3 + track_offset) { // Scale
                display.print(F(": "));
                if (numScaleNotes > 0) {
                    const char* noteNames[] = {"C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B"};
                    for (int j = 0; j < min(numScaleNotes, 6); j++) {
                        display.print(noteNames[scaleNotes[j]]);
                        if (j < min(numScaleNotes, 6) - 1) display.print(F(" "));
                    }
                }
            } else if (i == 4 + track_offset) { display.print(F(": ")); display.print(tempo); }
            else if (i == 5 + track_offset) { display.print(F(": ")); display.print(mutationEnabled ? F("ON") : F("OFF")); }
            else if (i == 6 + track_offset) { display.print(F(": ")); display.print(songModeEnabled ? F("ON") : F("OFF")); }
        }
        y += 9;
    }
}

void UIManager::drawTracker(int navSelection, EditMode editMode, int midiChannel, 
                 const Step sequence[][NUM_STEPS], int scrollOffset, int playbackStep, bool isPlaying,
                 PlayMode playMode, int currentTrack) {
    display.print(F("SEQ "));
    if (playMode == MODE_POLY) {
        switch(editMode) {
            case NAV_STEP: display.print(F("[STP]")); break;
            case NAV_TRACK: display.print(F("[TRK]")); break;
            case EDIT_NOTE: display.print(F("[EDT]")); break;
        }
    } else {
        if (navSelection > 0 && editMode == EDIT_NOTE) display.print(F("[EDT]"));
        else display.print(F("[NAV]"));
    }

    display.print(F(" CH:")); display.print(midiChannel);
    display.println();
    display.drawLine(0, 8, 128, 8, SSD1306_WHITE);

    int y = 10;
    for (int i = 0; i < 6; i++) {
        int itemIndex = i + scrollOffset;
        if (itemIndex > NUM_STEPS) break;

        if (itemIndex == navSelection && editMode != NAV_TRACK) {
            display.fillRect(0, y, 128, 8, SSD1306_WHITE);
            display.setTextColor(SSD1306_BLACK, SSD1306_WHITE);
        } else {
            display.setTextColor(SSD1306_WHITE);
        }
        display.setCursor(2, y);

        if (itemIndex == 0) {
            display.print(F(">> MENU"));
        } else {
            int stepIndex = itemIndex - 1;
            bool isPlayback = isPlaying && (stepIndex == playbackStep);
            if (isPlayback && editMode != NAV_TRACK) {
                if (itemIndex == navSelection) display.setTextColor(SSD1306_WHITE, SSD1306_BLACK); // Cursor on playback row
                else display.setTextColor(SSD1306_BLACK, SSD1306_WHITE);
            }
            if (stepIndex < 10) display.print(F("0"));
            display.print(stepIndex);
            if (isPlayback && editMode != NAV_TRACK) {
                if (itemIndex == navSelection) display.setTextColor(SSD1306_BLACK, SSD1306_WHITE);
                else display.setTextColor(SSD1306_WHITE);
            }
            display.print(F(" | "));

            int tracksToShow = (playMode == MODE_POLY) ? NUM_TRACKS : 1;
            for(int t=0; t<tracksToShow; t++) {
                if (playMode == MODE_POLY && t == currentTrack && (editMode == NAV_TRACK || (editMode == EDIT_NOTE && itemIndex == navSelection))) {
                    display.fillRect(display.getCursorX(), y, 28, 8, SSD1306_WHITE);
                    display.setTextColor(SSD1306_BLACK, SSD1306_WHITE);
                } else if (itemIndex == navSelection && editMode != NAV_TRACK) {
                    display.setTextColor(SSD1306_BLACK, SSD1306_WHITE);
                } else {
                    display.setTextColor(SSD1306_WHITE);
                }

                int n = sequence[t][stepIndex].note;
                if (n == -1) {
                    display.print(F(" ---"));
                } else {
                    const char* noteNames[] = {"C-", "C#", "D-", "D#", "E-", "F-", "F#", "G-", "G#", "A-", "A#", "B-"};
                    display.print(noteNames[n % 12]);
                    display.print(n / 12 - 1);
                }
                display.print(" ");
            }
        }
        y += 9;
    }
}

uint32_t UIManager::getNoteColor(int note, bool dim) {
    if (note == -1) return 0;
    uint16_t hue = 30000 + (note % 12) * 3628;
    return Adafruit_NeoPixel::ColorHSV(hue, 255, dim ? 10 : 50);
}

int UIManager::getPixelIndex(int x, int y) {
    return y * 8 + x;
}

void UIManager::updateLeds(const Step sequence[][NUM_STEPS], int navSelection, int playbackStep, bool isPlaying, 
                UIState currentState, EditMode editMode, bool songModeEnabled, 
                int songRepeatsRemaining, bool sequenceChangeScheduled, PlayMode playMode, int currentTrack,
                int numScaleNotes, const int* scaleNotes, const bool* trackMute) {
    pixels.clear();
    
    const uint32_t COLOR_PLAYHEAD = pixels.Color(0, 255, 0);
    const uint32_t COLOR_PLAYHEAD_DIM = pixels.Color(0, 32, 0);
    const uint32_t COLOR_MUTED_PLAYHEAD = pixels.Color(0, 0, 255);
    const uint32_t COLOR_CURSOR = pixels.Color(255, 255, 255);
    const uint32_t COLOR_CURSOR_DIM = pixels.Color(32, 0, 0);

    if(playMode == MODE_POLY) {
        for(int t=0; t<NUM_TRACKS; t++) {
            for(int s=0; s<NUM_STEPS; s++) {
                int col = t * 2 + (s / 8);
                int row = s % 8;
                
                uint32_t color = 0;
                int note = sequence[t][s].note;
                
                if (note != -1) {
                    color = getNoteColor(note, !sequence[t][s].accent);
                }
                
                if (isPlaying && s == playbackStep) {
                    if (trackMute[t]) {
                        color = COLOR_MUTED_PLAYHEAD;
                    } else {
                        color = (note != -1) ? COLOR_PLAYHEAD : COLOR_PLAYHEAD_DIM;
                    }
                } else if (!isPlaying && currentState == UI_TRACKER && (navSelection - 1) == s && currentTrack == t) {
                    color = (note != -1) ? COLOR_CURSOR : COLOR_CURSOR_DIM;
                }
                
                pixels.setPixelColor(getPixelIndex(col, row), color);
            }
        }
    } else {
        // --- Mono Mode (original) ---
        for (int s = 0; s < NUM_STEPS; s++) {
            int x = s % 8;
            int yBase = (s / 8) * 4;
            uint32_t color = 0, dimColor = 0;
            if (sequence[0][s].note != -1) {
                color = getNoteColor(sequence[0][s].note, sequence[0][s].tie);
                dimColor = getNoteColor(sequence[0][s].note, true);
            }
            uint32_t c[4] = {0};
            if (sequence[0][s].note != -1) {
                int octave = sequence[0][s].note / 12;
                if (octave > 4) { c[0] = color; if (sequence[0][s].accent) c[1] = dimColor; }
                else if (octave < 4) { c[2] = color; if (sequence[0][s].accent) c[1] = dimColor; }
                else { c[1] = color; if (sequence[0][s].accent) { c[0] = dimColor; c[2] = dimColor; } }
            }
            int stepNavIndex = navSelection - 1;
            uint32_t cursorColor = 0;
            if (isPlaying) {
                 cursorColor = pixels.Color(0, 50, 0);
                 if (songModeEnabled && s >= 8) {
                     int repeats = min(songRepeatsRemaining, 8);
                     if (x >= (8 - repeats)) cursorColor = (songRepeatsRemaining == 1 && x == 7 && (millis()/250)%2) ? pixels.Color(255, 200, 0) : pixels.Color(100, 220, 40);
                 }
            } else if (currentState == UI_TRACKER) {
                 cursorColor = (editMode == EDIT_NOTE) ? pixels.Color(50, 0, 0) : pixels.Color(40, 40, 40);
            }
            
            bool isCursorHere = (isPlaying && s == playbackStep) || (!isPlaying && currentState == UI_TRACKER && s == stepNavIndex);
            if (cursorColor != 0) {
                if (isCursorHere) c[3] = cursorColor;
                else {
                    uint8_t r = (uint8_t)(cursorColor >> 16), g = (uint8_t)(cursorColor >> 8), b = (uint8_t)cursorColor;
                    c[3] = pixels.Color(r/5, g/5, b/5);
                }
            }
            for(int i=0; i<4; i++) pixels.setPixelColor(getPixelIndex(x, yBase + i), c[i]);
        }
    }
    if (sequenceChangeScheduled && (millis() / 125) % 2) pixels.setPixelColor(NUM_PIXELS - 1, pixels.Color(127, 50, 0));
    pixels.show();
}