#include <Adafruit_TinyUSB.h>

extern "C" {
#include "src/nostr_core/nostr_common.h"
#include "src/nostr_core/utils.h"
#include "src/nostr_core/nip006.h"
#include "src/nostr_core/nip001.h"
#include "src/cjson/cJSON.h"
}

#include "display.h"
#include "mnemonic_kb.h"

// ============================================================
// KB2040 Hidden Signer
// - Composite USB: HID Consumer + WebUSB vendor transport
// - Media mode default
// - Hidden signer mode via PLAY + PREV press chord
// ============================================================

static const uint8_t desc_hid_report[] = {
    TUD_HID_REPORT_DESC_CONSUMER(HID_REPORT_ID(1))
};

Adafruit_USBD_HID usb_hid;
Adafruit_USBD_WebUSB usb_web;
WEBUSB_URL_DEF(landing_url, 1 /* https */, "example.com/nsigner/kb2040");

// ------------------ Hardware pin mapping --------------------
static constexpr uint8_t PIN_ENC_A = 2;
static constexpr uint8_t PIN_ENC_B = 3;
static constexpr uint8_t PIN_BTN_PLAY = 5;
static constexpr uint8_t PIN_BTN_NEXT = 6;
static constexpr uint8_t PIN_BTN_PREV = 7;

static constexpr int32_t DETENTS_PER_STEP = 4;
static constexpr uint32_t DEBOUNCE_MS = 20;

static bool s_display_init_done = false;
static constexpr uint32_t DISPLAY_INIT_DELAY_MS = 750;

// ------------------ Transport diagnostics -------------------
// These counters let us confirm, without any host tooling, whether
// request bytes actually reach the device and whether the WebUSB
// vendor interface believes it is "connected".
static uint32_t s_diag_bytes_rx = 0;      // total bytes pulled from usb_web
static uint32_t s_diag_frames_rx = 0;     // complete framed requests parsed
static uint32_t s_diag_responses_tx = 0;  // responses written back
static bool s_diag_last_connected = false;
static uint32_t s_diag_last_report_ms = 0;
static constexpr uint32_t DIAG_REPORT_INTERVAL_MS = 500;

// ---------------------- Modes / state -----------------------
enum device_mode_t {
  MODE_MEDIA = 0,
  MODE_SIGNER = 1,
};

enum signer_ui_state_t {
  SIGNER_UI_MENU = 0,
  SIGNER_UI_READY = 1,
  SIGNER_UI_GENERATE_REVIEW = 2,
  SIGNER_UI_ENTER_GRID = 3,
};

static device_mode_t s_mode = MODE_MEDIA;
static signer_ui_state_t s_signer_ui_state = SIGNER_UI_MENU;
static bool s_mode_switch_latched = false;

// Demo signer state (placeholder until full n_signer crypto integration)
static char s_mnemonic[256] = "";
static bool s_seed_loaded = false;
static bool s_auto_approve = false;

// Signer UI state
static uint8_t s_signer_menu_selected = 0;  // 0=Generate, 1=Enter
static uint8_t s_signer_review_page = 0;    // legacy placeholder (review now renders all 12 words at once)
static char s_generated_mnemonic[256] = "";
static uint8_t s_enter_slot = 0;                 // 0..11
static uint8_t s_enter_stage = 1;                // 1=first letter, 2=second letter, 3=word select
static uint8_t s_enter_first_letter_index = 0;   // 0..25 => a..z
static char s_enter_second_letters[26] = {0};
static uint8_t s_enter_second_count = 0;
static uint8_t s_enter_second_index = 0;
static uint16_t s_enter_prefix_start = 0;
static uint16_t s_enter_prefix_count = 0;
static uint16_t s_enter_candidate_offset = 0;
static uint16_t s_enter_selected_indices[12] = {0};
static bool s_enter_committed[12] = {false};

static void signer_enter_ready_screen();
static void set_device_mode(device_mode_t mode, bool blink_led);

// -------------------- Transport framing ----------------------
enum rpc_transport_t {
  RPC_TRANSPORT_WEBUSB = 0,
  RPC_TRANSPORT_SERIAL = 1,
};

static rpc_transport_t s_rpc_reply_transport = RPC_TRANSPORT_WEBUSB;

static uint8_t s_rx_web_buf[2048];
static size_t s_rx_web_len = 0;

static uint8_t s_rx_serial_buf[2048];
static size_t s_rx_serial_len = 0;

static uint8_t s_tx_buf[2048];

// ---------------------- Encoder state ------------------------
volatile int8_t s_isr_last_ab = 0;
volatile int32_t s_detent_accumulator = 0;
uint8_t s_local_volume = 50;

static constexpr int8_t QUAD_TABLE[16] = {
  0, -1, +1,  0,
 +1,  0,  0, -1,
 -1,  0,  0, +1,
  0, +1, -1,  0
};

// ---------------------- Button state -------------------------
struct ButtonState {
  uint8_t pin;
  bool stable_level;   // HIGH=released, LOW=pressed
  bool last_sample;
  uint32_t last_change_ms;
};

