/*
 * uwb_tag — SaltyBot ESP32 UWB Pro tag firmware (DS-TWR initiator)
 * Issue #545 + display/ESP-NOW/e-stop extensions
 *
 * Hardware: Makerfabs ESP32 UWB Pro with Display (DW3000 + SSD1306 OLED)
 *
 * Role
 * ────
 *   Tag is worn by a person riding an EUC while SaltyBot follows.
 *   Initiates DS-TWR ranging with 2 anchors on the robot at 20 Hz.
 *   Shows distance/status on OLED.  Sends range data via ESP-NOW
 *   (no WiFi AP needed — peer-to-peer, ~1ms latency, works outdoors).
 *   GPIO 0 = emergency stop button (active low).
 *
 * Serial output (USB, 115200) — debug
 * ────────────────────────────────────
 *   +RANGE:<anchor_id>,<range_mm>,<rssi_dbm>\r\n
 *
 * ESP-NOW packet (broadcast, 20 bytes)
 * ─────────────────────────────────────
 *   [0-1]  magic  0x5B 0x01
 *   [2]    tag_id
 *   [3]    msg_type   0x10=range, 0x20=estop, 0x30=heartbeat
 *   [4]    anchor_id
 *   [5-8]  range_mm   (int32_t LE)
 *   [9-12] rssi_dbm   (float LE)
 *   [13-16] timestamp  (uint32_t millis)
 *   [17]   battery_pct (0-100 or 0xFF)
 *   [18]   flags       bit0=estop_active
 *   [19]   seq_num_lo  (uint8_t, rolling)
 *
 * Pin mapping — Makerfabs ESP32 UWB Pro with Display
 * ──────────────────────────────────────────────────
 *   SPI SCK  18   SPI MISO 19   SPI MOSI 23
 *   DW CS    21   DW RST   27   DW IRQ   34
 *   I2C SDA   4   I2C SCL   5   OLED addr 0x3C
 *   LED       2   E-STOP    0 (BOOT, active LOW)
 */

#include <Arduino.h>
#include <SPI.h>
#include <Wire.h>
#include <math.h>
#include <WiFi.h>
#include <esp_now.h>
#include <esp_wifi.h>

#include "dw3000.h"

#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>

/* ── Configurable ───────────────────────────────────────────────── */

#ifndef TAG_ID
#  define TAG_ID          0x01
#endif

#ifndef NUM_ANCHORS
#  define NUM_ANCHORS     2
#endif

#ifndef RANGE_INTERVAL_MS
#  define RANGE_INTERVAL_MS  50   /* 20 Hz round-robin */
#endif

#define SERIAL_BAUD       115200

/* ── Pins ───────────────────────────────────────────────────────── */

#define PIN_SCK   18
#define PIN_MISO  19
#define PIN_MOSI  23
#define PIN_CS    21
#define PIN_RST   27
#define PIN_IRQ   34

#define PIN_SDA    4
#define PIN_SCL    5

#define PIN_LED    2
#define PIN_ESTOP  0   /* BOOT button, active LOW */

/* ── OLED ───────────────────────────────────────────────────────── */

#define SCREEN_W  128
#define SCREEN_H   64
Adafruit_SSD1306 display(SCREEN_W, SCREEN_H, &Wire, -1);

/* ── ESP-NOW ────────────────────────────────────────────────────── */

#define ESPNOW_MAGIC_0   0x5B   /* "SB" */
#define ESPNOW_MAGIC_1   0x01   /* v1   */

#define MSG_RANGE         0x10
#define MSG_ESTOP         0x20
#define MSG_HEARTBEAT     0x30

static uint8_t broadcast_mac[6] = {0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF};
static uint8_t g_seq = 0;

#pragma pack(push, 1)
struct EspNowPacket {
    uint8_t  magic[2];
    uint8_t  tag_id;
    uint8_t  msg_type;
    uint8_t  anchor_id;
    int32_t  range_mm;
    float    rssi_dbm;
    uint32_t timestamp_ms;
    uint8_t  battery_pct;
    uint8_t  flags;
    uint8_t  seq_num;
    uint8_t  _pad;          /* pad to 20 bytes */
};
#pragma pack(pop)

