// ============================================================================
// Band-bridge module — implementation
// ============================================================================

#include "band_bridge.h"
#include <cstring>

#ifdef PDL_ENABLE_BAND_BRIDGE

#ifndef UNIT_TEST
#include <Arduino.h>
#include "mesh_publish.h"
#endif

// ----------------------------------------------------------------------------
// Local SHA-256 / HMAC (matches drone-companion/stanag_decoder.cpp). We
// duplicate here rather than depend on mbedtls include paths because T1
// firmware is ESP-IDF 4 + Arduino and includes paths vary by SDK spin.
// ----------------------------------------------------------------------------
namespace {

struct Sha256Ctx {
    uint32_t state[8];
    uint64_t bitlen;
    uint8_t  buffer[64];
    uint32_t buf_len;
};

constexpr uint32_t K[64] = {
    0x428a2f98,0x71374491,0xb5c0fbcf,0xe9b5dba5,0x3956c25b,0x59f111f1,0x923f82a4,0xab1c5ed5,
    0xd807aa98,0x12835b01,0x243185be,0x550c7dc3,0x72be5d74,0x80deb1fe,0x9bdc06a7,0xc19bf174,
    0xe49b69c1,0xefbe4786,0x0fc19dc6,0x240ca1cc,0x2de92c6f,0x4a7484aa,0x5cb0a9dc,0x76f988da,
    0x983e5152,0xa831c66d,0xb00327c8,0xbf597fc7,0xc6e00bf3,0xd5a79147,0x06ca6351,0x14292967,
    0x27b70a85,0x2e1b2138,0x4d2c6dfc,0x53380d13,0x650a7354,0x766a0abb,0x81c2c92e,0x92722c85,
    0xa2bfe8a1,0xa81a664b,0xc24b8b70,0xc76c51a3,0xd192e819,0xd6990624,0xf40e3585,0x106aa070,
    0x19a4c116,0x1e376c08,0x2748774c,0x34b0bcb5,0x391c0cb3,0x4ed8aa4a,0x5b9cca4f,0x682e6ff3,
    0x748f82ee,0x78a5636f,0x84c87814,0x8cc70208,0x90befffa,0xa4506ceb,0xbef9a3f7,0xc67178f2
};

inline uint32_t rotr(uint32_t x, uint32_t n) { return (x >> n) | (x << (32 - n)); }

void sha256_init(Sha256Ctx& c) {
    c.state[0]=0x6a09e667; c.state[1]=0xbb67ae85; c.state[2]=0x3c6ef372; c.state[3]=0xa54ff53a;
    c.state[4]=0x510e527f; c.state[5]=0x9b05688c; c.state[6]=0x1f83d9ab; c.state[7]=0x5be0cd19;
    c.bitlen = 0; c.buf_len = 0;
}
void sha256_transform(Sha256Ctx& c, const uint8_t* d) {
    uint32_t w[64];
    for (int i = 0; i < 16; i++)
        w[i] = ((uint32_t)d[i*4]<<24)|((uint32_t)d[i*4+1]<<16)|((uint32_t)d[i*4+2]<<8)|(uint32_t)d[i*4+3];
    for (int i = 16; i < 64; i++) {
        uint32_t s0 = rotr(w[i-15],7) ^ rotr(w[i-15],18) ^ (w[i-15] >> 3);
        uint32_t s1 = rotr(w[i-2],17) ^ rotr(w[i-2],19) ^ (w[i-2] >> 10);
        w[i] = w[i-16] + s0 + w[i-7] + s1;
    }
    uint32_t a=c.state[0],b=c.state[1],cc=c.state[2],d2=c.state[3];
    uint32_t e=c.state[4],f=c.state[5],g=c.state[6],h=c.state[7];
    for (int i = 0; i < 64; i++) {
        uint32_t S1 = rotr(e,6) ^ rotr(e,11) ^ rotr(e,25);
        uint32_t ch = (e & f) ^ ((~e) & g);
        uint32_t t1 = h + S1 + ch + K[i] + w[i];
        uint32_t S0 = rotr(a,2) ^ rotr(a,13) ^ rotr(a,22);
        uint32_t mj = (a & b) ^ (a & cc) ^ (b & cc);
        uint32_t t2 = S0 + mj;
        h=g; g=f; f=e; e=d2+t1; d2=cc; cc=b; b=a; a=t1+t2;
    }
    c.state[0]+=a; c.state[1]+=b; c.state[2]+=cc; c.state[3]+=d2;
    c.state[4]+=e; c.state[5]+=f; c.state[6]+=g; c.state[7]+=h;
}
void sha256_update(Sha256Ctx& c, const uint8_t* d, size_t len) {
    for (size_t i = 0; i < len; i++) {
        c.buffer[c.buf_len++] = d[i];
        if (c.buf_len == 64) { sha256_transform(c, c.buffer); c.bitlen += 512; c.buf_len = 0; }
    }
}
void sha256_final(Sha256Ctx& c, uint8_t out[32]) {
    uint32_t i = c.buf_len;
    c.buffer[i++] = 0x80;
    if (i > 56) { while (i < 64) c.buffer[i++] = 0; sha256_transform(c, c.buffer); i = 0; }
    while (i < 56) c.buffer[i++] = 0;
    c.bitlen += (uint64_t)c.buf_len * 8;
    for (int j = 7; j >= 0; j--) c.buffer[i++] = (uint8_t)(c.bitlen >> (j * 8));
    sha256_transform(c, c.buffer);
    for (int j = 0; j < 8; j++) {
        out[j*4]=(uint8_t)(c.state[j]>>24); out[j*4+1]=(uint8_t)(c.state[j]>>16);
        out[j*4+2]=(uint8_t)(c.state[j]>>8); out[j*4+3]=(uint8_t)c.state[j];
    }
}
void hmac_sha256(const uint8_t* key, size_t key_len,
                 const uint8_t* msg, size_t msg_len, uint8_t out[32]) {
    uint8_t k[64] = {0};
    if (key_len > 64) {
        Sha256Ctx kc; sha256_init(kc);
        sha256_update(kc, key, key_len);
        sha256_final(kc, k);
    } else std::memcpy(k, key, key_len);
    uint8_t ipad[64], opad[64];
    for (int i = 0; i < 64; i++) { ipad[i] = k[i] ^ 0x36; opad[i] = k[i] ^ 0x5c; }
    uint8_t inner[32];
    Sha256Ctx c1; sha256_init(c1);
    sha256_update(c1, ipad, 64);
    sha256_update(c1, msg, msg_len);
    sha256_final(c1, inner);
    Sha256Ctx c2; sha256_init(c2);
    sha256_update(c2, opad, 64);
    sha256_update(c2, inner, 32);
    sha256_final(c2, out);
}

// ----- CRC-16-CCITT (for MESH_ACK shape) -----
uint16_t crc16_ccitt(const uint8_t* data, size_t len) {
    uint16_t crc = 0;
    for (size_t i = 0; i < len; i++) {
        crc ^= (uint16_t)data[i] << 8;
        for (int b = 0; b < 8; b++) crc = (crc & 0x8000) ? ((crc << 1) ^ 0x1021) : (crc << 1);
    }
    return crc;
}

// ----- Module state -----
uint8_t g_hmac_key[32] = {0};
size_t  g_hmac_key_len = 0;
bool    g_initialized  = false;

bool verify_stanag_hmac(const uint8_t* frame32) {
    if (g_hmac_key_len == 0) return false;
    uint8_t mac[32];
    hmac_sha256(g_hmac_key, g_hmac_key_len, frame32, 28, mac);
    return std::memcmp(mac, &frame32[28], 4) == 0;
}

uint32_t now_ms_local() {
#ifndef UNIT_TEST
    return millis();
#else
    static uint32_t t = 0; t += 10; return t;
#endif
}

} // namespace


