RNSDiscovery.h

#pragma once

#if defined(HAS_RNS)

#include <microReticulum/Bytes.h>
#include <microReticulum/Destination.h>
#include <microReticulum/Identity.h>
#include <microReticulum/Reticulum.h>
#include <microReticulum/Transport.h>
#include <microReticulum/Cryptography/HKDF.h>
#include <microReticulum/Cryptography/Random.h>
#include <SHA256.h>

#if MCU_VARIANT == MCU_ESP32
  #include <esp_task_wdt.h>
#endif

#define RNS_DISCOVERY_WORKBLOCK_ROUNDS 20
#define RNS_DISCOVERY_STAMP_SIZE 32
#define RNS_DISCOVERY_YIELD_INTERVAL_MS 25
#define RNS_DISCOVERY_RETRY_COOLDOWN_MS (5UL * 60UL * 1000UL)
#define RNS_DISCOVERY_ANNOUNCE_FAIL_WINDOW_MS (2UL * 60UL * 1000UL)
#define RNS_DISCOVERY_SKIP_LOG_INTERVAL_MS 60000UL
#define RNS_DISCOVERY_MAX_RETRIES 2

static RNS::Destination rns_discovery_destination(RNS::Type::NONE);
static bool rns_discovery_destination_ready = false;
static uint32_t rns_discovery_last_announce_ms = 0;
static uint32_t rns_discovery_last_attempt_ms = 0;
static uint32_t rns_discovery_retry_pending_since_ms = 0;
static uint32_t rns_discovery_last_skip_log_ms = 0;
static bool rns_discovery_retry_pending = false;
static uint8_t rns_discovery_retry_attempts = 0;
static bool rns_discovery_cache_valid = false;
static RNS::Bytes rns_discovery_cached_infohash;
static RNS::Bytes rns_discovery_cached_app_data;

void rns_discovery_service_io();
uint32_t rns_discovery_interval_ms();

void rns_discovery_breadcrumb(const char *event, bool force, const char *reason, uint32_t now, uint32_t elapsed, uint32_t interval, size_t app_len) {
#if defined(ARDUINO)
	if (!event) event = "unknown";
	if (!reason) reason = "-";
	Serial.printf("URTNDBG discovery_event=%s force=%u pending=%u retries=%u reason=%s elapsed=%lu interval=%lu app_len=%u ms=%lu\r\n",
		event,
		force ? 1 : 0,
		rns_discovery_retry_pending ? 1 : 0,
		rns_discovery_retry_attempts,
		reason,
		(unsigned long)elapsed,
		(unsigned long)interval,
		(unsigned int)app_len,
		(unsigned long)now);
#else
	(void)event;
	(void)force;
	(void)reason;
	(void)now;
	(void)elapsed;
	(void)interval;
	(void)app_len;
#endif
}

uint32_t rns_discovery_elapsed_ms(uint32_t now) {
	return (uint32_t)(now - rns_discovery_last_announce_ms);
}

bool rns_discovery_ready() {
	return rns_discovery_destination_ready && rns_config_discoverable_enabled();
}

bool rns_discovery_due(uint32_t now) {
	if (!rns_discovery_ready()) return false;
	return rns_discovery_elapsed_ms(now) >= rns_discovery_interval_ms();
}

void rns_discovery_mark_retry_pending(const char *reason) {
	if (!rns_discovery_ready()) return;
	uint32_t now = millis();
	uint32_t interval = rns_discovery_interval_ms();
	uint32_t elapsed = rns_discovery_elapsed_ms(now);
	if (rns_discovery_retry_attempts >= RNS_DISCOVERY_MAX_RETRIES) {
		rns_discovery_breadcrumb("retry_suppressed", false, reason, now, elapsed, interval, 0);
		return;
	}
	if (!rns_discovery_retry_pending) {
		rns_discovery_retry_pending = true;
		rns_discovery_retry_pending_since_ms = now;
		rns_discovery_breadcrumb("retry_pending", false, reason, now, elapsed, interval, 0);
	}
}

