// =============================================================================
// PDL CONSULTING LLC — TRADE SECRET
// Copyright (c) 2026 PDL Consulting LLC. All rights reserved.
//
// src/talon_federated.h
// TALON Federated-Learning — T1 Receiver Side (v1)
//
// v1 scope: receive aggregated weight-delta updates from T2 over LoRa mesh,
// reassemble chunks in PSRAM candidate slot, verify integrity, apply atomically.
//
// v2 (DEFERRED): on-device gradient computation, federated PUSH back to T2.
// Stub function declarations are provided below for forward-compatibility;
// both return false unconditionally in v1.
//
// Wire format — MSG_TYPE_FED_UPDATE (0x1B) chunk layout (28-byte header):
//   [0]      = 0x1B  (msg_type)
//   [1..2]   = round_id   (uint16, big-endian)
//   [3..6]   = total_chunks (uint32, big-endian)
//   [7..10]  = chunk_idx   (uint32, big-endian)
//   [11..14] = delta_crc32 (uint32, big-endian; CRC of entire reassembled delta)
//   [15..18] = apply_after_ts (uint32, big-endian; Unix seconds; 0 = apply now)
//   [19..27] = reserved (ignored)
//   [28..]   = weight bytes for this chunk (up to 200 bytes)
//
// Memory:
//   Candidate model buffer — heap_caps_malloc(MALLOC_CAP_SPIRAM) at init.
//   Size = TALON_CANDIDATE_BUF_SIZE bytes.
//   Two static metadata structs (< 40 bytes SRAM) track reassembly state.
//
// Integration (main.cpp):
//   Add to the MSG switch:
//     case MSG_TYPE_FED_UPDATE:
//         talon_handle_chunk(rx_buf, (size_t)rx_len);
//         break;
//   After the switch, check talon_round_complete() and talon_apply_round().
//
// =============================================================================
#ifndef TALON_FEDERATED_H
#define TALON_FEDERATED_H

#include <cstdint>
#include <cstddef>

// ---------------------------------------------------------------------------
// Constants
// ---------------------------------------------------------------------------

// LoRa payload ceiling (RadioLib default MTU = 255 bytes; header = 28 bytes)
#define TALON_CHUNK_HEADER_BYTES  28u
#define TALON_CHUNK_MAX_PAYLOAD   200u

// Maximum reassembled delta size stored in PSRAM (64 KB — covers int8 deltas
// for the resident drone_classifier model family).
#define TALON_CANDIDATE_BUF_SIZE  (64u * 1024u)

// Maximum chunks we track per round (ceil(64 KB / 200) = 328; use 512 margin)
#define TALON_MAX_CHUNKS          512u

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

// Initialise the TALON receiver.
// Allocates TALON_CANDIDATE_BUF_SIZE bytes from PSRAM for the candidate slot.
// Safe to call from setup() before mesh_init().
void talon_init();

// Feed a raw LoRa payload into the reassembler.
// buf[0] must equal MSG_TYPE_FED_UPDATE (0x1B); len must be ≥ 28.
// Idempotent if the same chunk_idx is received twice (duplicate suppression).
void talon_handle_chunk(const uint8_t* buf, size_t len);

// Returns true when all chunks for the current round have been received
// and the full delta buffer's CRC32 matches the value carried in the header.
// Returns false if the CRC fails (candidate slot is poisoned and reset).
bool talon_round_complete();

// Apply the reassembled delta to the active model manager slot.
// Delegates to ota_signed_verify_chunk if available; falls back to CRC check.
// Atomically swaps the candidate slot content into the active model.
// Returns false if the round is not complete, signature/CRC fails, or
// apply_after_ts has not been reached yet.
// On success, resets round state so the next MSG_TYPE_FED_UPDATE starts fresh.
bool talon_apply_round();

// Current round_id (from the most recent MSG_TYPE_FED_UPDATE header).
// Returns 0 before any chunk has been received.
uint16_t talon_current_round();

// ---------------------------------------------------------------------------
// v2 stub declarations (DEFERRED — do NOT implement in v1)
//
// These forward-declare the federated PUSH side required when T1 sensors
// compute local gradients and upload them to T2 for server-side aggregation.
// Implementing these requires backprop over int8 activations; deferred until
// the model_manager exposes a gradient callback.
// ---------------------------------------------------------------------------

// [v2] Collect local-anchor samples for the current detection window.
// In v1: always returns false.
bool talon_collect_local_anchors();

// [v2] Compute a gradient delta from the local anchor set and stage it
// in the PSRAM candidate slot for upward transmission to T2.
// In v1: always returns false.
bool talon_compute_gradient();

#endif // TALON_FEDERATED_H
