#include <SPI.h>
#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>
#include <Adafruit_NeoPixel.h>

// --- BLUETOOTH CONFIGURATION (Earle Philhower Core) ---
#if defined(ARDUINO_ARCH_RP2040) && defined(ARDUINO_RASPBERRY_PI_PICO_W)
  #define ENABLE_BTSTACK
  #include <btstack.h>
  
  // BTStack System Locks (provided by the core)
  extern "C" {
    void __lockBluetooth();
    void __unlockBluetooth();
  }
#endif

// --- HARDWARE CONFIGURATION ---
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
#define OLED_RESET    -1 // Reset pin # (or -1 if sharing Arduino reset pin)
#define SCREEN_ADDRESS 0x3C // See datasheet for Address; 0x3D for 128x64, 0x3C for 128x32

// Pin Definitions for Raspberry Pi Pico (RP2040)
#define PIN_SDA 4
#define PIN_SCL 5

#define ENC_CLK 12
#define ENC_DT  13
#define ENC_SW  14

// NeoPixel Pin (any GPIO is fine, I've chosen 16)
#define PIN_NEOPIXEL 16
#define NUM_PIXELS   64 // For 8x8 WS2812B matrix

// --- TRACKER DATA ---
#define NUM_STEPS 16

struct Step {
  int8_t note; // MIDI Note (0-127), -1 for OFF
};

Step sequence[NUM_STEPS];

// --- STATE ---
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);
Adafruit_NeoPixel pixels(NUM_PIXELS, PIN_NEOPIXEL, NEO_GRB + NEO_KHZ800);

#ifdef ENABLE_BTSTACK
// BLE State
#define TARGET_DEVICE_NAME "WIDI Thru6"

static hci_con_handle_t con_handle = HCI_CON_HANDLE_INVALID;
static uint16_t midi_char_value_handle = 0; // Handle to write to on remote device
static btstack_packet_callback_registration_t hci_event_callback_registration;
bool wasConnected = false;

// UUIDs (Big Endian to match stack return values)
// Service: 03B80E5A-EDE8-4B33-A751-6CE34EC4C700
static const uint8_t midi_service_uuid[] = {0x03, 0xB8, 0x0E, 0x5A, 0xED, 0xE8, 0x4B, 0x33, 0xA7, 0x51, 0x6C, 0xE3, 0x4E, 0xC4, 0xC7, 0x00};
// Char: 7772E5DB-3868-4112-A1A9-F2669D106BF3 (Little Endian for packet search)
static const uint8_t midi_char_uuid[]    = {0xF3, 0x6B, 0x10, 0x9D, 0x66, 0xF2, 0xA9, 0xA1, 0x12, 0x41, 0x68, 0x38, 0xDB, 0xE5, 0x72, 0x77};
// Char: 7772E5DB-3868-4112-A1A9-F2669D106BF3 (Big Endian fallback)
static const uint8_t midi_char_uuid_be[] = {0x77, 0x72, 0xE5, 0xDB, 0x38, 0x68, 0x41, 0x12, 0xA1, 0xA9, 0xF2, 0x66, 0x9D, 0x10, 0x6B, 0xF3};

static gatt_client_service_t found_service;
static gatt_client_service_t target_service;
static gatt_client_characteristic_t target_characteristic;
static uint16_t midi_cccd_handle = 0;

enum DiscoveryState {
  DISCOVERY_IDLE,
  DISCOVERY_FINDING_SERVICE,
  DISCOVERY_FINDING_CHARACTERISTIC,
  DISCOVERY_FINDING_CCCD,
  DISCOVERY_ENABLING_NOTIFICATIONS,
  DISCOVERY_COMPLETE
};
static DiscoveryState discovery_state = DISCOVERY_IDLE;
static bool service_found = false;
static bool midi_ready = false;

// Uncomment to use Write Request (slower, reliable) instead of Write Command (faster, fire-and-forget)
// #define BLE_USE_WRITE_RESPONSE
#endif