void rns_discovery_note_tcp_offline(bool was_online) {
	if (!rns_discovery_ready()) return;
	uint32_t now = millis();
	uint32_t interval = rns_discovery_interval_ms();
	uint32_t elapsed = rns_discovery_elapsed_ms(now);
	bool due = elapsed >= interval;
	bool recent_attempt = rns_discovery_last_attempt_ms != 0 &&
		(uint32_t)(now - rns_discovery_last_attempt_ms) <= RNS_DISCOVERY_ANNOUNCE_FAIL_WINDOW_MS;
	const char *reason = due ? "tcp_offline_due" : (recent_attempt && was_online ? "tcp_drop_after_attempt" : "tcp_offline");

	if (due || (recent_attempt && was_online)) {
		rns_discovery_mark_retry_pending(reason);
	}
	if ((due || rns_discovery_retry_pending) &&
		(rns_discovery_last_skip_log_ms == 0 ||
		 (uint32_t)(now - rns_discovery_last_skip_log_ms) >= RNS_DISCOVERY_SKIP_LOG_INTERVAL_MS)) {
		rns_discovery_last_skip_log_ms = now;
		rns_discovery_breadcrumb("announce_skipped_tcp_offline", false, reason, now, elapsed, interval, 0);
	}
}

void rns_mp_append_uint(RNS::Bytes &out, uint32_t value) {
	if (value <= 0x7F) {
		out.append((uint8_t)value);
	} else if (value <= 0xFF) {
		out.append((uint8_t)0xCC);
		out.append((uint8_t)value);
	} else if (value <= 0xFFFF) {
		out.append((uint8_t)0xCD);
		out.append((uint8_t)(value >> 8));
		out.append((uint8_t)value);
	} else {
		out.append((uint8_t)0xCE);
		out.append((uint8_t)(value >> 24));
		out.append((uint8_t)(value >> 16));
		out.append((uint8_t)(value >> 8));
		out.append((uint8_t)value);
	}
}

void rns_mp_append_bool(RNS::Bytes &out, bool value) {
	out.append((uint8_t)(value ? 0xC3 : 0xC2));
}

void rns_mp_append_nil(RNS::Bytes &out) {
	out.append((uint8_t)0xC0);
}

void rns_mp_append_map(RNS::Bytes &out, uint8_t count) {
	if (count <= 15) {
		out.append((uint8_t)(0x80 | count));
	} else {
		out.append((uint8_t)0xDE);
		out.append((uint8_t)0x00);
		out.append(count);
	}
}

void rns_mp_append_str(RNS::Bytes &out, const char *value) {
	if (!value) value = "";
	size_t len = strlen(value);
	if (len <= 31) {
		out.append((uint8_t)(0xA0 | len));
	} else {
		if (len > 255) len = 255;
		out.append((uint8_t)0xD9);
		out.append((uint8_t)len);
	}
	out.append((const uint8_t*)value, len);
}

void rns_mp_append_bin(RNS::Bytes &out, const uint8_t *data, size_t len) {
	if (len <= 255) {
		out.append((uint8_t)0xC4);
		out.append((uint8_t)len);
	} else {
		out.append((uint8_t)0xC5);
		out.append((uint8_t)(len >> 8));
		out.append((uint8_t)len);
	}
	out.append(data, len);
}

void rns_mp_append_double(RNS::Bytes &out, double value) {
	uint64_t bits = 0;
	memcpy(&bits, &value, sizeof(bits));
	out.append((uint8_t)0xCB);
	for (int8_t i = 7; i >= 0; i--) {
		out.append((uint8_t)(bits >> (8 * i)));
	}
}

void rns_mp_append_coord(RNS::Bytes &out, int32_t value_e7) {
	if (value_e7 == RNS_CFG_COORD_UNSET) {
		rns_mp_append_nil(out);
	} else {
		rns_mp_append_double(out, (double)value_e7 / 10000000.0);
	}
}

void rns_mp_append_height(RNS::Bytes &out, int16_t height_m) {
	if (height_m == RNS_CFG_HEIGHT_UNSET) {
		rns_mp_append_nil(out);
	} else {
		rns_mp_append_double(out, (double)height_m);
	}
}