ButtonState s_buttons[] = {
  { PIN_BTN_PLAY, true, true, 0 },
  { PIN_BTN_NEXT, true, true, 0 },
  { PIN_BTN_PREV, true, true, 0 },
};

static bool s_evt_play_pressed = false;
static bool s_evt_next_pressed = false;
static bool s_evt_prev_pressed = false;

// ---------------------- Utilities ----------------------------
static bool btn_pressed(int pin) {
  return digitalRead(pin) == LOW;
}

static void led_blink(int n, int on_ms, int off_ms) {
  for (int i = 0; i < n; i++) {
    digitalWrite(LED_BUILTIN, LOW);
    delay(on_ms);
    digitalWrite(LED_BUILTIN, HIGH);
    delay(off_ms);
  }
}

static void send_consumer_key(uint16_t usage) {
  if (!usb_hid.ready()) return;
  usb_hid.sendReport16(1, usage);
  delay(8);   // restore hold for better throughput while keeping reliability
  usb_hid.sendReport16(1, 0);
  delay(4);   // restore smaller inter-key gap for faster throughput
}

static int8_t read_ab() {
  const int a = digitalRead(PIN_ENC_A) ? 1 : 0;
  const int b = digitalRead(PIN_ENC_B) ? 1 : 0;
  return static_cast<int8_t>((a << 1) | b);
}

void on_encoder_edge() {
  const int8_t curr_ab = read_ab();
  const uint8_t idx = static_cast<uint8_t>((s_isr_last_ab << 2) | curr_ab);
  const int8_t delta = QUAD_TABLE[idx & 0x0F];

  if (delta != 0) {
    s_detent_accumulator += delta;
  }

  s_isr_last_ab = curr_ab;
}

// Pop full encoder detents from the ISR accumulator, preserving partial transitions.
static int32_t encoder_take_steps() {
  int32_t transitions = 0;

  noInterrupts();
  transitions = s_detent_accumulator;
  s_detent_accumulator = 0;
  interrupts();

  int32_t steps = 0;
  while (transitions >= DETENTS_PER_STEP) {
    transitions -= DETENTS_PER_STEP;
    steps += 1;
  }

  while (transitions <= -DETENTS_PER_STEP) {
    transitions += DETENTS_PER_STEP;
    steps -= 1;
  }

  if (transitions != 0) {
    noInterrupts();
    s_detent_accumulator += transitions;
    interrupts();
  }

  return steps;
}

static int8_t volume_change_by(int8_t delta) {
  int16_t next = static_cast<int16_t>(s_local_volume) + static_cast<int16_t>(delta);
  if (next < 0) {
    next = 0;
  } else if (next > 100) {
    next = 100;
  }

  const uint8_t clamped = static_cast<uint8_t>(next);
  const int8_t applied = static_cast<int8_t>(static_cast<int16_t>(clamped) - static_cast<int16_t>(s_local_volume));
  s_local_volume = clamped;
  return applied;
}

static void write_u32_be(uint8_t *p, uint32_t v) {
  p[0] = (uint8_t)((v >> 24) & 0xFF);
  p[1] = (uint8_t)((v >> 16) & 0xFF);
  p[2] = (uint8_t)((v >> 8) & 0xFF);
  p[3] = (uint8_t)(v & 0xFF);
}

static uint32_t read_u32_be(const uint8_t *p) {
  return ((uint32_t)p[0] << 24) | ((uint32_t)p[1] << 16) | ((uint32_t)p[2] << 8) | (uint32_t)p[3];
}

static bool json_extract_string(const char *json, const char *key, char *out, size_t out_sz) {
  char needle[64];
  snprintf(needle, sizeof(needle), "\"%s\"", key);
  const char *k = strstr(json, needle);
  if (!k) return false;

  const char *colon = strchr(k, ':');
  if (!colon) return false;

  const char *q1 = strchr(colon, '"');
  if (!q1) return false;
  q1++;

  const char *q2 = strchr(q1, '"');
  if (!q2) return false;

  size_t n = (size_t)(q2 - q1);
  if (n >= out_sz) n = out_sz - 1;
  memcpy(out, q1, n);
  out[n] = '\0';
  return true;
}

static bool json_extract_int(const char *json, const char *key, int *out) {
  if (!json || !key || !out) {
    return false;
  }

  char needle[64];
  snprintf(needle, sizeof(needle), "\"%s\"", key);
  const char *k = strstr(json, needle);
  if (!k) {
    return false;
  }

  const char *colon = strchr(k, ':');
  if (!colon) {
    return false;
  }

  colon++;
  while (*colon == ' ' || *colon == '\t' || *colon == '\r' || *colon == '\n') {
    colon++;
  }

  char *endptr = nullptr;
  long v = strtol(colon, &endptr, 10);
  if (endptr == colon) {
    return false;
  }

  *out = (int)v;
  return true;
}

