/*
 * uwb_anchor — SaltyBot ESP32 UWB Pro anchor firmware (TWR responder)
 * Issue #544
 *
 * Hardware: Makerfabs ESP32 UWB Pro (DW3000 chip)
 *
 * Role
 * ────
 *   Anchor sits on SaltyBot body, USB-connected to Jetson Orin.
 *   Two anchors per robot (anchor-0 port side, anchor-1 starboard).
 *   Person-worn tags initiate ranging; anchors respond.
 *
 * Protocol: Double-Sided TWR (DS-TWR)
 * ────────────────────────────────────
 *   Tag → Anchor   POLL   (msg_type 0x01)
 *   Anchor → Tag   RESP   (msg_type 0x02, payload: T_poll_rx, T_resp_tx)
 *   Tag → Anchor   FINAL  (msg_type 0x03, payload: Ra, Da, Db timestamps)
 *   Anchor computes range via DS-TWR formula, emits +RANGE on Serial.
 *
 * Serial output (115200 8N1, USB-CDC to Jetson)
 * ──────────────────────────────────────────────
 *   +RANGE:<anchor_id>,<range_mm>,<rssi_dbm>\r\n   (on each successful range)
 *
 * AT commands (host → anchor)
 * ───────────────────────────
 *   AT+RANGE?                  → returns last buffered +RANGE line
 *   AT+RANGE_ADDR=<hex_addr>   → pair with specific tag (filter others)
 *   AT+RANGE_ADDR=             → clear pairing (accept all tags)
 *   AT+ID?                     → returns +ID:<anchor_id>
 *   AT+PEER_RANGE=<id>         → inter-anchor DS-TWR (for auto-calibration)
 *                                → +PEER_RANGE:<my>,<peer>,<mm>,<rssi>
 *
 * Pin mapping — Makerfabs ESP32 UWB Pro
 * ──────────────────────────────────────
 *   SPI SCK   18    SPI MISO  19    SPI MOSI  23
 *   DW CS     21    DW RST    27    DW IRQ    34
 *
 * Build
 * ──────
 *   pio run -e anchor0 --target upload   (port side)
 *   pio run -e anchor1 --target upload   (starboard)
 */

#include <Arduino.h>
#include <SPI.h>
#include <math.h>
#include "dw3000.h"   // Makerfabs MaUWB_DW3000 library

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

#ifndef ANCHOR_ID
#  define ANCHOR_ID  0          /* 0 = port, 1 = starboard */
#endif

#define SERIAL_BAUD  115200

/* ── Pin map (Makerfabs ESP32 UWB Pro) ─────────────────────────── */

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

/* ── DW3000 channel / PHY config ───────────────────────────────── */

static dwt_config_t dw_cfg = {
    5,                 /* channel 5 (6.5 GHz, best penetration) */
    DWT_PLEN_128,      /* preamble length                       */
    DWT_PAC8,          /* PAC size                              */
    9,                 /* TX preamble code                      */
    9,                 /* RX preamble code                      */
    1,                 /* SFD type (IEEE 802.15.4z)             */
    DWT_BR_6M8,        /* data rate 6.8 Mbps                    */
    DWT_PHR_MODE_STD,  /* standard PHR                          */
    DWT_PHR_RATE_DATA,
    (129 + 8 - 8),     /* SFD timeout                           */
    DWT_STS_MODE_OFF,  /* STS off — standard TWR                */
    DWT_STS_LEN_64,
    DWT_PDOA_M0,       /* no PDoA                               */
};

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

/* Byte layout for all frames:
 *   [0]     frame_type  (FTYPE_*)
 *   [1]     src_id      (tag 8-bit addr, or ANCHOR_ID)
 *   [2]     dst_id
 *   [3..]   payload
 *   (FCS appended automatically by DW3000 — 2 bytes)
 */

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

#define FRAME_HDR   3
#define FCS_LEN     2

/* RESP payload: T_poll_rx(5 B) + T_resp_tx(5 B) */
#define RESP_PAYLOAD   10
#define RESP_FRAME_LEN (FRAME_HDR + RESP_PAYLOAD + FCS_LEN)