static_assert(sizeof(EspNowPacket) == 20, "packet must be 20 bytes");

/* ── DW3000 PHY config (must match anchor) ──────────────────────── */

static dwt_config_t dw_cfg = {
    5,                 /* channel 5 */
    DWT_PLEN_128,
    DWT_PAC8,
    9, 9,              /* TX/RX preamble code */
    1,                 /* SFD type */
    DWT_BR_6M8,
    DWT_PHR_MODE_STD,
    DWT_PHR_RATE_DATA,
    (129 + 8 - 8),
    DWT_STS_MODE_OFF,
    DWT_STS_LEN_64,
    DWT_PDOA_M0,
};

/* ── Frame format ──────────────────────────────────────────────── */

#define FTYPE_POLL   0x01
#define FTYPE_RESP   0x02
#define FTYPE_FINAL  0x03

#define FRAME_HDR   3
#define FCS_LEN     2

#define POLL_FRAME_LEN   (FRAME_HDR + FCS_LEN)
#define RESP_PAYLOAD     10
#define RESP_FRAME_LEN   (FRAME_HDR + RESP_PAYLOAD + FCS_LEN)
#define FINAL_PAYLOAD    15
#define FINAL_FRAME_LEN  (FRAME_HDR + FINAL_PAYLOAD + FCS_LEN)

/* ── Timing ────────────────────────────────────────────────────── */

#define FINAL_TX_DLY_US     500UL
#define DWT_TICKS_PER_US    63898UL
#define FINAL_TX_DLY_TICKS  (FINAL_TX_DLY_US * DWT_TICKS_PER_US)
#define RESP_RX_TIMEOUT_US  3000

#define SPEED_OF_LIGHT  299702547.0
#define DWT_TS_MASK     0xFFFFFFFFFFULL
#define ANT_DELAY       16385

/* ── ISR state ──────────────────────────────────────────────────── */

static volatile bool g_rx_ok   = false;
static volatile bool g_tx_done = false;
static volatile bool g_rx_err  = false;
static volatile bool g_rx_to   = false;

static uint8_t  g_rx_buf[128];
static uint32_t g_rx_len = 0;

static void cb_tx_done(const dwt_cb_data_t *) { g_tx_done = true; }
static void cb_rx_ok(const dwt_cb_data_t *d) {
    g_rx_len = d->datalength;
    if (g_rx_len > sizeof(g_rx_buf)) g_rx_len = sizeof(g_rx_buf);
    dwt_readrxdata(g_rx_buf, g_rx_len, 0);
    g_rx_ok = true;
}
static void cb_rx_err(const dwt_cb_data_t *) { g_rx_err = true; }
static void cb_rx_to(const dwt_cb_data_t  *) { g_rx_to  = true; }

/* ── Timestamp helpers ──────────────────────────────────────────── */

static uint64_t ts_read(const uint8_t *p) {
    uint64_t v = 0;
    for (int i = 4; i >= 0; i--) v = (v << 8) | p[i];
    return v;
}

static void ts_write(uint8_t *p, uint64_t v) {
    for (int i = 0; i < 5; i++, v >>= 8) p[i] = (uint8_t)(v & 0xFF);
}

static inline uint64_t ts_diff(uint64_t later, uint64_t earlier) {
    return (later - earlier) & DWT_TS_MASK;
}

static inline double ticks_to_s(uint64_t t) {
    return (double)t / (128.0 * 499200000.0);
}

static float rx_power_dbm(void) {
    dwt_rxdiag_t d;
    dwt_readdiagnostics(&d);
    if (d.maxGrowthCIR == 0 || d.rxPreamCount == 0) return 0.0f;
    float f = (float)d.maxGrowthCIR;
    float n = (float)d.rxPreamCount;
    return 10.0f * log10f((f * f) / (n * n)) - 121.74f;
}