int currentStep = 0;
bool isEditing = false;
int scrollOffset = 0;
bool isPlaying = false;
unsigned long lastStepTime = 0;
int tempo = 120; // BPM


// Encoder State
volatile int encoderDelta = 0;
static uint8_t prevNextCode = 0;
static uint16_t store = 0;

// Button State
bool lastButtonState = HIGH;
unsigned long lastDebounceTime = 0;
bool buttonActive = false;
bool buttonConsumed = false;
unsigned long buttonPressTime = 0;

// Bluetooth Icon (8x8)
const unsigned char bluetooth_icon[] PROGMEM = {
  0x10, // ...1....
  0x18, // ...11...
  0x14, // ...1.1..
  0x52, // .1.1..1.
  0x38, // ..111...
  0x52, // .1.1..1.
  0x14, // ...1.1..
  0x18  // ...11...
};

// --- 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(ENC_DT)) prevNextCode |= 0x02;
  if (digitalRead(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++;
  }
}

#ifdef ENABLE_BTSTACK

// GATT Client Event Handler
void handle_gatt_client_event(uint8_t packet_type, uint16_t channel, uint8_t *packet, uint16_t size) {
  switch(hci_event_packet_get_type(packet)) {
    case GATT_EVENT_SERVICE_QUERY_RESULT:
      gatt_event_service_query_result_get_service(packet, &found_service);
      Serial.print(F("Svc UUID: "));
      if (found_service.uuid16) {
        Serial.println(found_service.uuid16, HEX);
      } else {
        for(int i=0; i<16; i++) { Serial.print(found_service.uuid128[i], HEX); Serial.print(" "); }
        Serial.println();
        if (memcmp(found_service.uuid128, midi_service_uuid, 16) == 0) {
           target_service = found_service;
           service_found = true;
        }
      }
      break;
    case GATT_EVENT_CHARACTERISTIC_QUERY_RESULT:
      {
        Serial.print(F("Char Packet: "));
        for(int i=0; i<size; i++) {
           Serial.print(packet[i], HEX); Serial.print(" ");
        }
        Serial.println();

        // Scan packet for the MIDI Characteristic UUID
        // The packet structure varies, so we search for the 128-bit UUID sequence.
        for(int i=0; i <= size - 16; i++) {
           if (memcmp(&packet[i], midi_char_uuid, 16) == 0 || memcmp(&packet[i], midi_char_uuid_be, 16) == 0) {
              Serial.println(F("Found MIDI Char UUID!"));
              // The Value Handle is typically located 2 bytes before the UUID (in this specific event format)
              midi_char_value_handle = little_endian_read_16(packet, i - 2);
              Serial.print(F("Handle: 0x")); Serial.println(midi_char_value_handle, HEX);
              // Capture the characteristic struct ONLY if this is the correct one
              gatt_event_characteristic_query_result_get_characteristic(packet, &target_characteristic);
           }
        }
      }
      break;
    case GATT_EVENT_ALL_CHARACTERISTIC_DESCRIPTORS_QUERY_RESULT:
      {
        Serial.print(F("Desc Packet: "));
        for(int i=0; i<size; i++) {
           Serial.print(packet[i], HEX); Serial.print(" ");
        }
        Serial.println();

        uint16_t handle = little_endian_read_16(packet, 8);
        const uint8_t *uuid = &packet[10];
        Serial.print(F("Desc Handle: 0x")); Serial.print(handle, HEX);

        // Check for CCCD (0x2902) in standard 128-bit base UUID (Little Endian)
        uint8_t cccd_uuid128[] = {0xFB, 0x34, 0x9B, 0x5F, 0x80, 0x00, 0x00, 0x80, 0x00, 0x10, 0x00, 0x00, 0x02, 0x29, 0x00, 0x00};

        if (memcmp(uuid, cccd_uuid128, 16) == 0) {
          midi_cccd_handle = handle;
          Serial.println(F(" -> Found CCCD!"));
        } else {
          Serial.println();
        }
      }
      break;
    case GATT_EVENT_NOTIFICATION:
      {
        uint16_t value_length = gatt_event_notification_get_value_length(packet);
        const uint8_t *value = gatt_event_notification_get_value(packet);
        Serial.print(F("RX MIDI: "));
        for (int i = 0; i < value_length; i++) {
          Serial.print(value[i], HEX); Serial.print(" ");
        }
        Serial.println();
      }
      break;
    case GATT_EVENT_QUERY_COMPLETE:
      if (discovery_state == DISCOVERY_FINDING_SERVICE) {
        if (service_found) {
           Serial.println(F("Service found, searching for characteristic..."));
           discovery_state = DISCOVERY_FINDING_CHARACTERISTIC;
           // Discover ALL characteristics to avoid endianness issues in filters
           gatt_client_discover_characteristics_for_service(handle_gatt_client_event, con_handle, &target_service);
        } else {
           Serial.println(F("MIDI Service not found on device."));
           discovery_state = DISCOVERY_IDLE;
        }
      } else if (discovery_state == DISCOVERY_FINDING_CHARACTERISTIC) {
        if (midi_char_value_handle != 0) {
           Serial.println(F("MIDI Char found. Searching for CCCD..."));
           discovery_state = DISCOVERY_FINDING_CCCD;
           gatt_client_discover_characteristic_descriptors(handle_gatt_client_event, con_handle, &target_characteristic);
        } else {
           Serial.println(F("MIDI Characteristic not found."));
           discovery_state = DISCOVERY_IDLE;
        }
      } else if (discovery_state == DISCOVERY_FINDING_CCCD) {
        if (midi_cccd_handle != 0) {
           Serial.println(F("CCCD found. Enabling Notifications..."));
           discovery_state = DISCOVERY_ENABLING_NOTIFICATIONS;
           uint8_t data[] = {0x01, 0x00}; // Enable Notification (Little Endian)
           gatt_client_write_characteristic_descriptor_using_descriptor_handle(handle_gatt_client_event, con_handle, midi_cccd_handle, 2, data);
        } else {
           Serial.println(F("CCCD not found. Sending without notifications."));
           discovery_state = DISCOVERY_COMPLETE;
           midi_ready = true;
        }
      } else if (discovery_state == DISCOVERY_ENABLING_NOTIFICATIONS) {
         Serial.println(F("Notifications Enabled. BLE MIDI Ready!"));
         discovery_state = DISCOVERY_COMPLETE;
         midi_ready = true;
      }
      break;
  }
}

