#include <mutex>
#include "UIThread.h"
#include "SharedState.h"
#include <Arduino.h>
#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>
#include "synth_engine.h"
#include <EEPROM.h>

extern SynthEngine* globalSynth;

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

// I2C Pins (GP4/GP5)
#define PIN_SDA 4
#define PIN_SCL 5

// Encoder Pins
#define PIN_ENC_CLK 12
#define PIN_ENC_DT  13
#define PIN_ENC_SW 14

Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);

volatile int8_t encoderDelta = 0;
static uint8_t prevNextCode = 0;
static uint16_t store = 0;

// Button state
static int lastButtonReading = HIGH;
static int currentButtonState = HIGH;
static unsigned long lastDebounceTime = 0;
static bool buttonClick = false;

void handleInput();
void drawUI();

// --- ENCODER INTERRUPT ---
// Robust Rotary Encoder reading
void readEncoder() {
  static int8_t rot_enc_table[] = {0,1,1,0,1,0,0,1,1,0,0,1,0,1,1,0};

  prevNextCode <<= 2;
  if (digitalRead(PIN_ENC_DT)) prevNextCode |= 0x02;
  if (digitalRead(PIN_ENC_CLK)) prevNextCode |= 0x01;
  prevNextCode &= 0x0f;

  // If valid state
  if (rot_enc_table[prevNextCode]) {
    store <<= 4;
    store |= prevNextCode;
    if ((store & 0xff) == 0x2b) encoderDelta--;
    if ((store & 0xff) == 0x17) encoderDelta++;
  }
}

void saveGridToEEPROM() {
    if (!globalSynth) return;
    uint8_t buf[SynthEngine::SERIALIZED_GRID_SIZE];
    globalSynth->exportGrid(buf);
    
    EEPROM.write(0, 'N');
    EEPROM.write(1, 'S');
    for (size_t i = 0; i < sizeof(buf); i++) {
        EEPROM.write(2 + i, buf[i]);
    }
    EEPROM.commit();
}

void loadGridFromEEPROM() {
    if (!globalSynth) return;
    if (EEPROM.read(0) == 'N' && EEPROM.read(1) == 'S') {
        uint8_t buf[SynthEngine::SERIALIZED_GRID_SIZE];
        for (size_t i = 0; i < sizeof(buf); i++) {
            buf[i] = EEPROM.read(2 + i);
        }
        globalSynth->importGrid(buf);
    } else {
        globalSynth->loadPreset(1); // Default to preset 1
    }
}

void setupUI() {
  Wire.setSDA(PIN_SDA);
  Wire.setSCL(PIN_SCL);
  Wire.begin();

    pinMode(PIN_ENC_CLK, INPUT_PULLUP);
    pinMode(PIN_ENC_DT, INPUT_PULLUP);
    pinMode(PIN_ENC_SW, INPUT_PULLUP);
    attachInterrupt(digitalPinToInterrupt(PIN_ENC_CLK), readEncoder, CHANGE);
    attachInterrupt(digitalPinToInterrupt(PIN_ENC_DT), readEncoder, CHANGE);

  if(!display.begin(SSD1306_SWITCHCAPVCC, SCREEN_ADDRESS)) {
    Serial.println(F("SSD1306 allocation failed"));
    for(;;);
  }

  display.clearDisplay();
  display.display();
  
  // Initialize EEPROM
  EEPROM.begin(512);

  // Check for safety clear (Button held on startup)
  if (digitalRead(PIN_ENC_SW) == LOW) {
      display.setCursor(0, 0);
      display.setTextColor(SSD1306_WHITE);
      display.println(F("CLEARING DATA..."));
      display.display();
      EEPROM.write(0, 0); // Invalidate magic
      EEPROM.commit();
      delay(1000);
  }

  loadGridFromEEPROM();
}