uint8_t rns_discovery_stamp_value(const RNS::Bytes &workblock, const RNS::Bytes &stamp) {
	SHA256 sha256;
	sha256.reset();
	sha256.update(workblock.data(), workblock.size());
	sha256.update(stamp.data(), stamp.size());
	RNS::Bytes digest;
	sha256.finalize(digest.writable(32), 32);
	uint8_t value = 0;
	for (size_t i = 0; i < digest.size(); i++) {
		uint8_t b = digest.data()[i];
		if (b == 0x00) {
			value += 8;
			continue;
		}
		for (int8_t bit = 7; bit >= 0; bit--) {
			if (b & (1 << bit)) return value;
			value++;
		}
	}
	return value;
}

void rns_discovery_yield() {
	rns_discovery_service_io();
	yield();
#if MCU_VARIANT == MCU_ESP32
	esp_task_wdt_reset();
#endif
}

RNS::Bytes rns_discovery_workblock(const RNS::Bytes &material) {
	RNS::Bytes workblock;
	for (uint8_t n = 0; n < RNS_DISCOVERY_WORKBLOCK_ROUNDS; n++) {
		RNS::Bytes salt_material = material;
		rns_mp_append_uint(salt_material, n);
		RNS::Bytes salt = RNS::Identity::full_hash(salt_material);
		workblock.append(RNS::Cryptography::hkdf(256, material, salt));
		rns_discovery_yield();
	}
	return workblock;
}

RNS::Bytes rns_discovery_generate_stamp(const RNS::Bytes &infohash, uint8_t target_value) {
	RNS::Bytes workblock = rns_discovery_workblock(infohash);
	uint32_t rounds = 0;
	uint32_t last_yield = millis();
	while (true) {
		RNS::Bytes stamp = RNS::Cryptography::random(RNS_DISCOVERY_STAMP_SIZE);
		if (rns_discovery_stamp_value(workblock, stamp) >= target_value) {
			TRACEF("Generated discovery stamp with target %u after %u rounds", target_value, rounds);
			return stamp;
		}
		rounds++;
		uint32_t now = millis();
		if ((uint32_t)(now - last_yield) >= RNS_DISCOVERY_YIELD_INTERVAL_MS) {
			last_yield = now;
			rns_discovery_yield();
		}
	}
}

void rns_discovery_effective_name(char *dst, size_t size) {
	if (rns_config.discovery_name[0] != 0x00) {
		rns_config_set_cstr(dst, size, rns_config.discovery_name);
	} else {
		snprintf(dst, size, "microReticulum %s", device_uid_str);
	}
}

RNS::Bytes rns_discovery_build_app_data() {
	char name[RNS_CFG_NAME_SIZE];
	rns_discovery_effective_name(name, sizeof(name));

	RNS::Bytes packed;
	rns_mp_append_map(packed, 11);
	rns_mp_append_uint(packed, 0x00);
	rns_mp_append_str(packed, "RNodeInterface");
	rns_mp_append_uint(packed, 0x01);
	rns_mp_append_bool(packed, RNS::Reticulum::transport_enabled());
	rns_mp_append_uint(packed, 0xFE);
	const RNS::Bytes &transport_id = RNS::Transport::identity().hash();
	rns_mp_append_bin(packed, transport_id.data(), transport_id.size());
	rns_mp_append_uint(packed, 0xFF);
	rns_mp_append_str(packed, name);
	rns_mp_append_uint(packed, 0x03);
	rns_mp_append_coord(packed, rns_config.latitude_e7);
	rns_mp_append_uint(packed, 0x04);
	rns_mp_append_coord(packed, rns_config.longitude_e7);
	rns_mp_append_uint(packed, 0x05);
	rns_mp_append_height(packed, rns_config.height_m);
	rns_mp_append_uint(packed, 0x09);
	rns_mp_append_uint(packed, lora_freq);
	rns_mp_append_uint(packed, 0x0A);
	rns_mp_append_uint(packed, lora_bw);
	rns_mp_append_uint(packed, 0x0B);
	rns_mp_append_uint(packed, (uint32_t)lora_sf);
	rns_mp_append_uint(packed, 0x0C);
	rns_mp_append_uint(packed, (uint32_t)lora_cr);

	RNS::Bytes infohash = RNS::Identity::full_hash(packed);
	if (rns_discovery_cache_valid && infohash == rns_discovery_cached_infohash) {
		return rns_discovery_cached_app_data;
	}

	RNS::Bytes stamp = rns_discovery_generate_stamp(infohash, rns_config.discovery_stamp_value);
	RNS::Bytes app_data;
	app_data.append((uint8_t)0x00);
	app_data.append(packed);
	app_data.append(stamp);
	rns_discovery_cached_infohash = infohash;
	rns_discovery_cached_app_data = app_data;
	rns_discovery_cache_valid = true;
	return app_data;
}

