/*
 * uwb_anchor — SaltyBot ESP32 UWB Pro anchor (DW1000 TWR responder)
 *
 * Hardware: Makerfabs ESP32 UWB Pro (DW1000, 200m range)
 * Sits on robot body, USB-serial to Jetson Orin.
 * Two anchors: anchor-0 (port), anchor-1 (starboard).
 *
 * Serial output to Jetson (115200):
 *   +RANGE:<anchor_id>,<tag_short_addr>,<range_mm>,<rssi_dbm>\r\n
 *
 * Also receives ESP-NOW packets from tag and forwards:
 *   +ESPNOW:<tag_id>,<msg_type>,<range_mm>,<rssi>,<batt>,<flags>,<seq>\r\n
 *   +ESTOP:<tag_id>\r\n  (priority)
 *
 * AT commands (host → anchor):
 *   AT+ID?    → +ID:<anchor_id>
 *
 * Pin map — ESP32 UWB Pro (no display):
 *   SPI: SCK=18 MISO=19 MOSI=23 CS=4
 *   DW1000: RST=27 IRQ=34
 */

#include <Arduino.h>
#include <SPI.h>
#include <string.h>
#include <WiFi.h>
#include <esp_now.h>
#include <esp_wifi.h>
#include "DW1000Ranging.h"

/* ── Config ─────────────────────────────────────────────────────── */

#ifndef ANCHOR_ID
#  define ANCHOR_ID  0
#endif

#define SERIAL_BAUD  115200

/* Unique DW1000 address per anchor (last 2 bytes differ) */
#if ANCHOR_ID == 0
static char ANCHOR_ADDR[] = "86:17:5B:D5:A9:9A:E2:00";
#else
static char ANCHOR_ADDR[] = "86:17:5B:D5:A9:9A:E2:01";
#endif

/* ── Pins (ESP32 UWB Pro, no display) ──────────────────────────── */

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

/* ── ESP-NOW packet format (shared with tag) ───────────────────── */

#define ESPNOW_MAGIC_0   0x5B
#define ESPNOW_MAGIC_1   0x01
#define MSG_RANGE         0x10
#define MSG_ESTOP         0x20
#define MSG_HEARTBEAT     0x30
#define MSG_IMU           0x40
#define MSG_FALL          0x50

#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;
};
#pragma pack(pop)

/* ISR ring buffer for ESP-NOW */
#define ESPNOW_Q_SIZE  8
static EspNowPacket g_enq[ESPNOW_Q_SIZE];
static volatile int g_en_head = 0, g_en_tail = 0;

static void IRAM_ATTR espnow_rx_cb(const esp_now_recv_info_t *info, const uint8_t *data, int len) {
    if (len < (int)sizeof(EspNowPacket)) return;
    const EspNowPacket *p = (const EspNowPacket *)data;
    if (p->magic[0] != ESPNOW_MAGIC_0 || p->magic[1] != ESPNOW_MAGIC_1) return;
    int next = (g_en_head + 1) % ESPNOW_Q_SIZE;
    if (next == g_en_tail) return;
    memcpy(&g_enq[g_en_head], p, sizeof(EspNowPacket));
    g_en_head = next;
}

/* ── ESP-NOW TX: broadcast own range to other anchor + tag ──────── */

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

/* Track the other anchor's range (received via ESP-NOW) */
static int32_t  g_other_range_mm = -1;
static float    g_other_rssi = -100.0f;
static uint32_t g_other_last_ms = 0;

static void espnow_send_range(int32_t range_mm, float rssi, uint16_t tag_addr) {
    EspNowPacket pkt = {};
    pkt.magic[0]     = ESPNOW_MAGIC_0;
    pkt.magic[1]     = ESPNOW_MAGIC_1;
    pkt.tag_id       = (uint8_t)(tag_addr & 0xFF);
    pkt.msg_type     = MSG_RANGE;
    pkt.anchor_id    = ANCHOR_ID;
    pkt.range_mm     = range_mm;
    pkt.rssi_dbm     = rssi;
    pkt.timestamp_ms = millis();
    pkt.battery_pct  = 0xFF;
    pkt.flags        = 0x00;
    pkt.seq_num      = g_seq++;
    esp_now_send(broadcast_mac, (uint8_t *)&pkt, sizeof(pkt));
}