/* FINAL payload: Ra(5 B) + Da(5 B) + Db(5 B) */
#define FINAL_PAYLOAD   15
#define FINAL_FRAME_LEN (FRAME_HDR + FINAL_PAYLOAD + FCS_LEN)

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

/* Turnaround delay: anchor waits 500 µs after poll_rx before tx_resp.
 * DW3000 tick = 1/(128×499.2e6) ≈ 15.65 ps  →  500 µs = ~31.95M ticks.
 * Stored as uint32 shifted right 8 bits for dwt_setdelayedtrxtime. */
#define RESP_TX_DLY_US      500UL
#define DWT_TICKS_PER_US    63898UL          /* 1µs in DW3000 ticks (×8 prescaler) */
#define RESP_TX_DLY_TICKS   (RESP_TX_DLY_US * DWT_TICKS_PER_US)

/* How long anchor listens for FINAL after sending RESP */
#define FINAL_RX_TIMEOUT_US  3000

/* Speed of light (m/s) */
#define SPEED_OF_LIGHT  299702547.0

/* DW3000 40-bit timestamp mask */
#define DWT_TS_MASK  0xFFFFFFFFFFULL

/* Antenna delay (factory default; calibrate per unit for best accuracy) */
#define ANT_DELAY  16385

/* ── Interrupt flags (set in ISR, polled in main) ──────────────── */

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;

/* ── State ──────────────────────────────────────────────────────── */

/* Last successful range (serves AT+RANGE? queries) */
static int32_t g_last_range_mm = -1;
static char    g_last_range_line[72] = {};

/* Optional tag pairing: 0 = accept all tags */
static uint8_t g_paired_tag_id = 0;

/* ── DW3000 ISR callbacks ───────────────────────────────────────── */

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);
}

/* Estimate receive power from CIR diagnostics (dBm) */
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;
}

/* ── Peer-anchor ranging (initiator role, for auto-calibration) ─── */

/* Timeout waiting for peer's RESP during inter-anchor ranging */
#define PEER_RX_RESP_TIMEOUT_US   3500
/* Block up to this many ms waiting for the interrupt flags */
#define PEER_WAIT_MS              20

/* Initiate a single DS-TWR exchange toward peer anchor `peer_id`.
 * Returns range in mm (>=0) on success, or -1 on timeout/error.
 * RSSI is stored in *rssi_out if non-null.
 *
 * Exchange:
 *   This anchor  → peer  POLL
 *   peer         → This  RESP  (carries T_poll_rx, T_resp_tx)
 *   This anchor  → peer  FINAL (carries Ra, Da)
 *   This side computes its own range estimate from Ra/Da.
 */
