// ============================================================================
// reticulum_identity.cpp — see reticulum_identity.h for design notes.
// ============================================================================

#include "reticulum_identity.h"
#include "ota_signed.h"   // ed25519_keypair_from_seed, ed25519_sign, x25519_keypair_from_seed

#include <cstring>
#include <cstdio>

// ---------------------------------------------------------------------------
// SHA-256 (FIPS 180-4). Self-contained, stack-only -- same convention as
// ota_signed.cpp's embedded SHA-512 (no shared state, no dynamic alloc).
// ---------------------------------------------------------------------------
namespace {

struct Sha256Ctx {
    uint32_t state[8];
    uint8_t  buf[64];
    uint64_t total_len;
    size_t   buf_len;
};

static const uint32_t kSha256K[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,
};

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

static void sha256_compress(Sha256Ctx* ctx, const uint8_t block[64]) {
    uint32_t w[64];
    for (int i = 0; i < 16; i++) {
        w[i] = ((uint32_t)block[i*4] << 24) | ((uint32_t)block[i*4+1] << 16) |
               ((uint32_t)block[i*4+2] << 8) | (uint32_t)block[i*4+3];
    }
    for (int i = 16; i < 64; i++) {
        uint32_t s0 = rotr32(w[i-15], 7) ^ rotr32(w[i-15], 18) ^ (w[i-15] >> 3);
        uint32_t s1 = rotr32(w[i-2], 17) ^ rotr32(w[i-2], 19) ^ (w[i-2] >> 10);
        w[i] = w[i-16] + s0 + w[i-7] + s1;
    }

    uint32_t a = ctx->state[0], b = ctx->state[1], c = ctx->state[2], d = ctx->state[3];
    uint32_t e = ctx->state[4], f = ctx->state[5], g = ctx->state[6], h = ctx->state[7];

    for (int i = 0; i < 64; i++) {
        uint32_t S1 = rotr32(e, 6) ^ rotr32(e, 11) ^ rotr32(e, 25);
        uint32_t ch = (e & f) ^ ((~e) & g);
        uint32_t temp1 = h + S1 + ch + kSha256K[i] + w[i];
        uint32_t S0 = rotr32(a, 2) ^ rotr32(a, 13) ^ rotr32(a, 22);
        uint32_t maj = (a & b) ^ (a & c) ^ (b & c);
        uint32_t temp2 = S0 + maj;
        h = g; g = f; f = e; e = d + temp1;
        d = c; c = b; b = a; a = temp1 + temp2;
    }

    ctx->state[0] += a; ctx->state[1] += b; ctx->state[2] += c; ctx->state[3] += d;
    ctx->state[4] += e; ctx->state[5] += f; ctx->state[6] += g; ctx->state[7] += h;
}

static void sha256_init(Sha256Ctx* ctx) {
    static const uint32_t iv[8] = {
        0x6a09e667,0xbb67ae85,0x3c6ef372,0xa54ff53a,
        0x510e527f,0x9b05688c,0x1f83d9ab,0x5be0cd19,
    };
    memcpy(ctx->state, iv, sizeof(iv));
    ctx->total_len = 0;
    ctx->buf_len = 0;
}

static void sha256_update(Sha256Ctx* ctx, const uint8_t* data, size_t len) {
    ctx->total_len += len;
    while (len > 0) {
        size_t take = 64 - ctx->buf_len;
        if (take > len) take = len;
        memcpy(ctx->buf + ctx->buf_len, data, take);
        ctx->buf_len += take;
        data += take;
        len -= take;
        if (ctx->buf_len == 64) {
            sha256_compress(ctx, ctx->buf);
            ctx->buf_len = 0;
        }
    }
}

static void sha256_final(Sha256Ctx* ctx, uint8_t out[32]) {
    uint64_t bit_len = ctx->total_len * 8;
    uint8_t pad = 0x80;
    sha256_update(ctx, &pad, 1);
    uint8_t zero = 0x00;
    while (ctx->buf_len != 56) sha256_update(ctx, &zero, 1);
    uint8_t len_be[8];
    for (int i = 0; i < 8; i++) len_be[i] = (uint8_t)(bit_len >> (56 - 8*i));
    // Bypass the length-accounting update() wrapper for the length field
    // itself (it must not count toward total_len).
    memcpy(ctx->buf + 56, len_be, 8);
    sha256_compress(ctx, ctx->buf);
    for (int i = 0; i < 8; i++) {
        out[i*4]   = (uint8_t)(ctx->state[i] >> 24);
        out[i*4+1] = (uint8_t)(ctx->state[i] >> 16);
        out[i*4+2] = (uint8_t)(ctx->state[i] >> 8);
        out[i*4+3] = (uint8_t)(ctx->state[i]);
    }
}

}  // namespace