static bool json_extract_first_param_object(const char *json, char *out, size_t out_sz) {
  if (!json || !out || out_sz < 3) {
    return false;
  }

  const char *params = strstr(json, "\"params\"");
  if (!params) {
    return false;
  }

  const char *arr = strchr(params, '[');
  if (!arr) {
    return false;
  }

  const char *start = strchr(arr, '{');
  if (!start) {
    return false;
  }

  int depth = 0;
  bool in_string = false;
  bool escape = false;
  const char *p = start;

  for (; *p; ++p) {
    const char c = *p;

    if (in_string) {
      if (escape) {
        escape = false;
      } else if (c == '\\') {
        escape = true;
      } else if (c == '"') {
        in_string = false;
      }
      continue;
    }

    if (c == '"') {
      in_string = true;
      continue;
    }

    if (c == '{') {
      depth++;
    } else if (c == '}') {
      depth--;
      if (depth == 0) {
        const size_t n = (size_t)(p - start + 1);
        if (n >= out_sz) {
          return false;
        }
        memcpy(out, start, n);
        out[n] = '\0';
        return true;
      }
    }
  }

  return false;
}

static void secure_zero(void *p, size_t n) {
  if (!p || n == 0) return;
  volatile uint8_t *vp = reinterpret_cast<volatile uint8_t *>(p);
  while (n--) {
    *vp++ = 0;
  }
}

static void signer_apply_seed(const char *mnemonic) {
  if (!mnemonic) {
    return;
  }

  strncpy(s_mnemonic, mnemonic, sizeof(s_mnemonic) - 1);
  s_mnemonic[sizeof(s_mnemonic) - 1] = '\0';
  s_seed_loaded = strlen(s_mnemonic) > 0;
  // display_set_status(s_seed_loaded ? "Seed loaded" : "Seed cleared");
}

static bool wait_for_user_approval(uint32_t timeout_ms) {
  display_show_message("APPROVAL", "Host requests", "signature");
  // display_set_status("Waiting approval");

  uint32_t start = millis();
  while ((millis() - start) < timeout_ms) {
    if (btn_pressed(PIN_BTN_PLAY)) {
      while (btn_pressed(PIN_BTN_PLAY)) delay(5);
      // display_set_status("Approved");
      return true;
    }
    if (btn_pressed(PIN_BTN_PREV)) {
      while (btn_pressed(PIN_BTN_PREV)) delay(5);
      // display_set_status("Denied");
      return false;
    }
    display_tick();
    delay(5);
  }

  // display_set_status("Approval timeout");
  return false;
}

static void send_json_response(const char *json) {
  size_t n = strlen(json);
  if (n > (sizeof(s_tx_buf) - 4)) return;

  write_u32_be(s_tx_buf, (uint32_t)n);
  memcpy(s_tx_buf + 4, json, n);

  if (s_rpc_reply_transport == RPC_TRANSPORT_SERIAL) {
    Serial.write(s_tx_buf, n + 4);
    Serial.flush();
  } else {
    usb_web.write(s_tx_buf, n + 4);
    usb_web.flush();
  }

  s_diag_responses_tx++;
}

static void send_rpc_result(const char *id, const char *result_json) {
  char out[2048];
  const char *rpc_id = (id && id[0]) ? id : "0";
  const char *result = (result_json && result_json[0]) ? result_json : "null";
  snprintf(out, sizeof(out), "{\"jsonrpc\":\"2.0\",\"id\":\"%s\",\"result\":%s}", rpc_id, result);
  send_json_response(out);
}

static void send_rpc_error(const char *id, int code, const char *message) {
  char out[1024];
  const char *rpc_id = (id && id[0]) ? id : "0";
  const char *msg = (message && message[0]) ? message : "error";
  snprintf(out, sizeof(out),
           "{\"jsonrpc\":\"2.0\",\"id\":\"%s\",\"error\":{\"code\":%d,\"message\":\"%s\"}}",
           rpc_id, code, msg);
  send_json_response(out);
}

