/*
 * THUGS(red) Wardrive - wardriving firmware for the LilyGO T-Dongle-S3
 * -------------------------------------------------------------------
 * Target : LilyGO T-Dongle-S3 (ESP32-S3, 16 MB flash, 8 MB PSRAM)
 * GPS    : u-blox NEO-8M or NEO-6M on the dongle's 4-pin JST SH pads
 * Screen : the onboard ST7735 80x160 LCD, used landscape
 * API    : https://wardrive.thugs.red/api/v1/ingest
 *
 * Same firmware as thugs-wardrive-esp32c5.ino, same two modes, three
 * differences that come from the hardware:
 *
 *   1. The S3 is 2.4 GHz only. There is no 5 GHz band to alternate with, so a
 *      scan cycle is one scan. If you want 5 GHz, that is the ESP32-C5 build.
 *   2. Status goes to the onboard RGB LED (an APA102) as colour rather than to
 *      a single LED as a blink pattern - and to the LCD, which is not optional
 *      here because it is soldered on.
 *   3. The GPS lives on GPIO43/44, which are the pads LilyGO broke out. Those
 *      are also UART0, so the serial console moves to native USB.
 *
 * Behaviour
 *   Power on - plug it into a car USB socket - and it wardrives: scan, stamp
 *   with the GPS fix, buffer. The radio never associates while scanning.
 *
 *   Press the button while it is running and it stops scanning, joins the
 *   access point you configured below, and POSTs the buffer in batches. On
 *   success the buffer is cleared and it returns to scanning with a clean
 *   slate. On failure it keeps the buffer and returns to scanning anyway.
 *
 * Build
 *   Arduino IDE, esp32 board package 3.x, board "ESP32S3 Dev Module", and:
 *     USB CDC On Boot : Enabled      <- the log goes over USB, GPIO43/44 are the GPS
 *     Flash Size      : 16MB
 *     PSRAM           : OPI PSRAM
 *   Libraries: TinyGPSPlus, ArduinoJson v7, Adafruit ST7735 + Adafruit GFX.
 *
 *   Because UART0 is now the GPS, the board cannot fall into UART download
 *   mode. To flash it: unplug, hold the button, plug it in, release. The S3's
 *   native USB does the rest.
 */

#include <Arduino.h>
#include <WiFi.h>
#include <WiFiClientSecure.h>
#include <HTTPClient.h>
#include <SPI.h>
#include <TinyGPSPlus.h>
#include <ArduinoJson.h>
#include <Adafruit_GFX.h>
#include <Adafruit_ST7735.h>

// ===================== configure this ======================================

static const char *AP_SSID   = "your-hotspot";
static const char *AP_PASS   = "your-hotspot-password";
static const char *API_URL   = "https://wardrive.thugs.red/api/v1/ingest";
static const char *API_TOKEN = "wdrv_your_token_here";   // from /tokens

static const uint16_t BUFFER_SIZE = 8000;   // ~480 KB in PSRAM
static const uint16_t BATCH_SIZE  = 200;
static const float    MOVE_GATE_M = 15.0f;  // metres before a scan counts as new ground

static const uint32_t SCAN_MS_PER_CHANNEL = 300;
static const uint32_t SCAN_INTERVAL_MS    = 1500;

// ===================== pins ================================================
// The four pads on the dongle's JST SH connector are 3V3, GND, GPIO43, GPIO44.

#define GPS_RX_PIN 44   // GPS TX goes here
#define GPS_TX_PIN 43   // goes to GPS RX

#define BUTTON_PIN 0    // the dongle's button, and the S3 boot strapping pin

#define LED_DATA_PIN  40  // APA102 data
#define LED_CLOCK_PIN 39  // APA102 clock

#define TFT_CS   4
#define TFT_DC   2
#define TFT_RST  1
#define TFT_SCLK 5
#define TFT_MOSI 3
#define TFT_BL  38        // backlight. If your unit stays dark, invert BL_ON.

static const bool BL_ON = HIGH;
static const uint8_t LED_BRIGHTNESS = 6;   // APA102 global brightness, 1..31