static int32_t peer_range_once(uint8_t peer_id, float *rssi_out) {
    /* ── Reset interrupt flags ── */
    g_tx_done = g_rx_ok = g_rx_err = g_rx_to = false;

    /* ── Build POLL frame ── */
    uint8_t poll_buf[FRAME_HDR + FCS_LEN] = {
        FTYPE_POLL,
        (uint8_t)ANCHOR_ID,
        peer_id,
    };

    dwt_writetxdata(sizeof(poll_buf), poll_buf, 0);
    dwt_writetxfctrl(sizeof(poll_buf), 0, 1);
    dwt_setrxtimeout(PEER_RX_RESP_TIMEOUT_US);
    dwt_rxenable(DWT_START_RX_IMMEDIATE);
    if (dwt_starttx(DWT_START_TX_IMMEDIATE | DWT_RESPONSE_EXPECTED) != DWT_SUCCESS)
        return -1;

    /* Wait for TX done */
    uint32_t t0 = millis();
    while (!g_tx_done && !g_rx_err && !g_rx_to) {
        if (millis() - t0 > (uint32_t)PEER_WAIT_MS) return -1;
    }
    if (g_rx_err || g_rx_to) return -1;
    g_tx_done = false;

    /* Capture T_poll_tx */
    uint64_t T_poll_tx;
    dwt_readtxtimestamp((uint8_t *)&T_poll_tx);
    T_poll_tx &= DWT_TS_MASK;

    /* Wait for RESP */
    t0 = millis();
    while (!g_rx_ok && !g_rx_err && !g_rx_to) {
        if (millis() - t0 > (uint32_t)PEER_WAIT_MS) return -1;
    }
    if (!g_rx_ok) return -1;

    /* Validate RESP */
    if (g_rx_len < (uint32_t)(FRAME_HDR + RESP_PAYLOAD + FCS_LEN)) return -1;
    if (g_rx_buf[0] != FTYPE_RESP)              return -1;
    if (g_rx_buf[1] != peer_id)                  return -1;
    if (g_rx_buf[2] != (uint8_t)ANCHOR_ID)       return -1;

    uint64_t T_resp_rx;
    dwt_readrxtimestamp((uint8_t *)&T_resp_rx);
    T_resp_rx &= DWT_TS_MASK;

    /* Extract peer timestamps from RESP payload */
    const uint8_t *pl = g_rx_buf + FRAME_HDR;
    uint64_t T_poll_rx_peer = ts_read(pl);
    uint64_t T_resp_tx_peer = ts_read(pl + 5);

    /* DS-TWR Ra, Da */
    uint64_t Ra = ts_diff(T_resp_rx, T_poll_tx);
    uint64_t Da = ts_diff(T_resp_tx_peer, T_poll_rx_peer);

    g_rx_ok = g_rx_err = g_rx_to = false;

    /* ── Build FINAL frame ── */
    uint8_t final_buf[FRAME_HDR + FINAL_PAYLOAD + FCS_LEN];
    final_buf[0] = FTYPE_FINAL;
    final_buf[1] = (uint8_t)ANCHOR_ID;
    final_buf[2] = peer_id;
    ts_write(final_buf + FRAME_HDR,      Ra);
    ts_write(final_buf + FRAME_HDR + 5,  Da);
    ts_write(final_buf + FRAME_HDR + 10, (uint64_t)0);  /* Db placeholder */

    dwt_setrxtimeout(0);
    dwt_writetxdata(sizeof(final_buf), final_buf, 0);
    dwt_writetxfctrl(sizeof(final_buf), 0, 1);
    if (dwt_starttx(DWT_START_TX_IMMEDIATE) != DWT_SUCCESS) return -1;

    t0 = millis();
    while (!g_tx_done && !g_rx_err) {
        if (millis() - t0 > (uint32_t)PEER_WAIT_MS) return -1;
    }
    g_tx_done = false;

    /* Simplified one-sided range estimate: tof = Ra - Da/2 */
    double tof_s = ticks_to_s(Ra) - ticks_to_s(Da) / 2.0;
    if (tof_s < 0.0) tof_s = 0.0;
    int32_t range_mm = (int32_t)(tof_s * SPEED_OF_LIGHT * 1000.0);

    if (rssi_out) *rssi_out = rx_power_dbm();

    /* Re-enable normal RX for tag ranging */
    dwt_setrxtimeout(0);
    dwt_rxenable(DWT_START_RX_IMMEDIATE);

    return range_mm;
}

/* ── AT command handler ─────────────────────────────────────────── */

static char g_at_buf[64];
static int  g_at_idx = 0;

static void at_dispatch(const char *cmd) {
    if (strcmp(cmd, "AT+RANGE?") == 0) {
        if (g_last_range_mm >= 0)
            Serial.println(g_last_range_line);
        else
            Serial.println("+RANGE:NO_DATA");

    } else if (strcmp(cmd, "AT+ID?") == 0) {
        Serial.printf("+ID:%d\r\n", ANCHOR_ID);

    } else if (strncmp(cmd, "AT+RANGE_ADDR=", 14) == 0) {
        const char *v = cmd + 14;
        if (*v == '\0') {
            g_paired_tag_id = 0;
            Serial.println("+OK:UNPAIRED");
        } else {
            g_paired_tag_id = (uint8_t)strtoul(v, nullptr, 0);
            Serial.printf("+OK:PAIRED=0x%02X\r\n", g_paired_tag_id);
        }

    } else if (strncmp(cmd, "AT+PEER_RANGE=", 14) == 0) {
        /* Inter-anchor ranging for calibration.
         * Usage: AT+PEER_RANGE=<peer_anchor_id>
         * Response: +PEER_RANGE:<my_id>,<peer_id>,<range_mm>,<rssi_dbm>
         *       or: +PEER_RANGE:ERR,<peer_id>,TIMEOUT
         */
        uint8_t peer_id = (uint8_t)strtoul(cmd + 14, nullptr, 0);
        float rssi = 0.0f;
        int32_t mm = peer_range_once(peer_id, &rssi);
        if (mm >= 0) {
            Serial.printf("+PEER_RANGE:%d,%d,%ld,%.1f\r\n",
                          ANCHOR_ID, peer_id, mm, (double)rssi);
        } else {
            Serial.printf("+PEER_RANGE:ERR,%d,TIMEOUT\r\n", peer_id);
        }

    } else {
        Serial.println("+ERR:UNKNOWN_CMD");
    }
}