// BTStack Packet Handler
void packet_handler(uint8_t packet_type, uint16_t channel, uint8_t *packet, uint16_t size) {
  if (packet_type != HCI_EVENT_PACKET) return;
  
  switch (hci_event_packet_get_type(packet)) {
    case GAP_EVENT_ADVERTISING_REPORT: {
      // Check if this is the target device
      uint8_t event_type = gap_event_advertising_report_get_advertising_event_type(packet);
      uint8_t length = gap_event_advertising_report_get_data_length(packet);
      const uint8_t *data = gap_event_advertising_report_get_data(packet);
      bd_addr_t addr;
      gap_event_advertising_report_get_address(packet, addr);

      Serial.print(F("Adv [Type ")); Serial.print(event_type); 
      Serial.print(F("] from ")); Serial.print(bd_addr_to_str(addr));
      Serial.print(F(" len=")); Serial.println(length);
      
      // Parser to find Local Name (0x09 Complete, 0x08 Shortened)
      int i = 0;
      while (i < length) {
        uint8_t len = data[i];
        if (len == 0) break;
        uint8_t type = data[i+1];
        if (type == 0x09 || type == 0x08) { // Complete or Shortened Local Name
          Serial.print(F("  Name: "));
          Serial.write(&data[i+2], len - 1);
          Serial.println();
          
          // Check for match (ensure length matches to avoid partial matches)
          if ((len - 1) == strlen(TARGET_DEVICE_NAME) && 
              memcmp(&data[i+2], TARGET_DEVICE_NAME, len - 1) == 0) {
             Serial.println(F("Found WIDI Thru6! Connecting..."));
             gap_stop_scan();
             gap_connect(addr, (bd_addr_type_t)gap_event_advertising_report_get_address_type(packet));
             return;
          }
        }
        i += len + 1;
      }
      break;
    }
    case HCI_EVENT_LE_META:
      if (hci_event_le_meta_get_subevent_code(packet) == HCI_SUBEVENT_LE_CONNECTION_COMPLETE) {
        con_handle = hci_subevent_le_connection_complete_get_connection_handle(packet);
        Serial.println(F("Connected. Searching for MIDI Service..."));
        
        discovery_state = DISCOVERY_FINDING_SERVICE;
        service_found = false;
        midi_char_value_handle = 0;
        midi_cccd_handle = 0;
        midi_ready = false;
        
        // Discover ALL Services to debug/find MIDI
        gatt_client_discover_primary_services(handle_gatt_client_event, con_handle);
        
        // Optimize Connection
        gatt_client_send_mtu_negotiation(handle_gatt_client_event, con_handle);
        gap_request_connection_parameter_update(con_handle, 12, 24, 0, 3000); // 15ms - 30ms interval
      }
      break;
    case HCI_EVENT_DISCONNECTION_COMPLETE:
      con_handle = HCI_CON_HANDLE_INVALID;
      midi_char_value_handle = 0;
      discovery_state = DISCOVERY_IDLE;
      midi_ready = false;
      Serial.println(F("Disconnected. Restarting scan..."));
      gap_start_scan();
      break;
  }
}