void rns_full_hash(const uint8_t* data, size_t len, uint8_t out[RNS_HASHLEN]) {
    Sha256Ctx ctx;
    sha256_init(&ctx);
    if (data && len) sha256_update(&ctx, data, len);
    sha256_final(&ctx, out);
}

void rns_truncated_hash(const uint8_t* data, size_t len, uint8_t out[RNS_TRUNCATED_HASHLEN]) {
    uint8_t full[RNS_HASHLEN];
    rns_full_hash(data, len, full);
    memcpy(out, full, RNS_TRUNCATED_HASHLEN);
}

// ---------------------------------------------------------------------------
// Identity
// ---------------------------------------------------------------------------

void rns_identity_from_seeds(const uint8_t ed25519_seed[32], const uint8_t x25519_seed[32],
                             RnsIdentity* out) {
    memcpy(out->x25519_priv, x25519_seed, 32);
    x25519_keypair_from_seed(x25519_seed, out->x25519_pub);
    ed25519_keypair_from_seed(ed25519_seed, out->ed25519_pub, out->ed25519_sk);

    uint8_t pub64[RNS_PUBKEY_LEN];
    memcpy(pub64, out->x25519_pub, 32);
    memcpy(pub64 + 32, out->ed25519_pub, 32);
    rns_truncated_hash(pub64, RNS_PUBKEY_LEN, out->hash);
}

namespace {
RnsIdentity s_identity;
bool s_identity_ready = false;

#define RNS_IDENTITY_NVS_NS   "pdl_rnsid"
#define RNS_IDENTITY_NVS_KEY  "seeds"   // 64 bytes: ed25519_seed(32) || x25519_seed(32)
}  // namespace

#ifdef ARDUINO
#include <Arduino.h>
#include <esp_system.h>
#include <nvs_flash.h>
#include <nvs.h>

static bool nvs_load_seeds(uint8_t seeds64[64]) {
    nvs_handle_t h;
    if (nvs_open(RNS_IDENTITY_NVS_NS, NVS_READONLY, &h) != ESP_OK) return false;
    size_t len = 64;
    esp_err_t err = nvs_get_blob(h, RNS_IDENTITY_NVS_KEY, seeds64, &len);
    nvs_close(h);
    return (err == ESP_OK && len == 64);
}

static void nvs_save_seeds(const uint8_t seeds64[64]) {
    nvs_handle_t h;
    if (nvs_open(RNS_IDENTITY_NVS_NS, NVS_READWRITE, &h) != ESP_OK) return;
    nvs_set_blob(h, RNS_IDENTITY_NVS_KEY, seeds64, 64);
    nvs_commit(h);
    nvs_close(h);
}

void rns_identity_init() {
    if (s_identity_ready) return;
    uint8_t seeds[64];
    if (!nvs_load_seeds(seeds)) {
        esp_fill_random(seeds, 64);
        nvs_save_seeds(seeds);
        Serial.printf("[RnsIdentity] Generated a new Reticulum identity (first boot)\n");
    }
    rns_identity_from_seeds(seeds, seeds + 32, &s_identity);
    s_identity_ready = true;
}
#else  // native / test
void rns_identity_init() {
    if (s_identity_ready) return;
    // Fixed, deterministic identity for native builds/tests -- never
    // persisted, never used as a real identity off-device.
    uint8_t ed_seed[32], x_seed[32];
    for (int i = 0; i < 32; i++) { ed_seed[i] = (uint8_t)i; x_seed[i] = (uint8_t)(0xAA + i); }
    rns_identity_from_seeds(ed_seed, x_seed, &s_identity);
    s_identity_ready = true;
}
#endif

const RnsIdentity* rns_identity_get() {
    if (!s_identity_ready) rns_identity_init();
    return &s_identity;
}

#ifdef UNIT_TEST
void rns_identity_test_set(const RnsIdentity* id) {
    s_identity = *id;
    s_identity_ready = true;
}
#endif

// ---------------------------------------------------------------------------
// Destination hashing
// ---------------------------------------------------------------------------

void rns_name_hash(const char* app_name, const char* const* aspects, int n_aspects,
                   uint8_t out[RNS_NAME_HASHLEN]) {
    char name[128];
    size_t pos = 0;
    size_t alen = strlen(app_name);
    if (alen >= sizeof(name)) alen = sizeof(name) - 1;
    memcpy(name, app_name, alen);
    pos = alen;
    for (int i = 0; i < n_aspects && pos < sizeof(name) - 1; i++) {
        name[pos++] = '.';
        size_t asplen = strlen(aspects[i]);
        if (pos + asplen >= sizeof(name)) asplen = sizeof(name) - 1 - pos;
        memcpy(name + pos, aspects[i], asplen);
        pos += asplen;
    }
    uint8_t full[RNS_HASHLEN];
    rns_full_hash((const uint8_t*)name, pos, full);
    memcpy(out, full, RNS_NAME_HASHLEN);
}