static void serial_poll(void) {
    while (Serial.available()) {
        char c = (char)Serial.read();
        if (c == '\r') continue;
        if (c == '\n') {
            g_at_buf[g_at_idx] = '\0';
            if (g_at_idx > 0) at_dispatch(g_at_buf);
            g_at_idx = 0;
        } else if (g_at_idx < (int)(sizeof(g_at_buf) - 1)) {
            g_at_buf[g_at_idx++] = c;
        }
    }
}

/* ── DS-TWR anchor state machine ────────────────────────────────── */

/*
 * DS-TWR responder (one shot):
 *   1. Wait for POLL from tag
 *   2. Delayed-TX RESP (carry T_poll_rx + scheduled T_resp_tx)
 *   3. Wait for FINAL from tag  (tag embeds Ra, Da, Db)
 *   4. Compute: Rb = T_final_rx − T_resp_tx
 *               tof = (Ra·Rb − Da·Db) / (Ra+Rb+Da+Db)
 *               range_m = tof × c
 *   5. Print +RANGE line
 */
static void twr_cycle(void) {

    /* --- 1. Listen for POLL --- */
    dwt_setrxtimeout(0);
    dwt_rxenable(DWT_START_RX_IMMEDIATE);

    g_rx_ok = g_rx_err = false;
    uint32_t deadline = millis() + 2000;
    while (!g_rx_ok && !g_rx_err) {
        serial_poll();
        if (millis() > deadline) {
            /* restart RX if we've been stuck */
            dwt_rxenable(DWT_START_RX_IMMEDIATE);
            deadline = millis() + 2000;
        }
        yield();
    }
    if (!g_rx_ok || g_rx_len < FRAME_HDR) return;

    /* validate POLL */
    if (g_rx_buf[0] != FTYPE_POLL) return;
    uint8_t tag_id = g_rx_buf[1];
    if (g_paired_tag_id != 0 && tag_id != g_paired_tag_id) return;

    /* --- 2. Record T_poll_rx --- */
    uint8_t  poll_rx_raw[5];
    dwt_readrxtimestamp(poll_rx_raw);
    uint64_t T_poll_rx = ts_read(poll_rx_raw);

    /* Compute delayed TX time: poll_rx + turnaround, aligned to 512-tick grid */
    uint64_t resp_tx_sched = (T_poll_rx + RESP_TX_DLY_TICKS) & ~0x1FFULL;

    /* Build RESP frame */
    uint8_t resp[RESP_FRAME_LEN];
    resp[0] = FTYPE_RESP;
    resp[1] = ANCHOR_ID;
    resp[2] = tag_id;
    ts_write(&resp[3], T_poll_rx);      /* T_poll_rx  (tag uses this) */
    ts_write(&resp[8], resp_tx_sched);  /* scheduled T_resp_tx */

    dwt_writetxdata(RESP_FRAME_LEN - FCS_LEN, resp, 0);
    dwt_writetxfctrl(RESP_FRAME_LEN, 0, 1 /*ranging*/);
    dwt_setdelayedtrxtime((uint32_t)(resp_tx_sched >> 8));

    /* Enable RX after TX to receive FINAL */
    dwt_setrxaftertxdelay(300);
    dwt_setrxtimeout(FINAL_RX_TIMEOUT_US);

    /* Fire delayed TX */
    g_tx_done = g_rx_ok = g_rx_err = g_rx_to = false;
    if (dwt_starttx(DWT_START_TX_DELAYED | DWT_RESPONSE_EXPECTED) != DWT_SUCCESS) {
        dwt_forcetrxoff();
        return; /* TX window missed — try next cycle */
    }

    /* Wait for TX done (short wait, ISR fires fast) */
    uint32_t t0 = millis();
    while (!g_tx_done && millis() - t0 < 15) { yield(); }

    /* Read actual T_resp_tx */
    uint8_t resp_tx_raw[5];
    dwt_readtxtimestamp(resp_tx_raw);
    uint64_t T_resp_tx = ts_read(resp_tx_raw);

    /* --- 3. Wait for FINAL --- */
    t0 = millis();
    while (!g_rx_ok && !g_rx_err && !g_rx_to && millis() - t0 < 60) {
        serial_poll();
        yield();
    }
    if (!g_rx_ok || g_rx_len < FRAME_HDR + FINAL_PAYLOAD) return;
    if (g_rx_buf[0] != FTYPE_FINAL) return;
    if (g_rx_buf[1] != tag_id)      return;

    /* Extract DS-TWR timestamps from FINAL payload */
    uint64_t Ra = ts_read(&g_rx_buf[3]);   /* tag: T_resp_rx − T_poll_tx */
    uint64_t Da = ts_read(&g_rx_buf[8]);   /* tag: T_final_tx − T_resp_rx */
    /* g_rx_buf[13..17] = Db from tag (cross-check, unused here) */

    /* T_final_rx */
    uint8_t final_rx_raw[5];
    dwt_readrxtimestamp(final_rx_raw);
    uint64_t T_final_rx = ts_read(final_rx_raw);

    /* --- 4. DS-TWR formula --- */
    uint64_t Rb = ts_diff(T_final_rx, T_resp_tx);  /* anchor round-trip */
    uint64_t Db = ts_diff(T_resp_tx,  T_poll_rx);  /* anchor turnaround */

    double ra = ticks_to_s(Ra), rb = ticks_to_s(Rb);
    double da = ticks_to_s(Da), db = ticks_to_s(Db);

    double denom = ra + rb + da + db;
    if (denom < 1e-15) return;

    double tof    = (ra * rb - da * db) / denom;
    double range_m = tof * SPEED_OF_LIGHT;

    /* Validity window: 0.1 m – 130 m */
    if (range_m < 0.1 || range_m > 130.0) return;

    int32_t range_mm = (int32_t)(range_m * 1000.0 + 0.5);
    float   rssi     = rx_power_dbm();

    /* --- 5. Emit +RANGE --- */
    snprintf(g_last_range_line, sizeof(g_last_range_line),
             "+RANGE:%d,%ld,%.1f", ANCHOR_ID, (long)range_mm, rssi);
    g_last_range_mm = range_mm;
    Serial.println(g_last_range_line);
}