void setupBluetooth() {
  __lockBluetooth();

  l2cap_init();
  gatt_client_init();
  sm_init();
  sm_set_io_capabilities(IO_CAPABILITY_NO_INPUT_NO_OUTPUT);

  // Register Callback & Start
  hci_event_callback_registration.callback = &packet_handler;
  hci_add_event_handler(&hci_event_callback_registration);
  
  // Start Scanning
  gap_set_scan_parameters(1, 0x0030, 0x0030); // 1 = Active Scanning (requests Scan Response)
  gap_start_scan();
  Serial.println(F("Scanning started..."));
  
  hci_power_control(HCI_POWER_ON);

  __unlockBluetooth();
}
#endif

void setup() {
  Serial.begin(115200);
  delay(5000);
  Serial.println(F("Starting."));

  // 1. Setup Encoder
  pinMode(ENC_CLK, INPUT_PULLUP);
  pinMode(ENC_DT, INPUT_PULLUP);
  pinMode(ENC_SW, INPUT_PULLUP);
  
  attachInterrupt(digitalPinToInterrupt(ENC_CLK), readEncoder, CHANGE);
  attachInterrupt(digitalPinToInterrupt(ENC_DT), readEncoder, CHANGE);

  // 2. Setup Display
  // Note: Using default I2C pins (SDA=GP4, SCL=GP5) which works on both cores.
  Wire.begin();

  // SSD1306_SWITCHCAPVCC = generate display voltage from 3.3V internally
  if(!display.begin(SSD1306_SWITCHCAPVCC, SCREEN_ADDRESS)) {
    Serial.println(F("SSD1306 allocation failed"));
    for(;;); // Don't proceed, loop forever
  }

  // 3. Setup NeoPixel Matrix
  pixels.begin();
  pixels.setBrightness(40); // Set brightness to a medium-low value (0-255)
  pixels.clear();
  pixels.show();

  // 4. Init Sequence
  for(int i=0; i<NUM_STEPS; i++) {
    sequence[i].note = -1; // Default to empty
  }
  // Add a simple C-Major scale for testing
  sequence[0].note = 60; // C4
  sequence[2].note = 62; // D4
  sequence[4].note = 64; // E4
  sequence[6].note = 65; // F4
  sequence[8].note = 65; // F4
  sequence[9].note = 62; // D4
  sequence[12].note = 64; // E4

  display.clearDisplay();
  display.display();

  // 5. Setup BLE
#ifdef ENABLE_BTSTACK
  setupBluetooth();
  Serial.println(F("BTStack initialized."));
#else
  Serial.println(F("BLE Disabled (Requires Earle Philhower Core + Pico W)"));
#endif

  Serial.println(F("Started."));
}

