#include <SPI.h>
#include <SD.h>
#include <U8g2lib.h>
#include <Wire.h>
#include "KT0803.h"
#include <ESP_I2S.h>
#include "MP3DecoderHelix.h"

using namespace libhelix;

#define VERSION "1.2"
#define OLED_CS 5
#define OLED_DC 27
#define OLED_RST 14
#define JOY_X 34
#define JOY_Y 35
#define JOY_SW 32
#define SD_MOSI 23
#define SD_MISO 19
#define SD_SCK 18
#define SD_CS 33
#define I2S_BCLK 26
#define I2S_LRCK 25
#define I2S_DOUT 17

constexpr uint32_t SAMPLE_RATE = 44100;
constexpr float TONE_FREQ = 440.0f;
constexpr int MAX_MUSIC_FILES = 20;
constexpr int MAX_FILENAME_LEN = 48;

KT0803L fm(&Wire);
I2SClass I2S;
U8G2_SH1106_128X64_NONAME_F_4W_HW_SPI u8g2(U8G2_R0, OLED_CS, OLED_DC, OLED_RST);

struct CursorPos { int x, y; };
struct JoyStickStatus { bool up, down, left, right, pressed; };

bool swapJoyStickVertical = false;
bool swapJoyStickHorizontal = false;
int dacVolume = 12;
int selectOptionID = 1;
int currentOptionPage = 0;
int buttonPressCount = 0;
bool wasPressed = false;
unsigned long buttonPressTime = 0;

const char* menuNames[] = {
  "Debug Joystick", "FM Brocaster", "Debug PCM5102", "Volume", "Play Music"
};
constexpr int OPTION_COUNT = sizeof(menuNames) / sizeof(menuNames[0]);

char musicFiles[MAX_MUSIC_FILES][MAX_FILENAME_LEN];
int musicFileCount = 0;
int selectedMusic = 0;

File playingFile;
bool musicPlaying = false;
bool musicButtonWasPressed = false;
uint32_t currentAudioRate = SAMPLE_RATE;

void mp3DataCallback(MP3FrameInfo &info, int16_t *pcm, size_t len, void*) {
  if (!pcm || !len || info.nChans < 1) return;

  if ((uint32_t)info.samprate != currentAudioRate) {
    if (I2S.configureTX(info.samprate, I2S_DATA_BIT_WIDTH_16BIT, I2S_SLOT_MODE_STEREO))
      currentAudioRate = info.samprate;
  }

  constexpr size_t CHUNK_SAMPLES = 256;
  int16_t out[CHUNK_SAMPLES * 2];

  if (info.nChans == 1) {
    for (size_t pos = 0; pos < len;) {
      size_t n = min((size_t)CHUNK_SAMPLES, len - pos);
      for (size_t i = 0; i < n; i++) {
        int32_t v = (int32_t)pcm[pos + i] * dacVolume / 255;
        out[i * 2] = out[i * 2 + 1] = (int16_t)v;
      }
      I2S.write(
          reinterpret_cast<const uint8_t*>(out),
          n * 2 * sizeof(int16_t)
      );
      pos += n;
    }
  } else {
    for (size_t pos = 0; pos < len;) {
      size_t n = min((size_t)(CHUNK_SAMPLES * 2), len - pos);
      for (size_t i = 0; i < n; i++)
        out[i] = (int16_t)((int32_t)pcm[pos + i] * dacVolume / 255);
      I2S.write(
          reinterpret_cast<const uint8_t*>(out),
          n * sizeof(int16_t)
      );
      pos += n;
    }
  }
}

MP3DecoderHelix mp3(mp3DataCallback);

void printScreen(int x, int y, const char* text, bool clearScreen = true, bool send = true) {
  if (clearScreen) u8g2.clearBuffer();
  u8g2.setCursor(x, y);
  u8g2.print(text);
  if (send) u8g2.sendBuffer();
}

int readJoystickAxis(int pin) {
  int sum = 0;
  for (int i = 0; i < 4; i++) sum += analogRead(pin);
  return sum / 4;
}