// ----------------------------------------------------------------------------
// Public API
// ----------------------------------------------------------------------------

void band_bridge_init(const uint8_t* hmac_key, size_t hmac_key_len) {
    if (hmac_key && hmac_key_len > 0 && hmac_key_len <= sizeof(g_hmac_key)) {
        std::memcpy(g_hmac_key, hmac_key, hmac_key_len);
        g_hmac_key_len = hmac_key_len;
    } else {
        g_hmac_key_len = 0;
    }
    g_initialized = true;
}

BridgeResult band_bridge_relay(BandId src_band, BandId dst_band,
                               const uint8_t* frame, size_t len,
                               bool require_ack) {
    (void)src_band;
    BridgeResult r = {};
    r.sequence = 0;
    r.drone_rssi = -127;

    if (!g_initialized) { r.status = BridgeStatus::NOT_ENABLED; return r; }
    if (frame == nullptr) { r.status = BridgeStatus::HMAC_INVALID; return r; }

    const uint8_t* f32;
    size_t total;
    uint8_t buf_no_prefix[32];
    if (len == 32) {
        f32 = frame;
        total = 32;
    } else if (len == 33 && frame[0] == 0x0A) {
        std::memcpy(buf_no_prefix, frame + 1, 32);
        f32 = buf_no_prefix;
        total = 32;
    } else {
        r.status = BridgeStatus::HMAC_INVALID;
        return r;
    }
    r.sequence = ((uint16_t)f32[1] << 8) | (uint16_t)f32[2];

    // HMAC check (mandatory unless key_len==0, test-mode)
    if (g_hmac_key_len != 0 && !verify_stanag_hmac(f32)) {
        r.status = BridgeStatus::HMAC_INVALID;
        return r;
    }

    // Destination band must have a driver
    if ((size_t)dst_band >= (size_t)BandId::BAND_COUNT ||
        g_radios[(size_t)dst_band].tx == nullptr) {
        r.status = BridgeStatus::NO_TARGET_RADIO;
        return r;
    }

    // Duty-cycle budget (airtime ~ bytes * 1ms @ SF10BW125 approx)
    size_t airtime = total + 10;
    if (!radio_duty_ok(dst_band, airtime)) {
        r.status = BridgeStatus::DUTY_EXCEEDED;
        return r;
    }

    // Transmit on the destination band — include the 0x0A prefix so the
    // drone's parser (which accepts both 32/33-byte shapes) sees it
    // uniformly with how the LoRa mesh delivers.
    uint8_t tx_frame[33];
    tx_frame[0] = 0x0A;
    std::memcpy(&tx_frame[1], f32, 32);
    if (!radio_tx(dst_band, tx_frame, 33, 500)) {
        r.status = BridgeStatus::TX_FAILED;
        return r;
    }

    if (!require_ack) {
        r.status = BridgeStatus::DELIVERED;
        return r;
    }

    // Wait up to 500 ms for the drone's 16-byte ACK on the same band.
    uint32_t deadline = now_ms_local() + 500;
    uint8_t  ack_buf[32];
    size_t   ack_len = 0;
    RadioInterface& ri = g_radios[(size_t)dst_band];
    while (now_ms_local() < deadline) {
        if (ri.rx && ri.rx(ack_buf, &ack_len, 50) && ack_len >= 16) {
            if (ack_buf[0] == 0x0B) {
                // Parse sequence and compare
                uint16_t got = ((uint16_t)ack_buf[1] << 8) | (uint16_t)ack_buf[2];
                if (got == r.sequence) {
                    r.drone_rssi = ri.rssi_last ? ri.rssi_last() : -120;
                    r.status = BridgeStatus::DELIVERED;
                    return r;
                }
            }
        }
    }

    r.status = BridgeStatus::NO_ACK_FROM_DRONE;
    return r;
}

