#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::MAX_SERIALIZED_GRID_SIZE];
    size_t size = globalSynth->exportGrid(buf);
    
    EEPROM.write(0, 'N');
    EEPROM.write(1, 'S');
    EEPROM.write(2, (size >> 8) & 0xFF);
    EEPROM.write(3, size & 0xFF);
    for (size_t i = 0; i < size; i++) {
        EEPROM.write(4 + i, buf[i]);
    }
    EEPROM.commit();
}

void loadGridFromEEPROM() {
    if (!globalSynth) return;
    if (EEPROM.read(0) == 'N' && EEPROM.read(1) == 'S') {
        size_t size = (EEPROM.read(2) << 8) | EEPROM.read(3);
        if (size > SynthEngine::MAX_SERIALIZED_GRID_SIZE) size = SynthEngine::MAX_SERIALIZED_GRID_SIZE;
        
        uint8_t buf[SynthEngine::MAX_SERIALIZED_GRID_SIZE];
        for (size_t i = 0; i < size; i++) {
            buf[i] = EEPROM.read(4 + i);
        }
        globalSynth->importGrid(buf, size);
    } 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(2048);

  // 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();

    {
        // Copy grid state to local buffer to minimize lock time
        struct MiniCell {
            uint8_t type;
            uint8_t rotation;
            float value;
        };
        MiniCell gridCopy[SynthEngine::GRID_W][SynthEngine::GRID_H];

        if (globalSynth) {
            SynthLockGuard<SynthMutex> lock(globalSynth->gridMutex);
            for(int x=0; x<SynthEngine::GRID_W; ++x) {
                for(int y=0; y<SynthEngine::GRID_H; ++y) {
                    gridCopy[x][y].type = (uint8_t)globalSynth->grid[x][y].type;
                    gridCopy[x][y].rotation = (uint8_t)globalSynth->grid[x][y].rotation;
                    gridCopy[x][y].value = globalSynth->grid[x][y].value;
                }
            }
        }

        int cellW = 8;
        int cellH = 5;
        int marginX = 2;
        int marginY = (SCREEN_HEIGHT - (SynthEngine::GRID_H * cellH)) / 2;

        for(int x=0; x<SynthEngine::GRID_W; ++x) {
            for(int y=0; y<SynthEngine::GRID_H; ++y) {
                int px = marginX + x * cellW;
                int py = marginY + y * cellH;
                int cx = px + cellW / 2;
                int cy = py + cellH / 2;

                uint8_t type = gridCopy[x][y].type;
                uint8_t rot = gridCopy[x][y].rotation;

                if (type == SynthEngine::GridCell::EMPTY) {
                    display.drawPixel(cx, cy, SSD1306_WHITE);
                } else if (type == SynthEngine::GridCell::SINK) {
                    display.fillRect(px + 1, py + 1, cellW - 2, cellH - 2, SSD1306_WHITE);
                } else {
                    // Draw direction line
                    int dx = 0, dy = 0;
                    switch(rot) {
                        case 0: dy = -2; break; // N
                        case 1: dx = 4; break;  // E
                        case 2: dy = 2; break;  // S
                        case 3: dx = -4; break; // W
                    }
                    display.drawLine(cx, cy, cx + dx, cy + dy, SSD1306_WHITE);

                    if (type == SynthEngine::GridCell::FORK) {
                        if (rot == 0 || rot == 2) display.drawLine(cx - 2, cy, cx + 2, cy, SSD1306_WHITE);
                        else display.drawLine(cx, cy - 2, cx, cy + 2, SSD1306_WHITE);
                    } else if (type >= SynthEngine::GridCell::FIXED_OSCILLATOR && type <= SynthEngine::GridCell::GATE_INPUT) {
                        // Sources: Filled rect
                        display.fillRect(cx - 1, cy - 1, 3, 3, SSD1306_WHITE);
                    } else if (type != SynthEngine::GridCell::WIRE) {
                        // Processors: Hollow rect
                        display.drawRect(cx - 1, cy - 1, 3, 3, SSD1306_WHITE);
                    }
                }
            }
        }

        // Draw Buffer Stats
        int barX = 110;
        int barY = 10;
        int barW = 8;
        int barH = 30;
        
        display.drawRect(barX, barY, barW, barH, SSD1306_WHITE);
        
        int usage = audioBufferUsage;
        int fillH = (usage * (barH - 2)) / AUDIO_BUFFER_SIZE;
        if (fillH > barH - 2) fillH = barH - 2;
        if (fillH < 0) fillH = 0;
        
        display.fillRect(barX + 1, barY + (barH - 1) - fillH, barW - 2, fillH, SSD1306_WHITE);
        
        // display.setCursor(barX - 4, barY + barH + 4);
        // display.print(usage);
    }

    display.display();
}

void checkSerial() {
    static int state = 0; // 0: Header, 1: Count, 2: Data, 3: EndCount
    static int headerIdx = 0;
    static const char* header = "NSGRID";
    static int loadHeaderIdx = 0;
    static const char* loadHeader = "NSLOAD";
    static uint8_t buffer[SynthEngine::MAX_SERIALIZED_GRID_SIZE];
    static int bufferIdx = 0;
    static uint8_t elementCount = 0;

    while (Serial.available()) {
        uint8_t b = Serial.read();
        if (state == 0) {
            if (b == header[headerIdx]) {
                headerIdx++;
                if (headerIdx == 6) {
                    state = 1; // Expect count next
                    bufferIdx = 0;
                    headerIdx = 0;
                    loadHeaderIdx = 0;
                }
            } else {
                headerIdx = 0;
                if (b == 'N') headerIdx = 1; 
            }
            
            if (state == 0) {
                if (b == loadHeader[loadHeaderIdx]) {
                    loadHeaderIdx++;
                    if (loadHeaderIdx == 6) {
                        if (globalSynth) {
                            static uint8_t exportBuf[SynthEngine::MAX_SERIALIZED_GRID_SIZE];
                            size_t size = globalSynth->exportGrid(exportBuf);
                            Serial.write("NSGRID", 6);
                            Serial.write(exportBuf, size);
                            Serial.flush();
                        }
                        loadHeaderIdx = 0;
                        headerIdx = 0;
                    }
                } else {
                    loadHeaderIdx = 0;
                    if (b == 'N') loadHeaderIdx = 1;
                }
            }
        } else if (state == 1) { // Count
            elementCount = b;
            if (1 + elementCount * 5 + 1 > sizeof(buffer)) {
                state = 0;
                bufferIdx = 0;
                Serial.println(F("ERROR: Grid too large"));
            } else {
                buffer[bufferIdx++] = b;
                state = (elementCount == 0) ? 3 : 2;
            }
        } else if (state == 2) { // Data
            buffer[bufferIdx++] = b;
            if (bufferIdx == 1 + elementCount * 5) {
                state = 3;
            }
        } else if (state == 3) { // End Count
            buffer[bufferIdx++] = b;
            if (globalSynth) {
                int result = globalSynth->importGrid(buffer, bufferIdx);
                if (result != 0) {
                    Serial.print("CRC ERROR "); Serial.println(result);
                    globalSynth->clearGrid();
                } else {
                    saveGridToEEPROM();
                    Serial.println(F("OK: Grid Received"));
                }
                state = 0;
                bufferIdx = 0;
            }
        }
    }
}

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