#include "maps.h"
#include <EEPROM.h>
#include <Wire.h>
#include <libGravity.h>

// Configuration
const int MAX_DENSITY = 255;
const int MAP_RESOLUTION = 255;
const int MAX_CHAOS = 255;

// EEPROM addrs
const int EEPROM_INIT_ADDR = 0;
const int EEPROM_CV1 = EEPROM_INIT_ADDR + sizeof(byte);
const int EEPROM_CV2 = EEPROM_CV1 + sizeof(int);
const int EEPROM_PATTERNS = EEPROM_CV2 + sizeof(int);
const byte EEPROM_INIT_FLAG = 0xAF; // Update flag to re-init properly allocated memory

// EEPROM Delay Save
const unsigned long SAVE_DELAY_MS = 5000;
bool eeprom_needs_save = false;
unsigned long last_param_change = 0;

// Menus and State
struct PatternState {
  int inst_density[3] = {128, 128, 128};
  int map_x = 0;
  int map_y = 127;
  int chaos_amount = 0;
};
PatternState patterns[5];

int selected_slot = 1; // 0 = Global, 1-5 = Pattern A-E
int active_pattern = 0; // Tracks playback pattern index (0-4)

enum GlobalParam {
  PARAM_GLOBAL_CLK_SRC = 0,
  PARAM_GLOBAL_BPM = 1,
  PARAM_GLOBAL_CV1_DEST = 2,
  PARAM_GLOBAL_CV2_DEST = 3,
  PARAM_GLOBAL_LAST = 4
};

enum CvDest {
  CV_DEST_NONE = 0,
  CV_DEST_KICK_DENS = 1,
  CV_DEST_SNARE_DENS = 2,
  CV_DEST_HHAT_DENS = 3,
  CV_DEST_CHAOS = 4,
  CV_DEST_MAP_X = 5,
  CV_DEST_MAP_Y = 6,
  CV_DEST_PRESET = 7,
  CV_DEST_LAST = 8
};

GlobalParam current_global_param = PARAM_GLOBAL_CLK_SRC;
CvDest cv1_dest = CV_DEST_MAP_X;
CvDest cv2_dest = CV_DEST_CHAOS;
Clock::Source selected_source = Clock::SOURCE_INTERNAL;

// UI & Navigation
enum SelectedParam {
  PARAM_KICK_DENS = 0,
  PARAM_SNARE_DENS = 1,
  PARAM_HIHAT_DENS = 2,
  PARAM_CHAOS = 3,
  PARAM_MAP_X = 4,
  PARAM_MAP_Y = 5,
  PARAM_LAST = 6
};

SelectedParam current_param = PARAM_KICK_DENS;
bool editing_param = false;
bool needs_redraw = true;
unsigned long last_redraw = 0;
const unsigned long REDRAW_DELAY_MS = 30; // ~33fps limit

// Sequencer State
int current_step = 0;

volatile int cv1_val = 0;
volatile int cv2_val = 0;
int last_mapped_slot = -1;

// LFSR State for Chaos
uint16_t lfsr = 0xACE1;

// Math Helper: 1D Linear Interpolation between two bytes
uint8_t lerp(uint8_t a, uint8_t b, uint8_t t) {
  // t is 0-255. returns a if t=0, b if t=255
  return a + (((b - a) * t) >> 8);
}