void handleInput() {
    // Handle Encoder Rotation
    int rotation = 0;
    noInterrupts();
    rotation = encoderDelta;
    encoderDelta = 0;
    interrupts();

    if (rotation != 0) {
        switch (currentState) {
            case UI_MENU:
                menuSelection += rotation;
                while (menuSelection < 0) menuSelection += NUM_MENU_ITEMS;
                menuSelection %= NUM_MENU_ITEMS;
                break;
            case UI_EDIT_SCALE_TYPE:
                currentScaleIndex += rotation;
                while (currentScaleIndex < 0) currentScaleIndex += NUM_SCALES;
                currentScaleIndex %= NUM_SCALES;
                break;
            case UI_EDIT_SCALE_KEY:
                currentKeyIndex += rotation;
                while (currentKeyIndex < 0) currentKeyIndex += NUM_KEYS;
                currentKeyIndex %= NUM_KEYS;
                break;
            case UI_EDIT_WAVETABLE:
                currentWavetableIndex += rotation;
                while (currentWavetableIndex < 0) currentWavetableIndex += NUM_WAVETABLES;
                currentWavetableIndex %= NUM_WAVETABLES;
                break;
        }
    }

    // Handle Button Click
    int reading = digitalRead(PIN_ENC_SW);
    if (reading != lastButtonReading) {
        lastDebounceTime = millis();
    }

    if ((millis() - lastDebounceTime) > 50) {
        if (reading != currentButtonState) {
            currentButtonState = reading;
            if (currentButtonState == LOW) {
                buttonClick = true;
            }
        }
    }
    lastButtonReading = reading;

    if (buttonClick) {
        buttonClick = false;
        if (currentState == UI_MENU) {
            currentState = MENU_ITEMS[menuSelection].editState;
        } else {
            currentState = UI_MENU;
        }
    }
}

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

    if (currentState == UI_MENU) {
        for (int i = 0; i < NUM_MENU_ITEMS; i++) {
            if (i == menuSelection) {
                display.fillRect(0, i * 10, SCREEN_WIDTH, 10, SSD1306_WHITE);
                display.setTextColor(SSD1306_BLACK, SSD1306_WHITE);
            } else {
                display.setTextColor(SSD1306_WHITE);
            }
            display.setCursor(2, i * 10 + 1);
            display.print(MENU_ITEMS[i].label);
            display.print(": ");

            // Display current value
            switch (MENU_ITEMS[i].editState) {
                case UI_EDIT_SCALE_TYPE: display.print(SCALES[currentScaleIndex].name); break;
                case UI_EDIT_SCALE_KEY: display.print(KEY_NAMES[currentKeyIndex]); break;
                case UI_EDIT_WAVETABLE: display.print(WAVETABLE_NAMES[currentWavetableIndex]); break;
                default: break;
            }
        }
    } else {
        // In an edit screen
        const char* title = MENU_ITEMS[menuSelection].label;
        const char* value = "";
        switch (currentState) {
            case UI_EDIT_SCALE_TYPE: value = SCALES[currentScaleIndex].name; break;
            case UI_EDIT_SCALE_KEY: value = KEY_NAMES[currentKeyIndex]; break;
            case UI_EDIT_WAVETABLE: value = WAVETABLE_NAMES[currentWavetableIndex]; break;
            default: break;
        }
        display.println(title);
        display.drawLine(0, 10, SCREEN_WIDTH, 10, SSD1306_WHITE);
        display.setCursor(10, 25);
        display.setTextSize(2);
        display.print(value);
        display.setTextSize(1);
        display.setCursor(0, 50);
        display.println(F("(Press to confirm)"));
    }

    display.display();
}

void checkSerial() {
    static int state = 0; // 0: Header, 1: Data
    static int headerIdx = 0;
    static const char* header = "NSGRID";
    static uint8_t buffer[SynthEngine::SERIALIZED_GRID_SIZE];
    static int bufferIdx = 0;

    while (Serial.available()) {
        uint8_t b = Serial.read();
        if (state == 0) {
            if (b == header[headerIdx]) {
                headerIdx++;
                if (headerIdx == 6) {
                    state = 1;
                    bufferIdx = 0;
                    headerIdx = 0;
                }
            } else {
                headerIdx = 0;
                if (b == 'N') headerIdx = 1; 
            }
        } else if (state == 1) {
            buffer[bufferIdx++] = b;
            if (bufferIdx == SynthEngine::SERIALIZED_GRID_SIZE) {
                if (globalSynth) {
                    globalSynth->importGrid(buffer);
                    saveGridToEEPROM();
                }
                state = 0;
                bufferIdx = 0;
            }
        }
    }
}

void loopUI() {
  handleInput();
  checkSerial();
  drawUI();
  delay(20); // Prevent excessive screen refresh
}