void sendMidi(uint8_t status, uint8_t note, uint8_t velocity) {
#ifdef ENABLE_BTSTACK
  if (con_handle != HCI_CON_HANDLE_INVALID && midi_char_value_handle != 0) {
    // Use 0 timestamp for "immediate" execution to rule out timing/sync issues
    // Header: 10xxxxxx (Bit 7=1, Bit 6=0 for MIDI, Bits 5-0 = high 6 bits of timestamp)
    // Timestamp: 1xxxxxxx (Bit 7=1, Bits 6-0 = low 7 bits of timestamp)
    uint8_t header = 0x80; 
    uint8_t timestamp = 0x80;
    uint8_t midiPacket[] = { header, timestamp, status, note, velocity };
    
    __lockBluetooth();
#ifdef BLE_USE_WRITE_RESPONSE
    uint8_t err = gatt_client_write_value_of_characteristic(handle_gatt_client_event, con_handle, midi_char_value_handle, sizeof(midiPacket), midiPacket);
#else
    uint8_t err = gatt_client_write_value_of_characteristic_without_response(con_handle, midi_char_value_handle, sizeof(midiPacket), midiPacket);
#endif
    __unlockBluetooth();
    
    if (err) {
      Serial.print(F("BLE Write Error: 0x")); Serial.println(err, HEX);
    } else {
      Serial.print(F("MIDI TX: ")); 
      for(int i=0; i<sizeof(midiPacket); i++) {
        Serial.print(midiPacket[i], HEX); Serial.print(" ");
      }
      Serial.println();
    }
  }
#endif
}

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

  if (delta != 0) {
    if (isEditing) {
      // Change Note
      int newNote = sequence[currentStep].note + delta;
      if (newNote < -1) newNote = -1;
      if (newNote > 127) newNote = 127;
      sequence[currentStep].note = newNote;
      Serial.print(F("Note changed: ")); Serial.println(newNote);
    } else {
      // Move Cursor
      currentStep += (delta > 0 ? 1 : -1);
      if (currentStep < 0) currentStep = NUM_STEPS - 1;
      if (currentStep >= NUM_STEPS) currentStep = 0;
      
      // Adjust Scroll to keep cursor in view
      if (currentStep < scrollOffset) scrollOffset = currentStep;
      if (currentStep >= scrollOffset + 6) scrollOffset = currentStep - 5;
    }
  }

  // Handle Button
  int reading = digitalRead(ENC_SW);
  
  if (reading != lastButtonState) {
    lastDebounceTime = millis();
  }
  
  if ((millis() - lastDebounceTime) > 50) {
    if (reading == LOW && !buttonActive) {
       // Button Pressed
       buttonActive = true;
       buttonPressTime = millis();
       buttonConsumed = false;
       Serial.println(F("Button Down"));
    }
    
    if (reading == HIGH && buttonActive) {
       // Button Released
       buttonActive = false;
       if (!buttonConsumed) {
         isEditing = !isEditing;
         Serial.print(F("Mode toggled: ")); Serial.println(isEditing ? F("EDIT") : F("NAV"));
       }
    }
  }
  
  // Check for Long Press (Start/Stop Playback)
  if (buttonActive && !buttonConsumed && (millis() - buttonPressTime > 600)) {
    isPlaying = !isPlaying;
    buttonConsumed = true; // Prevent short press action
    Serial.print(F("Playback: ")); Serial.println(isPlaying ? F("ON") : F("OFF"));
    if (!isPlaying) {
      // Send All Notes Off on stop (CC 123)
      sendMidi(0xB0, 123, 0); 
    }
  }
  
  lastButtonState = reading;
}