// Math Helper: Get threshold from 2D map via interpolation
uint8_t GetThreshold(int inst, int step, int x_pos, int y_pos) {
  // x_pos is 0-255 mapped across 4 nodes (0, 1, 2, 3). Distance is 85 (255 / 3)
  // y_pos is 0-255 mapped across 4 nodes (0, 1, 2, 3). Distance is 85 (255 / 3)

  int x_idx = x_pos / 85; 
  int y_idx = y_pos / 85;

  uint8_t x_frac = (x_pos % 85) * 3; // scale remainder 0-84 up to 0-255
  uint8_t y_frac = (y_pos % 85) * 3;

  // Guard against out of bounds if exactly 255
  if (x_idx >= 3) {
    x_idx = 2;
    x_frac = 255;
  }
  if (y_idx >= 3) {
    y_idx = 2;
    y_frac = 255;
  }

  // Read 4 corners from PROGMEM
  uint8_t p00 = pgm_read_byte(&PATTERN_MAPS[x_idx][y_idx][inst][step]);
  uint8_t p10 = pgm_read_byte(&PATTERN_MAPS[x_idx + 1][y_idx][inst][step]);
  uint8_t p01 = pgm_read_byte(&PATTERN_MAPS[x_idx][y_idx + 1][inst][step]);
  uint8_t p11 = pgm_read_byte(&PATTERN_MAPS[x_idx + 1][y_idx + 1][inst][step]);

  // Bilinear interpolation
  uint8_t lerp_top = lerp(p00, p10, x_frac);
  uint8_t lerp_bottom = lerp(p01, p11, x_frac);
  return lerp(lerp_top, lerp_bottom, y_frac);
}

void LoadState() {
  if (EEPROM.read(EEPROM_INIT_ADDR) == EEPROM_INIT_FLAG) {
    EEPROM.get(EEPROM_CV1, cv1_dest);
    EEPROM.get(EEPROM_CV2, cv2_dest);
    EEPROM.get(EEPROM_PATTERNS, patterns);
  }
}

void SaveState() {
  EEPROM.update(EEPROM_INIT_ADDR, EEPROM_INIT_FLAG);
  EEPROM.put(EEPROM_CV1, cv1_dest);
  EEPROM.put(EEPROM_CV2, cv2_dest);
  EEPROM.put(EEPROM_PATTERNS, patterns);
}

// LFSR random bit generator (returns 0 or 1, fast)
uint8_t GetRandomBit() {
  uint8_t bit = ((lfsr >> 0) ^ (lfsr >> 2) ^ (lfsr >> 3) ^ (lfsr >> 5)) & 1;
  lfsr = (lfsr >> 1) | (bit << 15);
  return bit;
}

// Get 8-bit pseudo-random number
uint8_t GetRandomByte() {
  uint8_t r = 0;
  for (int i = 0; i < 8; i++) {
    r = (r << 1) | GetRandomBit();
  }
  return r;
}