JoyStickStatus getJoyStickerStatus() {
  constexpr int CENTER_LOW = 1500, CENTER_HIGH = 2600;
  int x = readJoystickAxis(JOY_X), y = readJoystickAxis(JOY_Y);
  JoyStickStatus s = {
    x >= CENTER_LOW && x <= CENTER_HIGH && y < 500,
    x >= CENTER_LOW && x <= CENTER_HIGH && y > 3800,
    y >= CENTER_LOW && y <= CENTER_HIGH && x < 500,
    y >= CENTER_LOW && y <= CENTER_HIGH && x > 3800,
    digitalRead(JOY_SW) == LOW
  };
  if (swapJoyStickVertical) std::swap(s.up, s.down);
  if (swapJoyStickHorizontal) std::swap(s.left, s.right);
  return s;
}

void getJoyStickerDebugInfo(bool serialOut) {
  int x = readJoystickAxis(JOY_X), y = readJoystickAxis(JOY_Y);
  bool pressed = digitalRead(JOY_SW) == LOW;
  JoyStickStatus s = getJoyStickerStatus();

  u8g2.setCursor(0, 26); u8g2.print("X: "); u8g2.print(x);
  u8g2.setCursor(0, 38); u8g2.print("Y: "); u8g2.print(y);
  u8g2.setCursor(0, 50); u8g2.print("SW: "); u8g2.print(pressed ? "PRESSED" : "RELEASED");
  u8g2.setCursor(0, 62);
  if (s.up) u8g2.print("Direction: UP");
  else if (s.down) u8g2.print("Direction: DOWN");
  else if (s.left) u8g2.print("Direction: LEFT");
  else if (s.right) u8g2.print("Direction: RIGHT");

  if (serialOut) {
    Serial.print("X:"); Serial.print(x);
    Serial.print(" , Y:"); Serial.print(y);
    Serial.print(", SW:"); Serial.println(pressed ? "PRESSED" : "RELEASED");
  }
}

float roundToOneDecimal(float n) { return (int)(n * 10 + 0.5f) / 10.0f; }

void hanleQuitMenuEvent() {
  bool pressed = getJoyStickerStatus().pressed;
  if (pressed && !wasPressed) {
    buttonPressTime = millis();
    wasPressed = true;
  } else if (!pressed && wasPressed) {
    if (millis() - buttonPressTime >= 1000) currentOptionPage = 0;
    wasPressed = false;
  }
}

void setFM(float freq, bool mute = false) {
  fm.setUSA();
  fm.setFrequency(freq);
  fm.setMute(mute);
}

void drawFM(float freq, bool frame) {
  printScreen(0, 14, frame ? "Brocasting_" : "Brocasting ", true, false);
  u8g2.setCursor(0, 26); u8g2.print("Frequency: "); u8g2.print(freq);
  u8g2.setCursor(0, 38); u8g2.print("Now you can leave");
  u8g2.setCursor(0, 50); u8g2.print("this page.");
  u8g2.sendBuffer();
}

void fmBroadcast() {
  float freq = 98.1f, oldFreq = freq;
  if (!fm.begin(freq, false)) {
    printScreen(0, 14, "FM Module KT0803 not", true, false);
    u8g2.setCursor(0, 26); u8g2.print(" found. Exiting...");
    u8g2.sendBuffer();
    delay(3500);
    currentOptionPage = 0;
    return;
  }
  setFM(freq);
  bool frame = false;
  while (currentOptionPage == 2) {
    int offset = analogRead(JOY_X) - 2048;
    if (abs(offset) > 200) {
      freq = roundToOneDecimal(freq + (offset / 2048.0f) * 0.1f);
      freq = constrain(freq, 88.0f, 108.0f);
    }
    if (oldFreq != freq) {
      setFM(freq);
      oldFreq = freq;
      Serial.print("Frequency: "); Serial.println(freq);
    }
    drawFM(fm.getFrequency(), frame = !frame);
    delay(300);
    hanleQuitMenuEvent();
  }
}