/* ── Shared state for display ───────────────────────────────────── */

static int32_t g_anchor_range_mm[NUM_ANCHORS];   /* last range per anchor */
static float   g_anchor_rssi[NUM_ANCHORS];        /* last RSSI per anchor */
static uint32_t g_anchor_last_ok[NUM_ANCHORS];    /* millis() of last good range */
static bool    g_estop_active = false;

/* ── ESP-NOW send helper ────────────────────────────────────────── */

static void espnow_send(uint8_t msg_type, uint8_t anchor_id,
                         int32_t range_mm, float rssi) {
    EspNowPacket pkt = {};
    pkt.magic[0]     = ESPNOW_MAGIC_0;
    pkt.magic[1]     = ESPNOW_MAGIC_1;
    pkt.tag_id       = TAG_ID;
    pkt.msg_type     = msg_type;
    pkt.anchor_id    = anchor_id;
    pkt.range_mm     = range_mm;
    pkt.rssi_dbm     = rssi;
    pkt.timestamp_ms = millis();
    pkt.battery_pct  = 0xFF;   /* TODO: read ADC battery voltage */
    pkt.flags        = g_estop_active ? 0x01 : 0x00;
    pkt.seq_num      = g_seq++;

    esp_now_send(broadcast_mac, (uint8_t *)&pkt, sizeof(pkt));
}

/* ── E-Stop handling ────────────────────────────────────────────── */

static uint32_t g_estop_last_tx = 0;

static void estop_check(void) {
    bool pressed = (digitalRead(PIN_ESTOP) == LOW);

    if (pressed && !g_estop_active) {
        /* Just pressed — enter e-stop */
        g_estop_active = true;
        Serial.println("+ESTOP:ACTIVE");
    }

    if (g_estop_active && pressed) {
        /* While held: send e-stop at 10 Hz */
        if (millis() - g_estop_last_tx >= 100) {
            espnow_send(MSG_ESTOP, 0xFF, 0, 0.0f);
            g_estop_last_tx = millis();
        }
    }

    if (!pressed && g_estop_active) {
        /* Released: send 3x clear packets, resume */
        for (int i = 0; i < 3; i++) {
            g_estop_active = false;  /* clear flag before sending so flags=0 */
            espnow_send(MSG_ESTOP, 0xFF, 0, 0.0f);
            delay(10);
        }
        g_estop_active = false;
        Serial.println("+ESTOP:CLEAR");
    }
}

/* ── OLED display update (5 Hz) ─────────────────────────────────── */

static uint32_t g_display_last = 0;