void handlePlayback() {
  if (!isPlaying) return;

  unsigned long interval = 15000 / tempo; // 16th notes (60000 / tempo / 4)
  if (millis() - lastStepTime > interval) {
    lastStepTime = millis();
    
    // Note Off for current step (before advancing)
    if (sequence[currentStep].note != -1) {
      sendMidi(0x80, sequence[currentStep].note, 0);
    }
    
    currentStep++;
    if (currentStep >= NUM_STEPS) currentStep = 0;
    
    // Note On for new step
    if (sequence[currentStep].note != -1) {
      sendMidi(0x90, sequence[currentStep].note, 100);
    }
    
    // Auto-scroll logic is handled in drawUI based on currentStep
    if (currentStep < scrollOffset) scrollOffset = currentStep;
    if (currentStep >= scrollOffset + 6) scrollOffset = currentStep - 5;
  }
}

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

  // Header
  display.print(F("TRACKER  "));
  display.print(isEditing ? F("[EDIT]") : F("[NAV]"));
  display.println();
  display.drawLine(0, 8, 128, 8, SSD1306_WHITE);

  // Bluetooth Icon
#ifdef ENABLE_BTSTACK
  if (con_handle != HCI_CON_HANDLE_INVALID) {
    display.drawBitmap(120, 0, bluetooth_icon, 8, 8, SSD1306_WHITE);
    if (midi_ready) {
      display.fillCircle(114, 4, 2, SSD1306_WHITE); // Small dot to indicate MIDI Ready
    }
  }
#endif

  // Steps
  int y = 10;
  for (int i = scrollOffset; i < min(scrollOffset + 6, NUM_STEPS); i++) {
    
    // Draw Cursor
    if (i == currentStep) {
      display.fillRect(0, y, 128, 8, SSD1306_WHITE);
      display.setTextColor(SSD1306_BLACK, SSD1306_WHITE); // Invert text
    } else {
      display.setTextColor(SSD1306_WHITE);
    }

    display.setCursor(2, y);
    
    // Step Number
    if (i < 10) display.print(F("0"));
    display.print(i);
    display.print(F(" | "));

    // Note Value
    int n = sequence[i].note;
    if (n == -1) {
      display.print(F("---"));
    } else {
      // Basic Note to String conversion
      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); // Octave
    }

    y += 9;
  }

  display.display();
}

// Helper to convert X,Y to pixel index for an 8x8 matrix.
// Assumes row-major wiring (NOT serpentine).
// If your matrix is wired differently, you'll need to change this function.
int getPixelIndex(int x, int y) {
  return y * 8 + x;
}

void updateLeds() {
  pixels.clear(); // Clear buffer

  for (int s = 0; s < NUM_STEPS; s++) {
    int blockX = (s % 4) * 2;
    int blockY = (s / 4) * 2;

    uint32_t color;

    if (s == currentStep) {
      if (isEditing) {
        color = pixels.Color(50, 0, 0); // Dim Red for editing
      } else {
        color = pixels.Color(40, 40, 40); // Dim White for current step
      }
    } else {
      if (sequence[s].note != -1) {
        color = pixels.Color(0, 0, 50); // Dim Blue for step with note
      } else {
        color = 0; // Off
      }
    }

    // Set the 4 pixels for the 2x2 block
    pixels.setPixelColor(getPixelIndex(blockX, blockY), color);
    pixels.setPixelColor(getPixelIndex(blockX + 1, blockY), color);
    pixels.setPixelColor(getPixelIndex(blockX, blockY + 1), color);
    pixels.setPixelColor(getPixelIndex(blockX + 1, blockY + 1), color);
  }

  pixels.show();
}

void loop() {
  // BTStack runs in background, no poll needed.
  
#ifdef ENABLE_BTSTACK
  bool connected = (con_handle != HCI_CON_HANDLE_INVALID);
  if (connected != wasConnected) {
    wasConnected = connected;
    if (connected) Serial.println(F("BLE: Connected"));
    else Serial.println(F("BLE: Disconnected"));
  }
#endif

  handleInput();
  handlePlayback();
  drawUI();
  updateLeds();
  delay(10); // Small delay to prevent screen tearing/excessive refresh
}