void joyStickDebugMenu() {
  int dots = 1;
  char title[24];
  while (currentOptionPage == 1) {
    strcpy(title, "Debugging joystick");
    for (int i = 0; i < dots; i++) strcat(title, ".");
    printScreen(0, 14, title, true, false);
    getJoyStickerDebugInfo(true);
    u8g2.sendBuffer();
    dots = dots % 4 + 1;
    delay(300);
    hanleQuitMenuEvent();
  }
}

CursorPos printOptionMenu(int x, int y, int availableHeight, int textSize) {
  CursorPos pos = {x, y};
  int maxVisible = availableHeight / textSize;
  int start = selectOptionID - maxVisible / 2;
  start = constrain(start, 0, max(0, OPTION_COUNT - maxVisible));
  int end = min(start + maxVisible, OPTION_COUNT);
  for (int i = start; i < end; i++) {
    pos.y += 12;
    u8g2.setCursor(pos.x, pos.y);
    u8g2.print(i + 1); u8g2.print(": "); u8g2.print(menuNames[i]);
    u8g2.print(selectOptionID == i + 1 ? " [X]" : " [ ]");
  }
  return pos;
}

void handleOptionSelectEvent() {
  JoyStickStatus s = getJoyStickerStatus();
  if (s.down && selectOptionID < OPTION_COUNT) selectOptionID++;
  else if (s.up && selectOptionID > 1) selectOptionID--;
  if (s.pressed) buttonPressCount++;
  if (buttonPressCount >= 2) {
    buttonPressCount = 0;
    currentOptionPage = selectOptionID;
    selectOptionID = 1;
  }
}

void printLogo(int x, int y) {
  u8g2.drawBox(x, y, 128 - x, 1);
  u8g2.setCursor(x + 2, y + 10);
  u8g2.print("FM Brocaster v" VERSION);
}

void homeMenu() {
  bool frame = false;
  unsigned long displayAt = millis(), inputAt = millis(), animateAt = millis();
  while (currentOptionPage == 0) {
    unsigned long now = millis();
    if (now - inputAt >= 500) { inputAt = now; handleOptionSelectEvent(); }
    if (now - animateAt >= 500) { animateAt = now; frame = !frame; }
    if (now - displayAt >= 33) {
      displayAt = now;
      printScreen(0, 14, frame ? "Running_" : "Running ", true, false);
      CursorPos p = printOptionMenu(0, 14, 36, 12);
      printLogo(p.x, p.y + 2);
      u8g2.sendBuffer();
    }
    delay(1);
  }
}

void test_pcm5102(float sineFreq, float squareFreq, float triangleFreq) {
  static float phase[3] = {0, 0, 0};
  constexpr size_t FRAMES = 256;
  int16_t samples[FRAMES * 2];
  float step[3] = {
    TWO_PI * sineFreq / SAMPLE_RATE,
    TWO_PI * squareFreq / SAMPLE_RATE,
    TWO_PI * triangleFreq / SAMPLE_RATE
  };

  for (size_t i = 0; i < FRAMES; i++) {
    float mixed = 0;
    if (sineFreq > 0) mixed += sinf(phase[0]) * 2500.0f;
    if (squareFreq > 0) mixed += (phase[1] < PI ? 1.0f : -1.0f) * 1000.0f;
    if (triangleFreq > 0)
      mixed += (phase[2] < PI ? 2.0f * phase[2] / PI - 1.0f : 3.0f - 2.0f * phase[2] / PI) * 1000.0f;

    int16_t sample = (int16_t)constrain(mixed, -32768.0f, 32767.0f);
    samples[i * 2] = samples[i * 2 + 1] = sample;
    for (int p = 0; p < 3; p++) {
      phase[p] += step[p];
      if (phase[p] >= TWO_PI) phase[p] -= TWO_PI;
    }
  }
  I2S.write((uint8_t*)samples, sizeof(samples));
}