static void display_update(void) {
    if (millis() - g_display_last < 200) return;
    g_display_last = millis();

    display.clearDisplay();

    if (g_estop_active) {
        /* Big E-STOP warning */
        display.setTextSize(3);
        display.setTextColor(SSD1306_WHITE);
        display.setCursor(10, 4);
        display.println(F("E-STOP"));
        display.setTextSize(1);
        display.setCursor(20, 48);
        display.println(F("RELEASE TO CLEAR"));
        display.display();
        return;
    }

    uint32_t now = millis();

    /* Find closest anchor */
    int32_t min_range = INT32_MAX;
    for (int i = 0; i < NUM_ANCHORS; i++) {
        if (g_anchor_range_mm[i] > 0 && g_anchor_range_mm[i] < min_range)
            min_range = g_anchor_range_mm[i];
    }

    /* Line 1: Big distance to nearest anchor */
    display.setTextSize(3);
    display.setTextColor(SSD1306_WHITE);
    display.setCursor(0, 0);
    if (min_range < INT32_MAX && min_range > 0) {
        float m = min_range / 1000.0f;
        if (m < 10.0f)
            display.printf("%.1fm", m);
        else
            display.printf("%.0fm", m);
    } else {
        display.println(F("---"));
    }

    /* Line 2: Both anchor ranges */
    display.setTextSize(1);
    display.setCursor(0, 30);
    for (int i = 0; i < NUM_ANCHORS && i < 2; i++) {
        if (g_anchor_range_mm[i] > 0) {
            float m = g_anchor_range_mm[i] / 1000.0f;
            display.printf("A%d:%.1fm ", i, m);
        } else {
            display.printf("A%d:--- ", i);
        }
    }

    /* Line 3: Connection status */
    display.setCursor(0, 42);
    bool any_linked = false;
    for (int i = 0; i < NUM_ANCHORS; i++) {
        if (g_anchor_last_ok[i] > 0 && (now - g_anchor_last_ok[i]) < 2000) {
            any_linked = true;
            break;
        }
    }

    if (any_linked) {
        /* RSSI bar: map -90..-30 dBm to 0-5 bars */
        float best_rssi = -100.0f;
        for (int i = 0; i < NUM_ANCHORS; i++) {
            if (g_anchor_rssi[i] > best_rssi) best_rssi = g_anchor_rssi[i];
        }
        int bars = constrain((int)((best_rssi + 90.0f) / 12.0f), 0, 5);

        display.print(F("LINKED "));
        /* Draw signal bars */
        for (int b = 0; b < 5; b++) {
            int x = 50 + b * 6;
            int h = 2 + b * 2;
            int y = 50 - h;
            if (b < bars)
                display.fillRect(x, y, 4, h, SSD1306_WHITE);
            else
                display.drawRect(x, y, 4, h, SSD1306_WHITE);
        }
        display.printf(" %.0fdB", best_rssi);
    } else {
        display.println(F("LOST  -- searching --"));
    }

    /* Line 4: Uptime */
    display.setCursor(0, 54);
    uint32_t secs = now / 1000;
    display.printf("UP %02d:%02d  seq:%d", secs / 60, secs % 60, g_seq);

    display.display();
}

/* ── DS-TWR initiator (one anchor, one cycle) ───────────────────── */

static int32_t twr_range_once(uint8_t anchor_id) {

    /* --- 1. TX POLL --- */
    uint8_t poll[POLL_FRAME_LEN];
    poll[0] = FTYPE_POLL;
    poll[1] = TAG_ID;
    poll[2] = anchor_id;

    dwt_writetxdata(POLL_FRAME_LEN - FCS_LEN, poll, 0);
    dwt_writetxfctrl(POLL_FRAME_LEN, 0, 1);

    dwt_setrxaftertxdelay(300);
    dwt_setrxtimeout(RESP_RX_TIMEOUT_US);

    g_tx_done = g_rx_ok = g_rx_err = g_rx_to = false;
    if (dwt_starttx(DWT_START_TX_IMMEDIATE | DWT_RESPONSE_EXPECTED) != DWT_SUCCESS)
        return -1;

    uint32_t t0 = millis();
    while (!g_tx_done && millis() - t0 < 15) yield();

    uint8_t poll_tx_raw[5];
    dwt_readtxtimestamp(poll_tx_raw);
    uint64_t T_poll_tx = ts_read(poll_tx_raw);

    /* --- 2. Wait for RESP --- */
    t0 = millis();
    while (!g_rx_ok && !g_rx_err && !g_rx_to && millis() - t0 < 60) yield();
    if (!g_rx_ok || g_rx_len < FRAME_HDR + RESP_PAYLOAD) return -1;
    if (g_rx_buf[0] != FTYPE_RESP)   return -1;
    if (g_rx_buf[2] != TAG_ID)       return -1;
    if (g_rx_buf[1] != anchor_id)    return -1;

    uint8_t resp_rx_raw[5];
    dwt_readrxtimestamp(resp_rx_raw);
    uint64_t T_resp_rx = ts_read(resp_rx_raw);

    uint64_t T_poll_rx_a = ts_read(&g_rx_buf[3]);
    uint64_t T_resp_tx_a = ts_read(&g_rx_buf[8]);

    /* --- 3. Compute DS-TWR values for FINAL --- */
    uint64_t Ra = ts_diff(T_resp_rx, T_poll_tx);
    uint64_t Db = ts_diff(T_resp_tx_a, T_poll_rx_a);

    uint64_t final_tx_sched = (T_resp_rx + FINAL_TX_DLY_TICKS) & ~0x1FFULL;
    uint64_t Da = ts_diff(final_tx_sched, T_resp_rx);

    /* --- 4. TX FINAL --- */
    uint8_t final_buf[FINAL_FRAME_LEN];
    final_buf[0] = FTYPE_FINAL;
    final_buf[1] = TAG_ID;
    final_buf[2] = anchor_id;
    ts_write(&final_buf[3],  Ra);
    ts_write(&final_buf[8],  Da);
    ts_write(&final_buf[13], Db);

    dwt_writetxdata(FINAL_FRAME_LEN - FCS_LEN, final_buf, 0);
    dwt_writetxfctrl(FINAL_FRAME_LEN, 0, 1);
    dwt_setdelayedtrxtime((uint32_t)(final_tx_sched >> 8));

    g_tx_done = false;
    if (dwt_starttx(DWT_START_TX_DELAYED) != DWT_SUCCESS) {
        dwt_forcetrxoff();
        return -1;
    }
    t0 = millis();
    while (!g_tx_done && millis() - t0 < 15) yield();

    /* --- 5. Local range estimate (debug) --- */
    uint8_t final_tx_raw[5];
    dwt_readtxtimestamp(final_tx_raw);
    /* uint64_t T_final_tx = ts_read(final_tx_raw); -- unused, tag uses SS estimate */

    double ra = ticks_to_s(Ra);
    double db = ticks_to_s(Db);
    double tof    = (ra - db) / 2.0;
    double range_m = tof * SPEED_OF_LIGHT;

    if (range_m < 0.1 || range_m > 130.0) return -1;
    return (int32_t)(range_m * 1000.0 + 0.5);
}