static void handle_rpc(const char *req_json) {
  char rpc_id[96] = {0};
  (void)json_extract_string(req_json, "id", rpc_id, sizeof(rpc_id));

  int nostr_index = 0;
  (void)json_extract_int(req_json, "nostr_index", &nostr_index);

  char method[64] = {0};
  if (!json_extract_string(req_json, "method", method, sizeof(method))) {
    send_rpc_error(rpc_id, -32600, "invalid request");
    // display_set_status("RPC invalid request");
    return;
  }

  if (strcmp(method, "ping") == 0) {
    send_rpc_result(rpc_id, "\"pong\"");
    // display_set_status("RPC ping");
    return;
  }

  if (strcmp(method, "get_status") == 0) {
    char result[256];
    snprintf(result, sizeof(result),
             "{\"mode\":\"%s\",\"seed_loaded\":%s,\"auto_approve\":%s,\"nostr_index\":%d}",
             s_mode == MODE_MEDIA ? "media" : "signer",
             s_seed_loaded ? "true" : "false",
             s_auto_approve ? "true" : "false",
             nostr_index);
    send_rpc_result(rpc_id, result);
    // display_set_status("RPC get_status");
    return;
  }

  if (strcmp(method, "set_mnemonic") == 0) {
    char phrase[256] = {0};
    if (!json_extract_string(req_json, "mnemonic", phrase, sizeof(phrase))) {
      send_rpc_error(rpc_id, -32602, "missing mnemonic");
      // display_set_status("Mnemonic missing");
      return;
    }

    signer_apply_seed(phrase);
    if (s_mode == MODE_SIGNER && s_seed_loaded) {
      signer_enter_ready_screen();
    }
    send_rpc_result(rpc_id, "\"ok\"");
    return;
  }

  if (strcmp(method, "set_auto_approve") == 0) {
    if (strstr(req_json, "\"value\":true") != nullptr) {
      s_auto_approve = true;
      send_rpc_result(rpc_id, "true");
      // display_set_status("Auto-approve ON");
    } else if (strstr(req_json, "\"value\":false") != nullptr) {
      s_auto_approve = false;
      send_rpc_result(rpc_id, "false");
      // display_set_status("Auto-approve OFF");
    } else {
      send_rpc_error(rpc_id, -32602, "missing value");
      // display_set_status("Auto-approve invalid");
    }
    return;
  }

  if (strcmp(method, "get_public_key") == 0) {
    if (!s_seed_loaded) {
      send_rpc_error(rpc_id, 2014, "seed not loaded");
      return;
    }

    unsigned char private_key[32] = {0};
    unsigned char public_key[32] = {0};
    if (nostr_derive_keys_from_mnemonic(s_mnemonic, nostr_index, private_key, public_key) != NOSTR_SUCCESS) {
      secure_zero(private_key, sizeof(private_key));
      secure_zero(public_key, sizeof(public_key));
      send_rpc_error(rpc_id, -32603, "key derivation failed");
      return;
    }

    char pubhex[65] = {0};
    nostr_bytes_to_hex(public_key, 32, pubhex);

    secure_zero(private_key, sizeof(private_key));
    secure_zero(public_key, sizeof(public_key));

    char result[96];
    snprintf(result, sizeof(result), "\"%s\"", pubhex);
    send_rpc_result(rpc_id, result);
    return;
  }

  if (strcmp(method, "sign_event") == 0) {
    if (!s_seed_loaded) {
      send_rpc_error(rpc_id, 2014, "seed not loaded");
      return;
    }

    if (s_mode != MODE_SIGNER) {
      send_rpc_error(rpc_id, 2015, "not in signer mode");
      return;
    }

    bool approved = s_auto_approve ? true : wait_for_user_approval(30000);
    if (!approved) {
      send_rpc_error(rpc_id, 2001, "user denied or timeout");
      if (s_signer_ui_state == SIGNER_UI_MENU) {
        display_show_signer_menu(s_signer_menu_selected);
      } else if (s_signer_ui_state == SIGNER_UI_READY) {
        signer_enter_ready_screen();
      }
      return;
    }

    char event_obj[1400] = {0};
    if (!json_extract_first_param_object(req_json, event_obj, sizeof(event_obj))) {
      send_rpc_error(rpc_id, -32602, "missing event object");
      return;
    }

    cJSON *event_in = cJSON_Parse(event_obj);
    if (!event_in || !cJSON_IsObject(event_in)) {
      if (event_in) cJSON_Delete(event_in);
      send_rpc_error(rpc_id, -32602, "invalid event object");
      return;
    }

    cJSON *kind_item = cJSON_GetObjectItemCaseSensitive(event_in, "kind");
    cJSON *content_item = cJSON_GetObjectItemCaseSensitive(event_in, "content");
    cJSON *tags_item = cJSON_GetObjectItemCaseSensitive(event_in, "tags");
    cJSON *created_at_item = cJSON_GetObjectItemCaseSensitive(event_in, "created_at");

    if (!cJSON_IsNumber(kind_item) || !cJSON_IsString(content_item)) {
      cJSON_Delete(event_in);
      send_rpc_error(rpc_id, -32602, "event missing kind/content");
      return;
    }

    int kind = kind_item->valueint;
    const char *content = content_item->valuestring ? content_item->valuestring : "";
    time_t created_at = (time_t)time(nullptr);
    if (cJSON_IsNumber(created_at_item)) {
      created_at = (time_t)created_at_item->valuedouble;
    }

    cJSON *tags_dup = nullptr;
    if (cJSON_IsArray(tags_item)) {
      tags_dup = cJSON_Duplicate(tags_item, 1);
      if (!tags_dup) {
        cJSON_Delete(event_in);
        send_rpc_error(rpc_id, -32603, "tags allocation failed");
        return;
      }
    }

    unsigned char private_key[32] = {0};
    unsigned char public_key[32] = {0};
    if (nostr_derive_keys_from_mnemonic(s_mnemonic, nostr_index, private_key, public_key) != NOSTR_SUCCESS) {
      if (tags_dup) cJSON_Delete(tags_dup);
      cJSON_Delete(event_in);
      secure_zero(private_key, sizeof(private_key));
      secure_zero(public_key, sizeof(public_key));
      send_rpc_error(rpc_id, -32603, "key derivation failed");
      return;
    }

    cJSON *signed_event = nostr_create_and_sign_event(kind, content, tags_dup, private_key, created_at);

    secure_zero(private_key, sizeof(private_key));
    secure_zero(public_key, sizeof(public_key));
    if (tags_dup) cJSON_Delete(tags_dup);
    cJSON_Delete(event_in);

    if (!signed_event) {
      send_rpc_error(rpc_id, -32603, "event signing failed");
      return;
    }

    char *signed_event_json = cJSON_PrintUnformatted(signed_event);
    cJSON_Delete(signed_event);

    if (!signed_event_json) {
      send_rpc_error(rpc_id, -32603, "event serialization failed");
      return;
    }

    send_rpc_result(rpc_id, signed_event_json);
    free(signed_event_json);

    if (s_signer_ui_state == SIGNER_UI_MENU) {
      display_show_signer_menu(s_signer_menu_selected);
    } else if (s_signer_ui_state == SIGNER_UI_READY) {
      signer_enter_ready_screen();
    }
    return;
  }

  send_rpc_error(rpc_id, -32601, "method not found");
  // display_set_status("RPC method missing");
}