void pcm5102DebugMenu() {
  unsigned long displayAt = 0;
  int frame = 0;
  float sine = TONE_FREQ, square = 0, triangle = 0;
  while (currentOptionPage == 3) {
    test_pcm5102(sine, square, triangle);
    int offset = readJoystickAxis(JOY_X) - 2048;
    if (abs(offset) > 200) sine = roundToOneDecimal(sine + (offset / 2048.0f) * 0.1f);

    offset = readJoystickAxis(JOY_Y) - 2048;
    float speed = abs(offset) / 2048.0f * 0.1f;
    if (offset > 0) triangle = roundToOneDecimal(triangle + speed);
    else if (offset < 0) square = roundToOneDecimal(square + speed);

    if (millis() - displayAt >= 300) {
      displayAt = millis();
      char title[24] = "Debugging PCM5102";
      for (int i = 0; i <= frame; i++) strcat(title, ".");
      frame = (frame + 1) % 4;
      printScreen(0, 14, title, true, false);
      u8g2.setCursor(0, 28); u8g2.print("Sine Freq: "); u8g2.print(sine);
      u8g2.setCursor(0, 40); u8g2.print("Square Freq: "); u8g2.print(square);
      u8g2.setCursor(0, 52); u8g2.print("Triangle Freq: "); u8g2.print(triangle);
      u8g2.sendBuffer();
      hanleQuitMenuEvent();
    }
  }
}

void volumeMenu() {
  unsigned long updateAt = 0;
  while (currentOptionPage == 4) {
    if (millis() - updateAt < 30) continue;
    updateAt = millis();
    JoyStickStatus s = getJoyStickerStatus();
    dacVolume = constrain(dacVolume + (s.left ? 1 : s.right ? -1 : 0), 0, 255);
    int barWidth = dacVolume * 84 / 255;
    u8g2.clearBuffer();
    u8g2.setCursor(0, 14); u8g2.print("Volume");
    u8g2.drawFrame(20, 30, 88, 12);
    u8g2.drawBox(22, 32, barWidth, 8);
    u8g2.sendBuffer();
    hanleQuitMenuEvent();
    delay(1);
  }
}

void scanMusicFiles() {
  musicFileCount = 0;
  File dir = SD.open("/music");
  if (!dir || !dir.isDirectory()) return;
  while (musicFileCount < MAX_MUSIC_FILES) {
    File file = dir.openNextFile();
    if (!file) break;
    if (!file.isDirectory()) {
      const char* name = file.name();
      const char* ext = strrchr(name, '.');
      if (ext && !strcasecmp(ext, ".mp3")) {
        strncpy(musicFiles[musicFileCount], name, MAX_FILENAME_LEN - 1);
        musicFiles[musicFileCount++][MAX_FILENAME_LEN - 1] = '\0';
      }
    }
    file.close();
  }
  dir.close();
}

void handleMusicSelect() {
  int y = readJoystickAxis(JOY_Y);
  if (musicFileCount == 0) return;
  if (y < 500) { selectedMusic = (selectedMusic + musicFileCount - 1) % musicFileCount; delay(200); }
  else if (y > 3800) { selectedMusic = (selectedMusic + 1) % musicFileCount; delay(200); }
}

void drawMusicMenu() {
  u8g2.clearBuffer();
  u8g2.setCursor(0, 10); u8g2.print("Music");
  if (!musicFileCount) {
    u8g2.setCursor(0, 30); u8g2.print("No MP3 files");
  } else {
    for (int row = -1; row <= 1; row++) {
      int index = selectedMusic + row;
      if (index < 0 || index >= musicFileCount) continue;
      int y = 32 + row * 12;
      if (index == selectedMusic) { u8g2.setCursor(0, y); u8g2.print('>'); }
      const char* name = musicFiles[index];
      if (!strncmp(name, "/music/", 7)) name += 7;
      char shown[18];
      strncpy(shown, name, 17); shown[17] = '\0';
      u8g2.setCursor(10, y); u8g2.print(shown);
    }
  }
  u8g2.sendBuffer();
}