size_t band_bridge_build_mesh_ack(uint16_t sequence, BridgeStatus status,
                                  uint8_t* out, size_t out_len) {
    if (out_len < 16) return 0;
    std::memset(out, 0, 16);
    out[0] = PDL_MRSH_MSG_TYPE_MESH_ACK;
    out[1] = (uint8_t)((sequence >> 8) & 0xFF);
    out[2] = (uint8_t)(sequence & 0xFF);
    out[3] = (uint8_t)status;   // delivered / dropped / etc
    out[4] = 0;                 // result_code placeholder
    // [5..12] vehicle_hash (unknown at this layer; zero-filled)
    out[13] = 0x01;  // relayed flag (bit0)
    uint16_t crc = crc16_ccitt(out, 14);
    out[14] = (uint8_t)((crc >> 8) & 0xFF);
    out[15] = (uint8_t)(crc & 0xFF);
    return 16;
}

bool band_bridge_handle_route_delivery(const uint8_t* frame, size_t len,
                                       BandId drone_band) {
    BridgeResult r = band_bridge_relay(BandId::LORA_915, drone_band,
                                       frame, len, /*require_ack=*/true);
    // Emit a MESH_ACK back onto the 915 mesh for the originator
    uint8_t ack[16];
    size_t n = band_bridge_build_mesh_ack(r.sequence, r.status, ack, sizeof(ack));
#ifndef UNIT_TEST
    if (n > 0) mesh_publish_raw(ack, n);
#else
    (void)n;
#endif
    return r.status == BridgeStatus::DELIVERED;
}

#else  // !PDL_ENABLE_BAND_BRIDGE — compile-out stubs

void band_bridge_init(const uint8_t*, size_t) {}
BridgeResult band_bridge_relay(BandId, BandId, const uint8_t*, size_t, bool) {
    BridgeResult r = {};
    r.status = BridgeStatus::NOT_ENABLED;
    return r;
}
size_t band_bridge_build_mesh_ack(uint16_t, BridgeStatus, uint8_t*, size_t) { return 0; }
bool band_bridge_handle_route_delivery(const uint8_t*, size_t, BandId) { return false; }

#endif  // PDL_ENABLE_BAND_BRIDGE