void ProcessSequencerTick(uint32_t tick) {
  // Assuming 96 PPQN clock. We want 16th notes.
  // 96 pulses per quarter note / 4 = 24 pulses per 16th note.
  const int PULSES_PER_16TH = 24;

  // Pulse logic outputs low halfway through the 16th note (12 pulses)
  if (tick % PULSES_PER_16TH == 12) {
    for (int i = 0; i < 6; i++) {
      gravity.outputs[i].Low();
    }
  }

  // Handle new 16th note step
  if (tick % PULSES_PER_16TH == 0) {

    // Evaluate CV Preset changes FIRST so params pull from the right struct
    int preset_cv = 0;
    if (cv1_dest == CV_DEST_PRESET) preset_cv += cv1_val;
    if (cv2_dest == CV_DEST_PRESET) preset_cv += cv2_val;
    
    if (cv1_dest == CV_DEST_PRESET || cv2_dest == CV_DEST_PRESET) {
      int mapped_slot = constrain(map(preset_cv, 0, 512, 0, 4), 0, 4);
      if (mapped_slot != last_mapped_slot) {
        last_mapped_slot = mapped_slot;
        if (active_pattern != mapped_slot) {
          active_pattern = mapped_slot;
          selected_slot = mapped_slot + 1;
          needs_redraw = true;
        }
      }
    } else {
      last_mapped_slot = -1;
    }

    PatternState &p = patterns[active_pattern];
    int active_map_x = p.map_x;
    int active_map_y = p.map_y;
    int active_chaos = p.chaos_amount;
    int active_dens[3] = { p.inst_density[0], p.inst_density[1], p.inst_density[2] };

    // Apply CV1
    switch (cv1_dest) {
      case CV_DEST_KICK_DENS: active_dens[0] = constrain(active_dens[0] + (cv1_val / 4), 0, 255); break;
      case CV_DEST_SNARE_DENS: active_dens[1] = constrain(active_dens[1] + (cv1_val / 4), 0, 255); break;
      case CV_DEST_HHAT_DENS: active_dens[2] = constrain(active_dens[2] + (cv1_val / 4), 0, 255); break;
      case CV_DEST_CHAOS: active_chaos = constrain(active_chaos + (cv1_val / 4), 0, 255); break;
      case CV_DEST_MAP_X: active_map_x = constrain(active_map_x + (cv1_val / 4), 0, 255); break;
      case CV_DEST_MAP_Y: active_map_y = constrain(active_map_y + (cv1_val / 4), 0, 255); break;
      default: break;
    }

    // Apply CV2
    switch (cv2_dest) {
      case CV_DEST_KICK_DENS: active_dens[0] = constrain(active_dens[0] + (cv2_val / 4), 0, 255); break;
      case CV_DEST_SNARE_DENS: active_dens[1] = constrain(active_dens[1] + (cv2_val / 4), 0, 255); break;
      case CV_DEST_HHAT_DENS: active_dens[2] = constrain(active_dens[2] + (cv2_val / 4), 0, 255); break;
      case CV_DEST_CHAOS: active_chaos = constrain(active_chaos + (cv2_val / 4), 0, 255); break;
      case CV_DEST_MAP_X: active_map_x = constrain(active_map_x + (cv2_val / 4), 0, 255); break;
      case CV_DEST_MAP_Y: active_map_y = constrain(active_map_y + (cv2_val / 4), 0, 255); break;
      default: break;
    }

    // Evaluate hits for Kick, Snare, HiHats
    for (int inst = 0; inst < 3; inst++) {
      uint8_t threshold = GetThreshold(inst, current_step, active_map_x, active_map_y);
      int active_density = active_dens[inst];

      // Inject chaos
      if (active_chaos > 0) {
        // Chaos randomly adds or subtracts from density.
        int r = GetRandomByte();
        int chaos_variance = map(active_chaos, 0, 255, 0, 128);
        if (GetRandomBit()) {
          active_density += map(r, 0, 255, 0, chaos_variance);
        } else {
          active_density -= map(r, 0, 255, 0, chaos_variance);
        }
        active_density = constrain(active_density, 0, 255);
      }

      // Fire Trigger?
      if (active_density > threshold) {
        // Output 1-3
        gravity.outputs[inst].High();

        // Fire Accent Trigger? (If density greatly exceeds threshold)
        if (active_density > threshold + 60) {
          // Output 4-6
          gravity.outputs[inst + 3].High();
        }
      }
    }

    current_step = (current_step + 1) % 16;
  }
}

void OnPlayPress() {
  if (!gravity.clock.IsPaused()) {
    gravity.clock.Stop();
    for (int i = 0; i < 6; i++)
      gravity.outputs[i].Low();
  } else {
    gravity.clock.Start();
  }
  needs_redraw = true;
}

void OnEncoderPress() {
  editing_param = !editing_param;
  needs_redraw = true;
}

void OnEncoderPressRotate(int val) {
  selected_slot = constrain(selected_slot + val, 0, 5);
  if (selected_slot > 0) {
    active_pattern = selected_slot - 1;
  }
  
  // Reset editing state when swapping pages
  editing_param = false;
  needs_redraw = true;
}

// Map CV Destination enum string labels
const char* GetCvDestName(CvDest dest) {
  switch(dest) {
    case CV_DEST_NONE: return "NONE";
    case CV_DEST_KICK_DENS: return "KICK DENS";
    case CV_DEST_SNARE_DENS: return "SNAR DENS";
    case CV_DEST_HHAT_DENS: return "HHAT DENS";
    case CV_DEST_CHAOS: return "CHAOS";
    case CV_DEST_MAP_X: return "MAP X";
    case CV_DEST_MAP_Y: return "MAP Y";
    case CV_DEST_PRESET: return "PRESET";
    default: return "";
  }
}