// ===================== state ==============================================

struct Observation {
  uint8_t  bssid[6];
  char     ssid[33];
  int8_t   rssi;
  uint8_t  channel;
  uint8_t  enc;       // index into ENC_NAMES
  float    lat;
  float    lon;
  int16_t  alt;
  uint32_t ts;        // unix seconds, 0 = let the server use its own clock
};

static const char *ENC_NAMES[] = {
  "unknown", "open", "wep", "wpa", "wpa2", "wpa3", "wpa2_wpa3"
};

static Observation *buffer   = nullptr;
static uint16_t     capacity = 0;
static uint16_t     buffered = 0;
static uint32_t     dropped  = 0;

static const uint16_t SEEN_SLOTS = 4096;
static uint32_t      *seen = nullptr;
static uint16_t       seenCount = 0;

static TinyGPSPlus gps;
static double lastKeptLat = 0, lastKeptLon = 0;
static bool   hasKeptFix  = false;

static uint16_t lastFound  = 0;
static uint32_t lastScanAt = 0;

static Adafruit_ST7735 tft(TFT_CS, TFT_DC, TFT_RST);

enum RigState { RIG_WAIT_FIX, RIG_WARDRIVE, RIG_BUFFER_HIGH, RIG_UPLOAD, RIG_OK, RIG_FAIL };
static RigState rigState = RIG_WAIT_FIX;
static uint32_t rigStateSince = 0;

static char lastHttp[24] = "-";
static uint32_t lastAccepted = 0, lastRejected = 0;

// ===================== RGB LED ============================================
//
// The APA102 wants a start frame, one 32-bit frame per LED and an end frame.
// Bit-banging it is a dozen lines and saves pulling in a library for one LED.

static void apaByte(uint8_t value) {
  for (int8_t bit = 7; bit >= 0; bit--) {
    digitalWrite(LED_DATA_PIN, (value >> bit) & 1);
    digitalWrite(LED_CLOCK_PIN, HIGH);
    digitalWrite(LED_CLOCK_PIN, LOW);
  }
}

static void ledColour(uint8_t r, uint8_t g, uint8_t b) {
  apaByte(0x00); apaByte(0x00); apaByte(0x00); apaByte(0x00);   // start frame
  apaByte(0xE0 | (LED_BRIGHTNESS & 0x1F));                      // 111 + brightness
  apaByte(b); apaByte(g); apaByte(r);                           // APA102 order is BGR
  apaByte(0xFF); apaByte(0xFF); apaByte(0xFF); apaByte(0xFF);   // end frame
}

static void ledOff() { ledColour(0, 0, 0); }

// Colour says what mode it is in; the blink rate says how it is doing. Both,
// because a colour alone is easy to misremember in a moving car.
struct LedLook { uint16_t slotMs; uint8_t slots; uint32_t mask; uint8_t r, g, b; };

static LedLook lookFor(RigState state) {
  switch (state) {
    case RIG_WAIT_FIX:    return { 100, 20, 0b1,      255, 140,   0 };  // amber, one blink
    case RIG_WARDRIVE:    return { 100, 20, 0b101,      0, 255,  90 };  // green, two blinks
    case RIG_BUFFER_HIGH: return { 100, 20, 0b10101, 255, 140,   0 };  // amber, three blinks
    case RIG_UPLOAD:      return { 100,  2, 0b01,      40, 120, 255 };  // blue, 5 Hz
    case RIG_OK:          return { 100, 20, 0xFFFFF,   0, 255,  90 };  // green, solid
    case RIG_FAIL:        return {  60, 10, 0b10101, 255,  30,  40 };  // red, stutter
  }
  return { 100, 20, 0b1, 255, 255, 255 };
}

static void setState(RigState state) {
  if (state != rigState) {
    rigState = state;
    rigStateSince = millis();
  }
}