static void pump_framed_requests(uint8_t *rx_buf, size_t *rx_len, rpc_transport_t transport) {
  while (*rx_len >= 4) {
    const uint32_t body_len = read_u32_be(rx_buf);
    if (body_len > (sizeof(s_rx_web_buf) - 4)) {
      *rx_len = 0;
      s_rpc_reply_transport = transport;
      send_json_response("{\"error\":{\"code\":-32000,\"message\":\"frame too large\"}}");
      // display_set_status("Frame too large");
      return;
    }

    if (*rx_len < (size_t)(4 + body_len)) {
      return;
    }

    s_diag_frames_rx++;

    char req[1536];
    size_t n = body_len;
    if (n >= sizeof(req)) n = sizeof(req) - 1;
    memcpy(req, rx_buf + 4, n);
    req[n] = '\0';

    const size_t remain = *rx_len - (4 + body_len);
    memmove(rx_buf, rx_buf + 4 + body_len, remain);
    *rx_len = remain;

    s_rpc_reply_transport = transport;
    handle_rpc(req);
  }
}

static void webusb_pump_frames() {
  while (usb_web.available()) {
    if (s_rx_web_len >= sizeof(s_rx_web_buf)) {
      s_rx_web_len = 0;
      break;
    }

    const int c = usb_web.read();
    if (c < 0) break;
    s_rx_web_buf[s_rx_web_len++] = (uint8_t)c;

    // DIAG: a byte arrived from the host. Pulse LED + count it so we can
    // confirm, with no host tooling, that request bytes reach the device.
    s_diag_bytes_rx++;
    digitalWrite(LED_BUILTIN, LOW);
  }

  pump_framed_requests(s_rx_web_buf, &s_rx_web_len, RPC_TRANSPORT_WEBUSB);
}

static void serial_pump_frames() {
  while (Serial.available() > 0) {
    if (s_rx_serial_len >= sizeof(s_rx_serial_buf)) {
      s_rx_serial_len = 0;
      break;
    }

    const int c = Serial.read();
    if (c < 0) break;
    s_rx_serial_buf[s_rx_serial_len++] = (uint8_t)c;

    s_diag_bytes_rx++;
    digitalWrite(LED_BUILTIN, LOW);
  }

  pump_framed_requests(s_rx_serial_buf, &s_rx_serial_len, RPC_TRANSPORT_SERIAL);
}

static void signer_enter_menu_screen() {
  s_signer_ui_state = SIGNER_UI_MENU;
  display_set_view(DISPLAY_VIEW_SIGNER);
  display_show_signer_menu(s_signer_menu_selected);
  // display_set_status("ENC scroll PLAY select");
}

static void signer_enter_ready_screen() {
  s_signer_ui_state = SIGNER_UI_READY;
  display_set_view(DISPLAY_VIEW_SIGNER);
  display_show_signer_ready();
  // display_set_status("Ready to sign");
}

static char signer_first_letter() {
  return (char)('a' + (s_enter_first_letter_index % 26));
}

static char signer_second_letter() {
  if (s_enter_second_count == 0 || s_enter_second_index >= s_enter_second_count) {
    return 'a';
  }
  return s_enter_second_letters[s_enter_second_index];
}

static void signer_reset_active_slot_state() {
  s_enter_stage = 1;
  s_enter_first_letter_index = 0;
  s_enter_second_count = 0;
  s_enter_second_index = 0;
  s_enter_prefix_start = 0;
  s_enter_prefix_count = 0;
  s_enter_candidate_offset = 0;
}

static bool signer_prepare_second_letters() {
  s_enter_second_count = 0;
  s_enter_second_index = 0;
  return mnemonic_second_letters(signer_first_letter(), s_enter_second_letters, &s_enter_second_count) && s_enter_second_count > 0;
}