void OnEncoderRotate(int val) {
  if (!editing_param) {
    // Navigate menu (clamp to edges, do not wrap)
    if (selected_slot > 0) {
      int next_param = (int)current_param + val;
      next_param = constrain(next_param, 0, PARAM_LAST - 1);
      current_param = (SelectedParam)next_param;
    } else {
      int next_param = (int)current_global_param + val;
      next_param = constrain(next_param, 0, PARAM_GLOBAL_LAST - 1);
      current_global_param = (GlobalParam)next_param;
    }
  } else {
    // Edit parameter
    if (selected_slot > 0) {
      int amt = val * 8; // Adjust by 8 values at a time for speed mapping
      PatternState &p = patterns[active_pattern];

      switch (current_param) {
      case PARAM_KICK_DENS:
        p.inst_density[0] = constrain(p.inst_density[0] + amt, 0, 255);
        break;
      case PARAM_SNARE_DENS:
        p.inst_density[1] = constrain(p.inst_density[1] + amt, 0, 255);
        break;
      case PARAM_HIHAT_DENS:
        p.inst_density[2] = constrain(p.inst_density[2] + amt, 0, 255);
        break;
      case PARAM_CHAOS:
        p.chaos_amount = constrain(p.chaos_amount + amt, 0, 255);
        break;
      case PARAM_MAP_X:
        p.map_x = constrain(p.map_x + amt, 0, 255);
        break;
      case PARAM_MAP_Y:
        p.map_y = constrain(p.map_y + amt, 0, 255);
        break;
      default:
        break;
      }
    } else {
      // Edit Global Parameter
      switch (current_global_param) {
      case PARAM_GLOBAL_CLK_SRC: {
        int src = (int)selected_source + val;
        src = constrain(src, 0, Clock::SOURCE_LAST - 1);
        selected_source = (Clock::Source)src;
        gravity.clock.SetSource(selected_source);
        break;
      }
      case PARAM_GLOBAL_BPM: {
        if (!gravity.clock.ExternalSource()) {
           gravity.clock.SetTempo(gravity.clock.Tempo() + val);
        }
        break;
      }
      case PARAM_GLOBAL_CV1_DEST: {
        int dest = (int)cv1_dest + val;
        dest = constrain(dest, 0, CV_DEST_LAST - 1);
        cv1_dest = (CvDest)dest;
        break;
      }
      case PARAM_GLOBAL_CV2_DEST: {
        int dest = (int)cv2_dest + val;
        dest = constrain(dest, 0, CV_DEST_LAST - 1);
        cv2_dest = (CvDest)dest;
        break;
      }
      default:
        break;
      }
    }

    eeprom_needs_save = true;
    last_param_change = millis();
  }
  needs_redraw = true;
}

#include "display_assets.h"

constexpr uint8_t SCREEN_CENTER_X = 32;
constexpr uint8_t MAIN_TEXT_Y = 26;
constexpr uint8_t SUB_TEXT_Y = 40;
constexpr uint8_t VISIBLE_MENU_ITEMS = 3;
constexpr uint8_t MENU_ITEM_HEIGHT = 14;
constexpr uint8_t MENU_BOX_PADDING = 4;
constexpr uint8_t MENU_BOX_WIDTH = 64;

void drawCenteredText(const char *text, int y, const uint8_t *font) {
  gravity.display.setFont(font);
  int textWidth = gravity.display.getStrWidth(text);
  gravity.display.drawStr(SCREEN_CENTER_X - (textWidth / 2), y, text);
}

void drawRightAlignedText(const char *text, int y) {
  int textWidth = gravity.display.getStrWidth(text);
  int drawX = (SCREEN_WIDTH - textWidth) - MENU_BOX_PADDING;
  gravity.display.drawStr(drawX, y, text);
}

void drawMainSelection() {
  gravity.display.setDrawColor(1);
  const int tickSize = 3;
  const int mainWidth = SCREEN_WIDTH / 2;
  const int mainHeight = 49;
  gravity.display.drawLine(0, 0, tickSize, 0);
  gravity.display.drawLine(0, 0, 0, tickSize);
  gravity.display.drawLine(mainWidth, 0, mainWidth - tickSize, 0);
  gravity.display.drawLine(mainWidth, 0, mainWidth, tickSize);
  gravity.display.drawLine(mainWidth, mainHeight, mainWidth,
                           mainHeight - tickSize);
  gravity.display.drawLine(mainWidth, mainHeight, mainWidth - tickSize,
                           mainHeight);
  gravity.display.drawLine(0, mainHeight, tickSize, mainHeight);
  gravity.display.drawLine(0, mainHeight, 0, mainHeight - tickSize);
}