/* ── Arduino setup ──────────────────────────────────────────────── */

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

    Serial.printf("\r\n[uwb_anchor] anchor_id=%d  starting\r\n", ANCHOR_ID);

    SPI.begin(PIN_SCK, PIN_MISO, PIN_MOSI, PIN_CS);

    /* Hardware reset */
    pinMode(PIN_RST, OUTPUT);
    digitalWrite(PIN_RST, LOW);
    delay(2);
    pinMode(PIN_RST, INPUT_PULLUP);
    delay(5);

    /* DW3000 probe + init (Makerfabs MaUWB_DW3000 library) */
    if (dwt_probe((struct dwt_probe_s *)&dw3000_probe_interf)) {
        Serial.println("[uwb_anchor] FATAL: DW3000 probe failed — check SPI wiring");
        for (;;) delay(1000);
    }

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

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

    dwt_setrxantennadelay(ANT_DELAY);
    dwt_settxantennadelay(ANT_DELAY);
    dwt_settxpower(0x0E080222UL);  /* max TX power for 120 m range */

    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);

    Serial.printf("[uwb_anchor] DW3000 ready  ch=%d  6.8Mbps  id=%d\r\n",
                  dw_cfg.chan, ANCHOR_ID);
    Serial.println("[uwb_anchor] Listening for tags...");
}

/* ── Arduino loop ───────────────────────────────────────────────── */

void loop(void) {
    serial_poll();
    twr_cycle();
}