/* ── 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+ID?") == 0) {
        Serial.printf("+ID:%d\r\n", ANCHOR_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;
        }
    }
}

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

static void espnow_process(void) {
    while (g_en_tail != g_en_head) {
        const EspNowPacket &pkt = (const EspNowPacket &)g_enq[g_en_tail];
        g_en_tail = (g_en_tail + 1) % ESPNOW_Q_SIZE;

        if (pkt.msg_type == MSG_ESTOP) {
            if (pkt.flags & 0x01)
                Serial.printf("+ESTOP:%d\r\n", pkt.tag_id);
            else
                Serial.printf("+ESTOP_CLEAR:%d\r\n", pkt.tag_id);
            continue;
        }

        /* If this is a range from the OTHER anchor, track it and print as +RANGE */
        if (pkt.msg_type == MSG_RANGE && pkt.anchor_id != ANCHOR_ID
            && pkt.anchor_id < 2) {
            g_other_range_mm = pkt.range_mm;
            g_other_rssi     = pkt.rssi_dbm;
            g_other_last_ms  = millis();
            /* Print as +RANGE so Orin sees both anchors from either serial port */
            Serial.printf("+RANGE:%d,%02X,%ld,%.1f\r\n",
                          pkt.anchor_id, pkt.tag_id,
                          (long)pkt.range_mm, pkt.rssi_dbm);
            continue;
        }

        /* Forward IMU and FALL packets to Jetson serial */
        if (pkt.msg_type == MSG_IMU || pkt.msg_type == MSG_FALL) {
            const uint8_t *raw = (const uint8_t *)&pkt;
            Serial.printf("+%s:%d,",
                          pkt.msg_type == MSG_FALL ? "FALL" : "IMU",
                          pkt.tag_id);
            /* Forward as hex for easy parsing: ax,ay,az,gx,gy,gz,mag */
            int16_t ax = (int16_t)(raw[4] | (raw[5] << 8));
            int16_t ay = (int16_t)(raw[6] | (raw[7] << 8));
            int16_t az = (int16_t)(raw[8] | (raw[9] << 8));
            int16_t gx = (int16_t)(raw[10] | (raw[11] << 8));
            int16_t gy = (int16_t)(raw[12] | (raw[13] << 8));
            int16_t gz = (int16_t)(raw[14] | (raw[15] << 8));
            uint8_t mag = raw[16];
            Serial.printf("%.3f,%.3f,%.3f,%.1f,%.1f,%.1f,%.1f\r\n",
                          ax/1000.0f, ay/1000.0f, az/1000.0f,
                          gx/10.0f, gy/10.0f, gz/10.0f,
                          mag/10.0f);
            continue;
        }

        Serial.printf("+ESPNOW:%d,%02X,%ld,%.1f,%d,%02X,%d\r\n",
                      pkt.tag_id, pkt.msg_type, (long)pkt.range_mm,
                      pkt.rssi_dbm, pkt.battery_pct, pkt.flags, pkt.seq_num);
    }
}

/* ── DW1000Ranging callbacks ────────────────────────────────────── */

static void newRange(void) {
    uint16_t tag_addr = DW1000Ranging.getDistantDevice()->getShortAddress();
    float    range_m  = DW1000Ranging.getDistantDevice()->getRange();
    float    rxPow    = DW1000Ranging.getDistantDevice()->getRXPower();

    if (range_m < 0.0f || range_m > 250.0f) return;

    int32_t range_mm = (int32_t)(range_m * 1000.0f + 0.5f);

    Serial.printf("+RANGE:%d,%04X,%ld,%.1f\r\n",
                  ANCHOR_ID, tag_addr, (long)range_mm, rxPow);

    /* Broadcast own range via ESP-NOW → other anchor + tag receive it */
    espnow_send_range(range_mm, rxPow, tag_addr);
}

static void newBlink(DW1000Device *device) {
    Serial.printf("+BLINK:%04X\r\n", device->getShortAddress());
}

static void inactiveDevice(DW1000Device *device) {
    Serial.printf("+GONE:%04X\r\n", device->getShortAddress());
}

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

void setup(void) {
    Serial.begin(SERIAL_BAUD);
    delay(300);
    Serial.printf("\r\n[uwb_anchor] id=%d  addr=%s  starting\r\n", ANCHOR_ID, ANCHOR_ADDR);

    /* ESP-NOW receiver */
    WiFi.mode(WIFI_STA);
    WiFi.disconnect();
    esp_wifi_set_channel(1, WIFI_SECOND_CHAN_NONE);
    if (esp_now_init() == ESP_OK) {
        esp_now_register_recv_cb(espnow_rx_cb);
        /* Add broadcast peer for TX */
        esp_now_peer_info_t peer = {};
        memcpy(peer.peer_addr, broadcast_mac, 6);
        peer.channel = 0;
        peer.encrypt = false;
        esp_now_add_peer(&peer);
        Serial.println("[uwb_anchor] ESP-NOW rx+tx ok");
    } else {
        Serial.println("[uwb_anchor] WARN: ESP-NOW failed");
    }

    /* DW1000 */
    SPI.begin(PIN_SCK, PIN_MISO, PIN_MOSI);
    DW1000Ranging.initCommunication(PIN_RST, PIN_CS, PIN_IRQ);

    DW1000Ranging.attachNewRange(newRange);
    DW1000Ranging.attachBlinkDevice(newBlink);
    DW1000Ranging.attachInactiveDevice(inactiveDevice);

    DW1000Ranging.startAsAnchor(ANCHOR_ADDR, DW1000.MODE_LONGDATA_RANGE_LOWPOWER, false);

    Serial.printf("[uwb_anchor] DW1000 ready  id=%d  MODE_LONGDATA_RANGE_LOWPOWER\r\n", ANCHOR_ID);
    Serial.println("[uwb_anchor] Listening for tags...");
}

/* ── Loop ───────────────────────────────────────────────────────── */

/*
 * Time-division multiplexing: 2 anchors share the same channel.
 * Each anchor only responds during its time slot to avoid RF collisions.
 * Slot duration = 50ms → each anchor gets 10 Hz effective rate.
 * Anchor 0 responds during even slots, anchor 1 during odd slots.
 */
#define SLOT_MS  50
#define NUM_ANCHORS_TOTAL  2

void loop(void) {
    serial_poll();
    espnow_process();

    uint32_t slot = (millis() / SLOT_MS) % NUM_ANCHORS_TOTAL;
    if (slot == ANCHOR_ID) {
        DW1000Ranging.loop();
    }
}