void drawMenuItems(String menu_items[], int menu_size, int current_item) {
  gravity.display.setFont(TEXT_FONT);

  int selectedBoxY = 0;
  if (menu_size >= VISIBLE_MENU_ITEMS && current_item == menu_size - 1) {
    selectedBoxY = MENU_ITEM_HEIGHT * min(2, current_item);
  } else if (current_item > 0) {
    selectedBoxY = MENU_ITEM_HEIGHT;
  }

  int boxX = MENU_BOX_WIDTH + 1;
  int boxY = selectedBoxY + 2;
  int boxWidth = MENU_BOX_WIDTH - 1;
  int boxHeight = MENU_ITEM_HEIGHT + 1;

  gravity.display.setDrawColor(1);
  if (editing_param) {
    gravity.display.drawBox(boxX, boxY, boxWidth, boxHeight);
    drawMainSelection();
  } else {
    gravity.display.drawFrame(boxX, boxY, boxWidth, boxHeight);
  }

  gravity.display.setDrawColor(2);
  int start_index = 0;
  if (menu_size >= VISIBLE_MENU_ITEMS && current_item == menu_size - 1) {
    start_index = menu_size - VISIBLE_MENU_ITEMS;
  } else if (current_item > 0) {
    start_index = current_item - 1;
  }

  for (int i = 0; i < min(menu_size, (int)VISIBLE_MENU_ITEMS); ++i) {
    int idx = start_index + i;
    drawRightAlignedText(menu_items[idx].c_str(),
                         MENU_ITEM_HEIGHT * (i + 1) - 1);
  }
}

void DisplayMainArea() {
  gravity.display.setFontMode(1);
  gravity.display.setDrawColor(2);

  String mainText;
  String subText;

  if (selected_slot > 0) {
    PatternState &p = patterns[active_pattern];
    switch (current_param) {
    case PARAM_KICK_DENS:
      mainText = String(map(p.inst_density[0], 0, 255, 0, 100)) + "%";
      subText = "KICK DENS";
      break;
    case PARAM_SNARE_DENS:
      mainText = String(map(p.inst_density[1], 0, 255, 0, 100)) + "%";
      subText = "SNAR DENS";
      break;
    case PARAM_HIHAT_DENS:
      mainText = String(map(p.inst_density[2], 0, 255, 0, 100)) + "%";
      subText = "HHAT DENS";
      break;
    case PARAM_CHAOS:
      mainText = String(map(p.chaos_amount, 0, 255, 0, 100)) + "%";
      subText = "CHAOS";
      break;
    case PARAM_MAP_X:
      mainText = String(p.map_x);
      subText = "MAP X";
      break;
    case PARAM_MAP_Y:
      mainText = String(p.map_y);
      subText = "MAP Y";
      break;
    default:
      break;
    }

    drawCenteredText(mainText.c_str(), MAIN_TEXT_Y, LARGE_FONT);
    drawCenteredText(subText.c_str(), SUB_TEXT_Y, TEXT_FONT);

    String menu_items[PARAM_LAST] = {
        "KICK", "SNARE", "HHAT", "CHAOS", "MAP X", "MAP Y"};
    drawMenuItems(menu_items, PARAM_LAST, (int)current_param);
    
  } else {
    switch (current_global_param) {
    case PARAM_GLOBAL_CLK_SRC:
      switch (selected_source) {
      case Clock::SOURCE_INTERNAL:
        mainText = F("INT");
        subText = F("CLOCK");
        break;
      case Clock::SOURCE_EXTERNAL_PPQN_24:
        mainText = F("EXT");
        subText = F("24 PPQN");
        break;
      case Clock::SOURCE_EXTERNAL_PPQN_4:
        mainText = F("EXT");
        subText = F("4 PPQN");
        break;
      case Clock::SOURCE_EXTERNAL_PPQN_2:
        mainText = F("EXT");
        subText = F("2 PPQN");
        break;
      case Clock::SOURCE_EXTERNAL_PPQN_1:
        mainText = F("EXT");
        subText = F("1 PPQN");
        break;
      case Clock::SOURCE_EXTERNAL_MIDI:
        mainText = F("EXT");
        subText = F("MIDI");
        break;
      }
      break;
    case PARAM_GLOBAL_BPM:
      if (gravity.clock.ExternalSource()) {
        mainText = F("EXT");
      } else {
        mainText = String(gravity.clock.Tempo());
      }
      subText = F("BPM");
      break;
    case PARAM_GLOBAL_CV1_DEST:
      mainText = F("CV1");
      subText = GetCvDestName(cv1_dest);
      break;
    case PARAM_GLOBAL_CV2_DEST:
      mainText = F("CV2");
      subText = GetCvDestName(cv2_dest);
      break;
    default:
      break;
    }
  
    drawCenteredText(mainText.c_str(), MAIN_TEXT_Y, LARGE_FONT);
    drawCenteredText(subText.c_str(), SUB_TEXT_Y, TEXT_FONT);

    String menu_items[PARAM_GLOBAL_LAST] = {
        "SOURCE", "TEMPO", "CV1 DEST", "CV2 DEST"};
    drawMenuItems(menu_items, PARAM_GLOBAL_LAST, (int)current_global_param);
  }
}