/* ── Setup ──────────────────────────────────────────────────────── */

void setup(void) {
    Serial.begin(SERIAL_BAUD);
    delay(300);

    /* E-Stop button */
    pinMode(PIN_ESTOP, INPUT_PULLUP);
    pinMode(PIN_LED, OUTPUT);
    digitalWrite(PIN_LED, LOW);

    Serial.printf("\r\n[uwb_tag] tag_id=0x%02X  num_anchors=%d  starting\r\n",
                  TAG_ID, NUM_ANCHORS);

    /* --- OLED init --- */
    Wire.begin(PIN_SDA, PIN_SCL);
    if (!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
        Serial.println("[uwb_tag] WARN: SSD1306 not found — running headless");
    } else {
        display.clearDisplay();
        display.setTextSize(2);
        display.setTextColor(SSD1306_WHITE);
        display.setCursor(0, 0);
        display.println(F("SaltyBot"));
        display.setTextSize(1);
        display.setCursor(0, 20);
        display.printf("Tag 0x%02X  v2.0", TAG_ID);
        display.setCursor(0, 35);
        display.println(F("DW3000 DS-TWR + ESP-NOW"));
        display.setCursor(0, 50);
        display.println(F("Initializing..."));
        display.display();
        Serial.println("[uwb_tag] OLED ok");
    }

    /* --- ESP-NOW init --- */
    WiFi.mode(WIFI_STA);
    WiFi.disconnect();
    /* Set WiFi channel to match anchors (default ch 1) */
    esp_wifi_set_channel(1, WIFI_SECOND_CHAN_NONE);

    if (esp_now_init() != ESP_OK) {
        Serial.println("[uwb_tag] FATAL: esp_now_init failed");
        for (;;) delay(1000);
    }

    /* Add broadcast peer */
    esp_now_peer_info_t peer = {};
    memcpy(peer.peer_addr, broadcast_mac, 6);
    peer.channel = 0;  /* use current channel */
    peer.encrypt = false;
    esp_now_add_peer(&peer);
    Serial.println("[uwb_tag] ESP-NOW ok");

    /* --- DW3000 init --- */
    SPI.begin(PIN_SCK, PIN_MISO, PIN_MOSI, PIN_CS);

    pinMode(PIN_RST, OUTPUT);
    digitalWrite(PIN_RST, LOW);
    delay(2);
    pinMode(PIN_RST, INPUT_PULLUP);
    delay(5);

    if (dwt_probe((struct dwt_probe_s *)&dw3000_probe_interf)) {
        Serial.println("[uwb_tag] FATAL: DW3000 probe failed");
        for (;;) delay(1000);
    }

    if (dwt_initialise(DWT_DW_INIT) != DWT_SUCCESS) {
        Serial.println("[uwb_tag] FATAL: dwt_initialise failed");
        for (;;) delay(1000);
    }

    if (dwt_configure(&dw_cfg) != DWT_SUCCESS) {
        Serial.println("[uwb_tag] FATAL: dwt_configure failed");
        for (;;) delay(1000);
    }

    dwt_setrxantennadelay(ANT_DELAY);
    dwt_settxantennadelay(ANT_DELAY);
    dwt_settxpower(0x0E080222UL);

    dwt_setcallbacks(cb_tx_done, cb_rx_ok, cb_rx_to, cb_rx_err,
                     nullptr, nullptr, nullptr);
    dwt_setinterrupt(
        DWT_INT_TXFRS | DWT_INT_RFCG | DWT_INT_RFTO |
        DWT_INT_RFSL  | DWT_INT_SFDT | DWT_INT_ARFE | DWT_INT_CPERR,
        0, DWT_ENABLE_INT_ONLY);

    attachInterrupt(digitalPinToInterrupt(PIN_IRQ),
                    []() { dwt_isr(); }, RISING);

    /* Init range state */
    for (int i = 0; i < NUM_ANCHORS; i++) {
        g_anchor_range_mm[i] = -1;
        g_anchor_rssi[i] = -100.0f;
        g_anchor_last_ok[i] = 0;
    }

    Serial.printf("[uwb_tag] DW3000 ready  ch=%d  6.8Mbps  tag=0x%02X\r\n",
                  dw_cfg.chan, TAG_ID);
    Serial.println("[uwb_tag] Ranging + ESP-NOW + display active");
}