void rns_destination_hash(const uint8_t name_hash[RNS_NAME_HASHLEN],
                          const uint8_t identity_hash[RNS_TRUNCATED_HASHLEN],
                          uint8_t out[RNS_TRUNCATED_HASHLEN]) {
    uint8_t material[RNS_NAME_HASHLEN + RNS_TRUNCATED_HASHLEN];
    memcpy(material, name_hash, RNS_NAME_HASHLEN);
    memcpy(material + RNS_NAME_HASHLEN, identity_hash, RNS_TRUNCATED_HASHLEN);
    rns_truncated_hash(material, sizeof(material), out);
}

// ---------------------------------------------------------------------------
// Announce packet
// ---------------------------------------------------------------------------

size_t rns_build_announce(const RnsIdentity* id,
                          const char* app_name, const char* const* aspects, int n_aspects,
                          const uint8_t* app_data, size_t app_data_len,
                          const uint8_t random_hash10[RNS_RANDOM_HASHLEN],
                          uint8_t* out, size_t out_max) {
    uint8_t name_hash[RNS_NAME_HASHLEN];
    rns_name_hash(app_name, aspects, n_aspects, name_hash);

    uint8_t dest_hash[RNS_TRUNCATED_HASHLEN];
    rns_destination_hash(name_hash, id->hash, dest_hash);

    uint8_t pubkey[RNS_PUBKEY_LEN];
    memcpy(pubkey, id->x25519_pub, 32);
    memcpy(pubkey + 32, id->ed25519_pub, 32);

    // signed_data = destination_hash || pubkey || name_hash || random_hash [|| app_data]
    // (no ratchet -- this module never enables ratchets, matching a real
    // RNS node that simply hasn't turned them on for this destination.)
    uint8_t signed_data[RNS_TRUNCATED_HASHLEN + RNS_PUBKEY_LEN + RNS_NAME_HASHLEN + RNS_RANDOM_HASHLEN + 64];
    size_t sd_len = 0;
    memcpy(signed_data + sd_len, dest_hash, RNS_TRUNCATED_HASHLEN);   sd_len += RNS_TRUNCATED_HASHLEN;
    memcpy(signed_data + sd_len, pubkey, RNS_PUBKEY_LEN);             sd_len += RNS_PUBKEY_LEN;
    memcpy(signed_data + sd_len, name_hash, RNS_NAME_HASHLEN);        sd_len += RNS_NAME_HASHLEN;
    memcpy(signed_data + sd_len, random_hash10, RNS_RANDOM_HASHLEN);  sd_len += RNS_RANDOM_HASHLEN;
    if (app_data && app_data_len > 0) {
        size_t room = sizeof(signed_data) - sd_len;
        size_t take = app_data_len < room ? app_data_len : room;
        memcpy(signed_data + sd_len, app_data, take);
        sd_len += take;
    }

    uint8_t sig[RNS_SIG_LEN];
    ed25519_sign(id->ed25519_sk, signed_data, sd_len, sig);

    // Wire packet: flags(1) + hops(1) + dest_hash(16) + context(1) +
    // [pubkey(64) + name_hash(10) + random_hash(10) + sig(64) + app_data]
    size_t total = 1 + 1 + RNS_TRUNCATED_HASHLEN + 1 +
                   RNS_PUBKEY_LEN + RNS_NAME_HASHLEN + RNS_RANDOM_HASHLEN + RNS_SIG_LEN +
                   (app_data ? app_data_len : 0);
    if (!out || total > out_max) return 0;

    size_t p = 0;
    // flags = (HEADER_1=0<<6)|(FLAG_UNSET=0<<5)|(BROADCAST=0<<4)|(SINGLE=0<<2)|(ANNOUNCE=1)
    out[p++] = 0x01;
    out[p++] = 0x00;   // hops
    memcpy(out + p, dest_hash, RNS_TRUNCATED_HASHLEN); p += RNS_TRUNCATED_HASHLEN;
    out[p++] = 0x00;   // context = NONE
    memcpy(out + p, pubkey, RNS_PUBKEY_LEN);            p += RNS_PUBKEY_LEN;
    memcpy(out + p, name_hash, RNS_NAME_HASHLEN);       p += RNS_NAME_HASHLEN;
    memcpy(out + p, random_hash10, RNS_RANDOM_HASHLEN); p += RNS_RANDOM_HASHLEN;
    memcpy(out + p, sig, RNS_SIG_LEN);                  p += RNS_SIG_LEN;
    if (app_data && app_data_len > 0) {
        memcpy(out + p, app_data, app_data_len);
        p += app_data_len;
    }
    return p;
}