void DisplayBottomBar() {
  int boxY = 50;
  int boxHeight = 14;

  gravity.display.setDrawColor(1);
  gravity.display.drawHLine(1, boxY, SCREEN_WIDTH - 2);
  gravity.display.drawVLine(SCREEN_WIDTH - 1, boxY, boxHeight);

  for (int i = 0; i <= 5; i++) {
    int x = i * 21;
    int bw = (i == 5) ? (SCREEN_WIDTH - x) : 21;

    gravity.display.setDrawColor(1);
    if (selected_slot == i) {
      gravity.display.drawBox(x, boxY, bw, boxHeight);
    } else {
      gravity.display.drawVLine(x, boxY, boxHeight);
    }

    gravity.display.setDrawColor(2);
    if (i == 0) {
      gravity.display.setBitmapMode(1);
      auto icon = gravity.clock.IsPaused() ? pause_icon : play_icon;
      int iconX = x + (bw / 2) - (play_icon_width / 2);
      gravity.display.drawXBMP(iconX, boxY, play_icon_width, play_icon_height, icon);
    } else {
      char label[2] = {(char)('A' + i - 1), '\0'};
      gravity.display.setFont(TEXT_FONT);
      int patW = gravity.display.getStrWidth(label);
      int patX = x + (bw / 2) - (patW / 2);
      gravity.display.drawStr(patX, 64 - 3, label);
    }
  }
}

void UpdateDisplay() {
  gravity.display.setFontMode(1);
  DisplayMainArea();
  DisplayBottomBar();
}

void setup() {
  gravity.Init();
  LoadState();

  gravity.play_button.AttachPressHandler(OnPlayPress);
  gravity.encoder.AttachPressHandler(OnEncoderPress);
  gravity.encoder.AttachRotateHandler(OnEncoderRotate);
  gravity.encoder.AttachPressRotateHandler(OnEncoderPressRotate);

  gravity.clock.AttachIntHandler(ProcessSequencerTick);
  // Default to 120 BPM internal
  gravity.clock.SetTempo(120);
  gravity.clock.SetSource(Clock::SOURCE_INTERNAL);

  // Speed up I2C for faster OLED refreshing
  Wire.setClock(400000);
}

void loop() {
  gravity.Process();

  // Apply CV modulation
  // CV1 modulates Map X, CV2 modulates Chaos
  cv1_val = gravity.cv1.Read(); // -512 to 512
  cv2_val = gravity.cv2.Read();

  if (eeprom_needs_save && (millis() - last_param_change > SAVE_DELAY_MS)) {
    SaveState();
    eeprom_needs_save = false;
  }

  if (needs_redraw && (millis() - last_redraw > REDRAW_DELAY_MS)) {
    needs_redraw = false;
    last_redraw = millis();
    gravity.display.firstPage();
    do {
      UpdateDisplay();
    } while (gravity.display.nextPage());
  }
}