static void updateLed() {
  LedLook look = lookFor(rigState);
  uint32_t slot = ((millis() - rigStateSince) / look.slotMs) % look.slots;
  if ((look.mask >> slot) & 1) ledColour(look.r, look.g, look.b);
  else                         ledOff();
}

static void ledHold(RigState state, uint32_t ms) {
  setState(state);
  uint32_t until = millis() + ms;
  while (millis() < until) {
    updateLed();
    delay(5);
  }
}

// ===================== screen =============================================

static const uint16_t INK   = ST77XX_WHITE;
static const uint16_t PAPER = ST77XX_BLACK;
static const uint16_t BRAND = ST77XX_RED;

static void screenBegin() {
  pinMode(TFT_BL, OUTPUT);
  digitalWrite(TFT_BL, BL_ON);

  SPI.begin(TFT_SCLK, -1, TFT_MOSI, TFT_CS);
  tft.initR(INITR_MINI160x80);
  tft.setRotation(1);            // landscape: 160 x 80
  tft.fillScreen(PAPER);
  tft.setTextWrap(false);
}

static void screenHeader(const char *mode) {
  tft.fillRect(0, 0, 160, 12, INK);
  tft.setTextColor(PAPER);
  tft.setTextSize(1);
  tft.setCursor(2, 2);
  tft.print("THUGS");
  tft.setTextColor(BRAND, INK);
  tft.print("(red)");
  tft.setTextColor(PAPER, INK);
  tft.print(" ");
  tft.print(mode);
}

static void screenLine(uint8_t row, const char *text) {
  const int16_t y = 16 + row * 11;
  tft.fillRect(0, y, 160, 10, PAPER);
  tft.setTextColor(INK);
  tft.setCursor(2, y);
  tft.print(text);
}

static void screenWardrive() {
  char line[28];
  screenHeader("WARDRIVE");

  if (gps.location.isValid()) {
    snprintf(line, sizeof(line), "FIX 3D  sat %d", (int) gps.satellites.value());
  } else {
    snprintf(line, sizeof(line), "no fix  sat %d", (int) gps.satellites.value());
  }
  screenLine(0, line);

  snprintf(line, sizeof(line), "last scan  %u", lastFound);
  screenLine(1, line);
  snprintf(line, sizeof(line), "buf  %u / %u", buffered, capacity);
  screenLine(2, line);
  snprintf(line, sizeof(line), "uniq %u", seenCount);
  screenLine(3, line);

  if (gps.location.isValid()) {
    snprintf(line, sizeof(line), "%.5f %.5f", gps.location.lat(), gps.location.lng());
    screenLine(4, line);
  }
}

static void screenUpload(uint16_t batch, uint16_t batches) {
  char line[28];
  screenHeader("UPLOAD");

  snprintf(line, sizeof(line), "AP  %.18s", AP_SSID);
  screenLine(0, line);
  snprintf(line, sizeof(line), "batch  %u / %u", batch, batches);
  screenLine(1, line);
  snprintf(line, sizeof(line), "acc %lu rej %lu", (unsigned long) lastAccepted, (unsigned long) lastRejected);
  screenLine(2, line);
  snprintf(line, sizeof(line), "HTTP %s", lastHttp);
  screenLine(3, line);
  screenLine(4, "");
}

static void screenBanner(const char *mode, const char *text) {
  screenHeader(mode);
  screenLine(0, text);
  for (uint8_t row = 1; row < 5; row++) screenLine(row, "");
}

// ===================== helpers ============================================

static float metresBetween(double lat1, double lon1, double lat2, double lon2) {
  const double R = 6371000.0;
  double latRad = radians((lat1 + lat2) / 2.0);
  double dx = radians(lon2 - lon1) * cos(latRad);
  double dy = radians(lat2 - lat1);
  return (float) (R * sqrt(dx * dx + dy * dy));
}