static bool signer_prepare_prefix_candidates() {
  s_enter_prefix_start = 0;
  s_enter_prefix_count = 0;
  s_enter_candidate_offset = 0;
  return mnemonic_prefix_range(signer_first_letter(), signer_second_letter(), &s_enter_prefix_start, &s_enter_prefix_count) &&
         s_enter_prefix_count > 0;
}

static void signer_render_enter_grid() {
  const char *slot_words[12] = {nullptr};
  for (uint8_t i = 0; i < 12; ++i) {
    if (s_enter_committed[i]) {
      slot_words[i] = mnemonic_word_at(s_enter_selected_indices[i]);
    } else {
      slot_words[i] = "";
    }
  }

  char building[16] = "";
  uint8_t cursor_pos = 0;

  if (s_enter_stage == 1) {
    building[0] = signer_first_letter();
    building[1] = '\0';
    cursor_pos = 1;
  } else if (s_enter_stage == 2) {
    building[0] = signer_first_letter();
    building[1] = signer_second_letter();
    building[2] = '\0';
    cursor_pos = 2;
  } else {
    const uint16_t idx = (uint16_t)(s_enter_prefix_start + s_enter_candidate_offset);
    const char *word = mnemonic_word_at(idx);
    if (word) {
      strncpy(building, word, sizeof(building) - 1);
      building[sizeof(building) - 1] = '\0';
    }
    cursor_pos = (uint8_t)strlen(building);
  }

  display_show_enter_mnemonic(slot_words, s_enter_slot, s_enter_stage, building, cursor_pos);
}

static void signer_enter_grid_screen() {
  s_signer_ui_state = SIGNER_UI_ENTER_GRID;
  display_set_view(DISPLAY_VIEW_SIGNER);
  signer_render_enter_grid();
}

static void signer_enter_review_screen() {
  s_signer_ui_state = SIGNER_UI_GENERATE_REVIEW;
  display_show_mnemonic_page(s_generated_mnemonic, 0);
  // display_set_status("PLAY confirm NEXT cancel");
}

static void build_entered_mnemonic(char *out, size_t out_sz) {
  if (!out || out_sz == 0) {
    return;
  }

  out[0] = '\0';
  size_t cursor = 0;

  for (uint8_t i = 0; i < 12; ++i) {
    if (!s_enter_committed[i]) {
      return;
    }

    const char *w = mnemonic_word_at(s_enter_selected_indices[i]);
    if (!w) w = "abandon";

    const size_t wlen = strlen(w);
    if (cursor != 0) {
      if (cursor + 1 >= out_sz) break;
      out[cursor++] = ' ';
    }
    if (cursor + wlen >= out_sz) break;
    memcpy(out + cursor, w, wlen);
    cursor += wlen;
    out[cursor] = '\0';
  }
}

static void set_device_mode(device_mode_t mode, bool blink_led) {
  s_mode = mode;

  display_set_mode(s_mode == MODE_SIGNER);
  if (s_mode == MODE_SIGNER) {
    display_set_view(DISPLAY_VIEW_SIGNER);
    signer_enter_menu_screen();
    // display_set_status("Signer mode");
  } else {
    display_set_view(DISPLAY_VIEW_VOLUME);
    // display_set_status("Media mode");
  }

  if (blink_led) {
    led_blink(s_mode == MODE_SIGNER ? 3 : 1, 80, 80);
  }
}

static void handle_mode_switch_chord() {
  const bool play = btn_pressed(PIN_BTN_PLAY);
  const bool prev = btn_pressed(PIN_BTN_PREV);

  if (play && prev) {
    if (s_mode_switch_latched) {
      return;
    }

    s_mode_switch_latched = true;
    set_device_mode((s_mode == MODE_MEDIA) ? MODE_SIGNER : MODE_MEDIA, true);
    return;
  }

  s_mode_switch_latched = false;
}

static void update_button_events() {
  s_evt_play_pressed = false;
  s_evt_next_pressed = false;
  s_evt_prev_pressed = false;

  const uint32_t now = millis();
  for (auto &btn : s_buttons) {
    const bool sample = digitalRead(btn.pin);

    if (sample != btn.last_sample) {
      btn.last_sample = sample;
      btn.last_change_ms = now;
    }

    if ((now - btn.last_change_ms) >= DEBOUNCE_MS && btn.stable_level != sample) {
      btn.stable_level = sample;

      // Active-low press edge
      if (!btn.stable_level) {
        if (btn.pin == PIN_BTN_PLAY) s_evt_play_pressed = true;
        if (btn.pin == PIN_BTN_NEXT) s_evt_next_pressed = true;
        if (btn.pin == PIN_BTN_PREV) s_evt_prev_pressed = true;
      }
    }
  }
}