void stopMusic() {
  if (playingFile) playingFile.close();
  if (musicPlaying) mp3.end();
  musicPlaying = false;
}

bool startMusic(int index) {
  stopMusic();

  char path[MAX_FILENAME_LEN + 8];
  const char* name = musicFiles[index];
  if (name[0] == '/') snprintf(path, sizeof(path), "%s", name);
  else if (!strncmp(name, "music/", 6)) snprintf(path, sizeof(path), "/%s", name);
  else snprintf(path, sizeof(path), "/music/%s", name);

  playingFile = SD.open(path, FILE_READ);
  if (!playingFile) {
    Serial.print("Unable to open: ");
    Serial.println(path);
    return false;
  }

  Serial.print("Playing: ");
  Serial.println(path);
  mp3.begin();
  musicPlaying = true;
  return true;
}

void feedMusicDecoder() {
  if (!musicPlaying || !playingFile) return;

  uint8_t input[1024];
  int n = playingFile.read(input, sizeof(input));
  if (n > 0) {
    mp3.write(input, n);
  } else {
    stopMusic();
    Serial.println("Playback finished");
  }
}

void musicMenu() {
  scanMusicFiles();
  musicButtonWasPressed = digitalRead(JOY_SW) == LOW;

  while (currentOptionPage == 5) {
    if (!musicPlaying) handleMusicSelect();
    drawMusicMenu();

    bool pressed = digitalRead(JOY_SW) == LOW;
    if (pressed && !musicButtonWasPressed && musicFileCount > 0)
      startMusic(selectedMusic);
    musicButtonWasPressed = pressed;

    feedMusicDecoder();
    hanleQuitMenuEvent();
    delay(1);
  }

  stopMusic();
}

void setup() {
  Serial.begin(115200);
  u8g2.begin();
  u8g2.setContrast(255);
  u8g2.enableUTF8Print();
  u8g2.setFont(u8g2_font_6x12_tf);
  Wire.begin();
  Wire.setClock(100000);

  if (psramFound()) {
    Serial.println("PSRAM found!");
    Serial.printf("PSRAM size: %u bytes\n", ESP.getPsramSize());
    Serial.printf("Free PSRAM: %u bytes\n", ESP.getFreePsram());
  } else Serial.println("No PSRAM found.");

  int nDevices = 0;
  for (byte address = 1; address < 127; address++) {
    Wire.beginTransmission(address);
    if (Wire.endTransmission() == 0) {
      Serial.printf("I2C device found at address 0x%02X\n", address);
      nDevices++;
    }
  }
  if (!nDevices) Serial.println("No I2C devices found");

  pinMode(JOY_SW, INPUT_PULLUP);
  analogReadResolution(12);

  I2S.setPins(I2S_BCLK, I2S_LRCK, I2S_DOUT, -1, -1);
  Serial.println(I2S.begin(I2S_MODE_STD, SAMPLE_RATE, I2S_DATA_BIT_WIDTH_16BIT, I2S_SLOT_MODE_STEREO)
    ? "Found PCM5102" : "Unable to init PCM5102!");

  SPI.begin(SD_SCK, SD_MISO, SD_MOSI);
  if (!SD.begin(SD_CS, SPI)) {
    Serial.println("Unable to init SD card module.");
    return;
  }

  if (SD.cardType() == CARD_NONE) Serial.println("No SD card detected.");
  else {
    Serial.printf("SD Card Size: %llu MB\n", SD.cardSize() / (1024ULL * 1024ULL));
    scanMusicFiles();
    for (int i = 0; i < musicFileCount; i++) Serial.println(musicFiles[i]);
  }
  Serial.println("Initialization Finished");
}

void loop() {
  switch (currentOptionPage) {
    case 0: homeMenu(); break;
    case 1: joyStickDebugMenu(); break;
    case 2: fmBroadcast(); break;
    case 3: pcm5102DebugMenu(); break;
    case 4: volumeMenu(); break;
    case 5: musicMenu(); break;
    default: currentOptionPage = 0;
  }
}