uint32_t rns_discovery_interval_ms() {
	return (uint32_t)rns_config.announce_interval_minutes * 60UL * 1000UL;
}

void rns_discovery_init_destination() {
	if (!rns_config_discoverable_enabled()) return;
	rns_discovery_destination = RNS::Destination(
		RNS::Transport::identity(),
		RNS::Type::Destination::IN,
		RNS::Type::Destination::SINGLE,
		"rnstransport",
		"discovery.interface"
	);
	rns_discovery_destination_ready = true;
	rns_discovery_last_announce_ms = millis() - rns_discovery_interval_ms();
	rns_discovery_last_attempt_ms = 0;
	rns_discovery_retry_pending_since_ms = 0;
	rns_discovery_last_skip_log_ms = 0;
	rns_discovery_retry_pending = false;
	rns_discovery_retry_attempts = 0;
	TRACEF("RNS interface discovery destination: %s", rns_discovery_destination.hash().toHex().c_str());
}

bool announce_rns_discovery(bool force) {
	if (!rns_discovery_ready()) return false;
	uint32_t now = millis();
	uint32_t interval = rns_discovery_interval_ms();
	uint32_t elapsed = rns_discovery_elapsed_ms(now);
	bool due = elapsed >= interval;
	bool retry_ready = rns_discovery_retry_pending &&
		(uint32_t)(now - rns_discovery_retry_pending_since_ms) >= RNS_DISCOVERY_RETRY_COOLDOWN_MS;
	if (!due && !retry_ready) return false;
	if (due && rns_discovery_last_attempt_ms != 0 &&
		(uint32_t)(now - rns_discovery_last_attempt_ms) < RNS_DISCOVERY_RETRY_COOLDOWN_MS) {
		return false;
	}

	const bool retrying = retry_ready && !due;
	rns_discovery_breadcrumb(due ? "announce_due" : "announce_retry_due", force, retrying ? "pending_retry" : "scheduled", now, elapsed, interval, 0);

	try {
		rns_discovery_last_attempt_ms = now;
		RNS::Bytes app_data = rns_discovery_build_app_data();
		rns_discovery_breadcrumb("announce_attempt", force, retrying ? "pending_retry" : "scheduled", now, elapsed, interval, app_data.size());
		TRACEF("Announcing RNS interface discovery with %u bytes of app_data", app_data.size());
		rns_discovery_destination.announce(app_data);
		now = millis();
		rns_discovery_last_announce_ms = now;
		if (retrying && rns_discovery_retry_attempts < 0xFF) {
			rns_discovery_retry_attempts++;
		} else if (!retrying) {
			rns_discovery_retry_attempts = 0;
		}
		rns_discovery_retry_pending = false;
		rns_discovery_retry_pending_since_ms = 0;
		rns_discovery_breadcrumb("announce_done", force, retrying ? "pending_retry" : "scheduled", now, 0, interval, app_data.size());
		return true;
	}
	catch (const std::bad_alloc&) {
		ERROR("RNS discovery announce failed: bad_alloc - out of memory");
		rns_discovery_mark_retry_pending("bad_alloc");
	}
	catch (std::exception& e) {
		ERRORF("RNS discovery announce failed: %s", e.what());
		rns_discovery_mark_retry_pending("exception");
	}
	return false;
}

#endif