static void media_mode_tick() {
  const int32_t steps = encoder_take_steps();
  if (steps > 0) {
    for (int32_t i = 0; i < steps; ++i) {
      volume_change_by(+1);  // keep display clamped to 0..100
      send_consumer_key(HID_USAGE_CONSUMER_VOLUME_INCREMENT);  // always emit HID step
      display_set_volume(s_local_volume);
    }
    // display_set_status("Volume up");
  } else if (steps < 0) {
    for (int32_t i = 0; i < -steps; ++i) {
      volume_change_by(-1);  // keep display clamped to 0..100
      send_consumer_key(HID_USAGE_CONSUMER_VOLUME_DECREMENT);  // always emit HID step
      display_set_volume(s_local_volume);
    }
    // display_set_status("Volume down");
  }

  if (s_evt_play_pressed) {
    send_consumer_key(HID_USAGE_CONSUMER_PLAY_PAUSE);
    // display_set_status("Play/Pause");
  }
  if (s_evt_next_pressed) {
    // Preserve original hardware remap: NEXT wiring is swapped.
    send_consumer_key(HID_USAGE_CONSUMER_SCAN_PREVIOUS_TRACK);
    // display_set_status("Prev track");
  }
  if (s_evt_prev_pressed) {
    send_consumer_key(HID_USAGE_CONSUMER_SCAN_NEXT_TRACK);
    // display_set_status("Next track");
  }
}

static void signer_mode_tick() {
  const int32_t steps = encoder_take_steps();

  // Suppress actions while chord is held for mode switching.
  if (btn_pressed(PIN_BTN_PLAY) && btn_pressed(PIN_BTN_PREV)) {
    return;
  }

  switch (s_signer_ui_state) {
    case SIGNER_UI_MENU: {
      if (steps != 0) {
        int16_t v = (int16_t)s_signer_menu_selected + (steps > 0 ? 1 : -1);
        if (v < 0) v = 0;
        if (v > 1) v = 1;
        s_signer_menu_selected = (uint8_t)v;
        display_show_signer_menu(s_signer_menu_selected);
      }

      if (s_evt_play_pressed) {
        if (s_signer_menu_selected == 0) {
          if (generate_mnemonic_12(s_generated_mnemonic, sizeof(s_generated_mnemonic)) == 0) {
            s_signer_review_page = 0;
            signer_enter_review_screen();
          } else {
            display_show_message("ERROR", "Generate failed", "Try again");
            // display_set_status("Generate failed");
          }
        } else {
          s_enter_slot = 0;
          for (uint8_t i = 0; i < 12; ++i) {
            s_enter_committed[i] = false;
            s_enter_selected_indices[i] = 0;
          }
          signer_reset_active_slot_state();
          signer_enter_grid_screen();
        }
      }
      break;
    }

    case SIGNER_UI_GENERATE_REVIEW: {
      if (s_evt_play_pressed) {
        signer_apply_seed(s_generated_mnemonic);
        signer_enter_ready_screen();
      }

      if (s_evt_next_pressed) {
        signer_enter_menu_screen();
        // display_set_status("Generate canceled");
      }

      if (s_evt_prev_pressed) {
        if (generate_mnemonic_12(s_generated_mnemonic, sizeof(s_generated_mnemonic)) == 0) {
          signer_enter_review_screen();
        } else {
          display_show_message("ERROR", "Generate failed", "Try again");
        }
      }
      break;
    }

    case SIGNER_UI_READY: {
      if (s_evt_next_pressed) {
        signer_enter_menu_screen();
      }
      break;
    }

    case SIGNER_UI_ENTER_GRID: {
      if (steps != 0) {
        if (s_enter_stage == 1) {
          int16_t idx = (int16_t)s_enter_first_letter_index + (steps > 0 ? 1 : -1);
          while (idx < 0) idx += 26;
          while (idx >= 26) idx -= 26;
          s_enter_first_letter_index = (uint8_t)idx;
        } else if (s_enter_stage == 2) {
          if (s_enter_second_count > 0) {
            int16_t idx = (int16_t)s_enter_second_index + (steps > 0 ? 1 : -1);
            while (idx < 0) idx += s_enter_second_count;
            while (idx >= s_enter_second_count) idx -= s_enter_second_count;
            s_enter_second_index = (uint8_t)idx;
          }
        } else if (s_enter_prefix_count > 0) {
          int32_t rel = (int32_t)s_enter_candidate_offset + steps;
          while (rel < 0) rel += s_enter_prefix_count;
          while (rel >= s_enter_prefix_count) rel -= s_enter_prefix_count;
          s_enter_candidate_offset = (uint16_t)rel;
        }

        signer_render_enter_grid();
      }

      if (s_evt_play_pressed) {
        if (s_enter_stage == 1) {
          if (signer_prepare_second_letters()) {
            s_enter_stage = 2;
            signer_render_enter_grid();
          }
        } else if (s_enter_stage == 2) {
          if (signer_prepare_prefix_candidates()) {
            s_enter_stage = 3;
            signer_render_enter_grid();
          }
        } else {
          s_enter_selected_indices[s_enter_slot] = (uint16_t)(s_enter_prefix_start + s_enter_candidate_offset);
          s_enter_committed[s_enter_slot] = true;

          if (s_enter_slot < 11) {
            s_enter_slot++;
            signer_reset_active_slot_state();
            signer_render_enter_grid();
          } else {
            char phrase[256];
            build_entered_mnemonic(phrase, sizeof(phrase));
            if (mnemonic_validate_basic(phrase)) {
              signer_apply_seed(phrase);
              signer_enter_ready_screen();
            } else {
              display_show_message("INVALID", "Bad mnemonic", "Fix and retry");
            }
          }
        }
      }

      if (s_evt_next_pressed) {
        if (s_enter_stage == 3) {
          s_enter_stage = 2;
          signer_render_enter_grid();
        } else if (s_enter_stage == 2) {
          s_enter_stage = 1;
          signer_render_enter_grid();
        } else {
          if (s_enter_slot == 0) {
            signer_enter_menu_screen();
          } else {
            s_enter_slot--;
            s_enter_committed[s_enter_slot] = false;
            s_enter_selected_indices[s_enter_slot] = 0;
            signer_reset_active_slot_state();
            signer_render_enter_grid();
          }
        }
      }
      break;
    }
  }
}