// GPS date/time to unix seconds, written out rather than trusting mktime()'s
// timezone handling.
static uint32_t gpsEpoch() {
  if (!gps.date.isValid() || !gps.time.isValid() || gps.date.year() < 2020) {
    return 0;
  }
  int32_t y = gps.date.year();
  int32_t m = gps.date.month();
  int32_t d = gps.date.day();

  y -= m <= 2;
  int32_t era = (y >= 0 ? y : y - 399) / 400;
  uint32_t yoe = (uint32_t) (y - era * 400);
  uint32_t doy = (153 * (m + (m > 2 ? -3 : 9)) + 2) / 5 + d - 1;
  uint32_t doe = yoe * 365 + yoe / 4 - yoe / 100 + doy;
  int32_t days = era * 146097 + (int32_t) doe - 719468;

  return (uint32_t) days * 86400u
       + gps.time.hour() * 3600u + gps.time.minute() * 60u + gps.time.second();
}

static uint8_t encIndex(wifi_auth_mode_t mode) {
  switch (mode) {
    case WIFI_AUTH_OPEN:           return 1;
    case WIFI_AUTH_WEP:            return 2;
    case WIFI_AUTH_WPA_PSK:        return 3;
    case WIFI_AUTH_WPA2_PSK:
    case WIFI_AUTH_WPA_WPA2_PSK:   return 4;
    case WIFI_AUTH_WPA3_PSK:       return 5;
    case WIFI_AUTH_WPA2_WPA3_PSK:  return 6;
    default:                       return 0;
  }
}

static uint16_t frequencyFor(uint8_t channel) {
  if (channel == 14) return 2484;
  return 2407 + 5 * channel;
}

static uint32_t bssidHash(const uint8_t *b) {
  uint32_t h = 2166136261u;
  for (int i = 0; i < 6; i++) { h ^= b[i]; h *= 16777619u; }
  return h ? h : 1;
}

static void seenAdd(const uint8_t *bssid) {
  uint32_t h = bssidHash(bssid);
  for (uint16_t probe = 0; probe < 64; probe++) {
    uint16_t slot = (h + probe) & (SEEN_SLOTS - 1);
    if (seen[slot] == h) return;
    if (seen[slot] == 0) { seen[slot] = h; seenCount++; return; }
  }
}

static void bufferReset() {
  buffered = 0;
  dropped = 0;
  seenCount = 0;
  memset(seen, 0, SEEN_SLOTS * sizeof(uint32_t));
  hasKeptFix = false;
}

// ===================== wardrive mode ======================================

static void feedGps(uint32_t ms) {
  uint32_t until = millis() + ms;
  do {
    while (Serial1.available()) gps.encode(Serial1.read());
    updateLed();
  } while (millis() < until);
}

static void scanCycle() {
  if (!gps.location.isValid()) {
    setState(RIG_WAIT_FIX);
    screenWardrive();
    lastScanAt = millis();
    return;
  }

  double lat = gps.location.lat();
  double lon = gps.location.lng();

  if (hasKeptFix && metresBetween(lastKeptLat, lastKeptLon, lat, lon) < MOVE_GATE_M) {
    lastScanAt = millis();
    return;                       // parked; the same corner again is not data
  }

  uint32_t ts = gpsEpoch();
  int found = WiFi.scanNetworks(false, true, false, SCAN_MS_PER_CHANNEL);
  if (found < 0) found = 0;
  lastFound = (uint16_t) found;

  for (int i = 0; i < found; i++) {
    if (buffered >= capacity) { dropped++; break; }

    Observation &o = buffer[buffered];
    memcpy(o.bssid, WiFi.BSSID(i), 6);
    strncpy(o.ssid, WiFi.SSID(i).c_str(), sizeof(o.ssid) - 1);
    o.ssid[sizeof(o.ssid) - 1] = '\0';
    o.rssi    = (int8_t) WiFi.RSSI(i);
    o.channel = (uint8_t) WiFi.channel(i);
    o.enc     = encIndex(WiFi.encryptionType(i));
    o.lat     = (float) lat;
    o.lon     = (float) lon;
    o.alt     = (int16_t) (gps.altitude.isValid() ? gps.altitude.meters() : 0);
    o.ts      = ts;

    seenAdd(o.bssid);
    buffered++;
  }

  WiFi.scanDelete();

  lastKeptLat = lat;
  lastKeptLon = lon;
  hasKeptFix = true;
  lastScanAt = millis();

  setState(buffered > (uint32_t) capacity * 9 / 10 ? RIG_BUFFER_HIGH : RIG_WARDRIVE);
  screenWardrive();

  Serial.printf("[scan] found=%d buffered=%u/%u uniq=%u dropped=%lu\n",
                found, buffered, capacity, seenCount, (unsigned long) dropped);
}

