#include "mnemonic_kb.h"

#include <Arduino.h>
#include <string.h>

#include "mnemonic_wordlist.h"

namespace {

static bool s_rng_seeded = false;

void ensure_rng_seeded() {
  if (s_rng_seeded) {
    return;
  }

  uint32_t seed = micros() ^ (millis() << 16);
  for (int i = 0; i < 16; ++i) {
    seed ^= ((uint32_t)analogRead(A0) & 0x3FFu) << (i % 22);
    delayMicroseconds(50);
  }

  randomSeed(seed);
  s_rng_seeded = true;
}

bool append_word(char* out, size_t out_len, size_t* cursor, const char* word) {
  const size_t wlen = strlen(word);

  if (*cursor >= out_len) {
    return false;
  }

  if (*cursor != 0) {
    if (*cursor + 1 >= out_len) {
      return false;
    }
    out[*cursor] = ' ';
    (*cursor)++;
  }

  if (*cursor + wlen >= out_len) {
    return false;
  }

  memcpy(out + *cursor, word, wlen);
  *cursor += wlen;
  out[*cursor] = '\0';
  return true;
}

}  // namespace

bool mnemonic_letter_range(char letter, uint16_t *out_start, uint16_t *out_count) {
  if (!out_start || !out_count) {
    return false;
  }

  if (letter >= 'A' && letter <= 'Z') {
    letter = (char)(letter - 'A' + 'a');
  }
  if (letter < 'a' || letter > 'z') {
    return false;
  }

  int first = -1;
  int last = -1;
  for (int i = 0; i < 2048; ++i) {
    const char c = BIP39_WORDLIST[i][0];
    if (c == letter) {
      if (first < 0) first = i;
      last = i;
    } else if (first >= 0) {
      break;
    }
  }

  if (first < 0 || last < first) {
    return false;
  }

  *out_start = (uint16_t)first;
  *out_count = (uint16_t)(last - first + 1);
  return true;
}

bool mnemonic_second_letters(char first, char *out, uint8_t *count) {
  if (!out || !count) {
    return false;
  }

  *count = 0;

  if (first >= 'A' && first <= 'Z') {
    first = (char)(first - 'A' + 'a');
  }
  if (first < 'a' || first > 'z') {
    return false;
  }

  uint16_t start = 0;
  uint16_t range_count = 0;
  if (!mnemonic_letter_range(first, &start, &range_count) || range_count == 0) {
    return false;
  }

  bool seen[26] = {false};
  uint8_t found = 0;
  for (uint16_t i = 0; i < range_count; ++i) {
    const char* w = BIP39_WORDLIST[start + i];
    char c = w[1];
    if (c >= 'A' && c <= 'Z') {
      c = (char)(c - 'A' + 'a');
    }
    if (c < 'a' || c > 'z') {
      continue;
    }

    const uint8_t idx = (uint8_t)(c - 'a');
    if (!seen[idx]) {
      seen[idx] = true;
      out[found++] = c;
    }
  }

  *count = found;
  return found > 0;
}

bool mnemonic_prefix_range(char first, char second, uint16_t *out_start, uint16_t *out_count) {
  if (!out_start || !out_count) {
    return false;
  }

  if (first >= 'A' && first <= 'Z') {
    first = (char)(first - 'A' + 'a');
  }
  if (second >= 'A' && second <= 'Z') {
    second = (char)(second - 'A' + 'a');
  }
  if (first < 'a' || first > 'z' || second < 'a' || second > 'z') {
    return false;
  }

  uint16_t start = 0;
  uint16_t range_count = 0;
  if (!mnemonic_letter_range(first, &start, &range_count) || range_count == 0) {
    return false;
  }

  int first_match = -1;
  int last_match = -1;
  for (uint16_t i = 0; i < range_count; ++i) {
    const char* w = BIP39_WORDLIST[start + i];
    char c = w[1];
    if (c >= 'A' && c <= 'Z') {
      c = (char)(c - 'A' + 'a');
    }

    if (c == second) {
      if (first_match < 0) {
        first_match = (int)(start + i);
      }
      last_match = (int)(start + i);
    } else if (first_match >= 0) {
      break;
    }
  }

  if (first_match < 0 || last_match < first_match) {
    return false;
  }

  *out_start = (uint16_t)first_match;
  *out_count = (uint16_t)(last_match - first_match + 1);
  return true;
}

const char* mnemonic_word_at(uint16_t index) {
  if (index >= 2048) {
    return nullptr;
  }
  return BIP39_WORDLIST[index];
}

int mnemonic_word_index(const char* word) {
  if (!word || !word[0]) {
    return -1;
  }

  for (int i = 0; i < 2048; ++i) {
    if (strcmp(word, BIP39_WORDLIST[i]) == 0) {
      return i;
    }
  }

  return -1;
}

int generate_mnemonic_12(char* out, size_t out_len) {
  if (!out || out_len == 0) {
    return -1;
  }

  out[0] = '\0';
  ensure_rng_seeded();

  size_t cursor = 0;
  for (int i = 0; i < 12; ++i) {
    const uint16_t idx = (uint16_t)random(0, 2048);
    const char* word = mnemonic_word_at(idx);
    if (!word) {
      return -1;
    }

    if (!append_word(out, out_len, &cursor, word)) {
      return -1;
    }
  }

  return 0;
}

bool mnemonic_validate_basic(const char* mnemonic) {
  if (!mnemonic || !mnemonic[0]) {
    return false;
  }

  char buf[256];
  strncpy(buf, mnemonic, sizeof(buf) - 1);
  buf[sizeof(buf) - 1] = '\0';

  int count = 0;
  char* saveptr = nullptr;
  char* tok = strtok_r(buf, " ", &saveptr);
  while (tok) {
    if (mnemonic_word_index(tok) < 0) {
      return false;
    }

    ++count;
    tok = strtok_r(nullptr, " ", &saveptr);
  }

  return count == 12;
}