void setup() {
  pinMode(LED_BUILTIN, OUTPUT);
  digitalWrite(LED_BUILTIN, LOW);

  pinMode(PIN_ENC_A, INPUT_PULLUP);
  pinMode(PIN_ENC_B, INPUT_PULLUP);
  pinMode(PIN_BTN_PLAY, INPUT_PULLUP);
  pinMode(PIN_BTN_NEXT, INPUT_PULLUP);
  pinMode(PIN_BTN_PREV, INPUT_PULLUP);

  Serial.begin(115200);

  if (nostr_init() != NOSTR_SUCCESS) {
    // Keep running transport/UI so host can still diagnose, but signing will fail.
  }

  usb_web.setLandingPage(&landing_url);
  usb_web.setStringDescriptor("n_signer WebUSB");
  usb_web.begin();

  usb_hid.setPollInterval(2);
  usb_hid.setReportDescriptor(desc_hid_report, sizeof(desc_hid_report));
  usb_hid.setStringDescriptor("KB2040 Media");
  usb_hid.begin();

  s_isr_last_ab = read_ab();
  attachInterrupt(digitalPinToInterrupt(PIN_ENC_A), on_encoder_edge, CHANGE);
  attachInterrupt(digitalPinToInterrupt(PIN_ENC_B), on_encoder_edge, CHANGE);

  const uint32_t now = millis();
  for (auto &btn : s_buttons) {
    const bool s = digitalRead(btn.pin);
    btn.stable_level = s;
    btn.last_sample = s;
    btn.last_change_ms = now;
  }

  while (!TinyUSBDevice.mounted()) {
    delay(1);
  }

  display_set_mode(false);
  display_set_view(DISPLAY_VIEW_VOLUME);
  // display_set_status("USB mounted");

  digitalWrite(LED_BUILTIN, HIGH);
  led_blink(2, 60, 60);
}

void loop() {
  if (!s_display_init_done && millis() >= DISPLAY_INIT_DELAY_MS) {
    s_display_init_done = true;
    display_init();
    display_set_mode(s_mode == MODE_SIGNER);
    display_set_volume(s_local_volume);
    display_set_view(s_mode == MODE_SIGNER ? DISPLAY_VIEW_SIGNER : DISPLAY_VIEW_VOLUME);
    // display_set_status("Ready");
    if (s_mode == MODE_SIGNER) {
      signer_enter_menu_screen();
    }
  }

  update_button_events();
  handle_mode_switch_chord();

  if (s_mode == MODE_MEDIA) {
    media_mode_tick();
  } else {
    signer_mode_tick();
  }

  serial_pump_frames();
  webusb_pump_frames();

  // DIAG: periodically surface transport activity on the OLED status line so
  // we can confirm whether host request bytes reach the device, and whether
  // the WebUSB vendor interface reports itself "connected". Format:
  //   D c<0|1> b<bytes> f<frames> t<tx>
  // - c1 => usb_web.connected() true (host opened the WebUSB session)
  // - b  => total bytes received from host on the vendor OUT endpoint
  // - f  => complete framed requests parsed
  // - t  => responses written back
  {
    const bool connected = usb_web.connected();
    const uint32_t now = millis();
    if (connected != s_diag_last_connected ||
        (now - s_diag_last_report_ms) >= DIAG_REPORT_INTERVAL_MS) {
      s_diag_last_connected = connected;
      s_diag_last_report_ms = now;

      char diag[24];
      snprintf(diag, sizeof(diag), "D c%d b%lu f%lu t%lu",
               connected ? 1 : 0,
               (unsigned long)s_diag_bytes_rx,
               (unsigned long)s_diag_frames_rx,
               (unsigned long)s_diag_responses_tx);
      // display_set_status(diag);

      // Idle LED state: ON (HIGH) when connected, OFF otherwise. The RX path
      // drives it LOW briefly on each byte, so a flicker == bytes arriving.
      digitalWrite(LED_BUILTIN, connected ? HIGH : LOW);
    }
  }

  display_tick();

  delay(2);
}