// ===================== upload mode ========================================

enum PostResult { POST_SENT, POST_RETRY, POST_ABORT };

static PostResult postRange(uint16_t from, uint16_t count) {
  JsonDocument doc;
  JsonArray arr = doc["observations"].to<JsonArray>();

  for (uint16_t i = from; i < from + count; i++) {
    const Observation &o = buffer[i];
    char mac[18];
    snprintf(mac, sizeof(mac), "%02x:%02x:%02x:%02x:%02x:%02x",
             o.bssid[0], o.bssid[1], o.bssid[2], o.bssid[3], o.bssid[4], o.bssid[5]);

    JsonObject r = arr.add<JsonObject>();
    r["bssid"] = mac;
    if (o.ssid[0] != '\0') r["ssid"] = o.ssid;
    r["enc"]  = ENC_NAMES[o.enc];
    r["rssi"] = o.rssi;
    r["ch"]   = o.channel;
    r["freq"] = frequencyFor(o.channel);
    r["lat"]  = serialized(String(o.lat, 6));
    r["lon"]  = serialized(String(o.lon, 6));
    if (o.alt != 0) r["alt"] = o.alt;
    if (o.ts != 0)  r["ts"]  = o.ts;
  }

  String body;
  serializeJson(doc, body);

  WiFiClientSecure client;
  // Not verifying the certificate: this rig has no clock and no CA bundle. The
  // token is the only secret on the wire and it is revocable from /tokens. Pin
  // the root with setCACert() instead if you would rather not make that trade.
  client.setInsecure();

  HTTPClient http;
  http.setTimeout(20000);
  if (!http.begin(client, API_URL)) return POST_RETRY;

  const char *wanted[] = { "Retry-After" };
  http.collectHeaders(wanted, 1);
  http.addHeader("Content-Type", "application/json");
  http.addHeader("Authorization", String("Bearer ") + API_TOKEN);
  http.addHeader("User-Agent", "thugs-wardrive-tdongle/1.0");

  int code = http.POST(body);
  snprintf(lastHttp, sizeof(lastHttp), "%d", code);
  String payload = http.getString();
  String retryAfter = http.header("Retry-After");
  http.end();

  Serial.printf("[post] %u records -> HTTP %d %s\n", count, code, payload.c_str());

  if (code == 200 || code == 202) {
    JsonDocument res;
    if (deserializeJson(res, payload) == DeserializationError::Ok) {
      lastAccepted += res["accepted"] | 0;
      lastRejected += res["rejected"] | 0;
    }
    return POST_SENT;
  }

  if (code == 429) {
    uint32_t wait = retryAfter.toInt();
    if (wait == 0 || wait > 120) wait = 60;
    Serial.printf("[post] rate limited, waiting %us\n", wait);
    uint32_t until = millis() + wait * 1000;
    while (millis() < until) { updateLed(); delay(10); }
    return POST_RETRY;
  }

  if (code == 401 || code == 422) return POST_ABORT;   // token or batch is wrong
  if (code == 503) return POST_ABORT;                  // ingest off, or maintenance
  return POST_RETRY;
}