/* ── Main loop ──────────────────────────────────────────────────── */

void loop(void) {
    static uint8_t  anchor_idx    = 0;
    static uint32_t last_range_ms = 0;
    static uint32_t last_hb_ms    = 0;

    /* E-Stop always has priority */
    estop_check();
    if (g_estop_active) {
        display_update();
        return;   /* skip ranging while e-stop active */
    }

    /* Heartbeat every 1 second */
    uint32_t now = millis();
    if (now - last_hb_ms >= 1000) {
        espnow_send(MSG_HEARTBEAT, 0xFF, 0, 0.0f);
        last_hb_ms = now;
    }

    /* Ranging at configured interval */
    if (now - last_range_ms >= RANGE_INTERVAL_MS) {
        last_range_ms = now;

        uint8_t anchor_id = anchor_idx % NUM_ANCHORS;
        int32_t range_mm  = twr_range_once(anchor_id);

        if (range_mm > 0) {
            float rssi = rx_power_dbm();

            /* Update shared state for display */
            g_anchor_range_mm[anchor_id] = range_mm;
            g_anchor_rssi[anchor_id]     = rssi;
            g_anchor_last_ok[anchor_id]  = now;

            /* Serial debug */
            Serial.printf("+RANGE:%d,%ld,%.1f\r\n",
                          anchor_id, (long)range_mm, rssi);

            /* ESP-NOW broadcast */
            espnow_send(MSG_RANGE, anchor_id, range_mm, rssi);

            /* LED blink */
            digitalWrite(PIN_LED, HIGH);
            delay(2);
            digitalWrite(PIN_LED, LOW);
        }

        anchor_idx++;
        if (anchor_idx >= NUM_ANCHORS) anchor_idx = 0;
    }

    /* Display at 5 Hz (non-blocking) */
    display_update();
}