static bool uploadBuffer() {
  if (buffered == 0) return true;

  setState(RIG_UPLOAD);
  lastAccepted = lastRejected = 0;
  strncpy(lastHttp, "...", sizeof(lastHttp));

  WiFi.scanDelete();
  WiFi.mode(WIFI_STA);
  WiFi.begin(AP_SSID, AP_PASS);

  screenBanner("UPLOAD", "joining AP...");
  uint32_t until = millis() + 20000;
  while (WiFi.status() != WL_CONNECTED && millis() < until) { updateLed(); delay(50); }

  if (WiFi.status() != WL_CONNECTED) {
    Serial.println("[upload] no AP; keeping the buffer");
    return false;
  }

  uint16_t batches = (buffered + BATCH_SIZE - 1) / BATCH_SIZE;
  uint16_t sentTo = 0;
  uint16_t batchNo = 0;

  while (sentTo < buffered) {
    uint16_t count = min<uint16_t>(BATCH_SIZE, buffered - sentTo);
    batchNo++;
    screenUpload(batchNo, batches);

    PostResult result = POST_RETRY;
    for (uint8_t attempt = 0; attempt < 4 && result == POST_RETRY; attempt++) {
      if (attempt > 0) {
        uint32_t wake = millis() + (1000u << attempt);   // 2s, 4s, 8s
        while (millis() < wake) { updateLed(); delay(10); }
      }
      result = postRange(sentTo, count);
      if (result == POST_RETRY && count > 1) count /= 2;
    }

    if (result != POST_SENT) {
      Serial.println("[upload] giving up; keeping what is left");
      return false;
    }
    sentTo += count;
  }

  return true;
}

static bool buttonPressed() {
  if (digitalRead(BUTTON_PIN) != LOW) return false;
  delay(80);
  if (digitalRead(BUTTON_PIN) != LOW) return false;
  while (digitalRead(BUTTON_PIN) == LOW) { updateLed(); delay(10); }
  return true;
}

static void enterWardriveMode() {
  WiFi.disconnect(true, true);
  WiFi.mode(WIFI_STA);
  lastScanAt = 0;
  setState(gps.location.isValid() ? RIG_WARDRIVE : RIG_WAIT_FIX);
  screenWardrive();
}

// ===================== setup / loop =======================================

void setup() {
  pinMode(LED_DATA_PIN, OUTPUT);
  pinMode(LED_CLOCK_PIN, OUTPUT);
  ledOff();
  pinMode(BUTTON_PIN, INPUT_PULLUP);

  Serial.begin(115200);                  // native USB, because UART0 is the GPS
  Serial1.begin(9600, SERIAL_8N1, GPS_RX_PIN, GPS_TX_PIN);

  screenBegin();
  screenBanner("WARDRIVE", "booting...");

  capacity = BUFFER_SIZE;
  buffer = (Observation *) ps_malloc(sizeof(Observation) * capacity);
  if (buffer == nullptr) {
    // PSRAM missing or not enabled in the board menu. Shrink rather than fail;
    // every bounds check reads `capacity`, never BUFFER_SIZE.
    capacity = 1500;
    Serial.println("[boot] no PSRAM; falling back to a smaller buffer");
    buffer = (Observation *) malloc(sizeof(Observation) * capacity);
  }
  seen = (uint32_t *) calloc(SEEN_SLOTS, sizeof(uint32_t));

  if (buffer == nullptr || seen == nullptr) {
    Serial.println("[boot] out of memory; halting");
    screenBanner("WARDRIVE", "out of memory");
    while (true) { ledHold(RIG_FAIL, 1000); }
  }

  Serial.println("THUGS(red) Wardrive - T-Dongle-S3");
  Serial.printf("[boot] buffer=%u records\n", capacity);

  enterWardriveMode();
}

void loop() {
  feedGps(200);

  if (buttonPressed()) {
    Serial.printf("[button] upload requested with %u buffered\n", buffered);
    bool ok = uploadBuffer();

    if (ok) {
      bufferReset();                     // clean slate: no records, no BSSIDs, no fix
      screenBanner("UPLOAD", "uploaded. clean.");
      ledHold(RIG_OK, 2000);
    } else {
      screenBanner("UPLOAD", "failed - kept buffer");
      ledHold(RIG_FAIL, 2000);
    }

    enterWardriveMode();
    return;
  }

  if (millis() - lastScanAt >= SCAN_INTERVAL_MS) {
    scanCycle();
  }
}