#include "ota_adapter.h"
#include "rf_classifier.h"
#include <cstring>

#ifdef ARDUINO
#include <Arduino.h>
#include <esp_heap_caps.h>
#else
// Host-side stubs so this module (previously buildable ONLY on-device) can
// be exercised by a native Unity test -- same pattern used throughout this
// sweep (see health_report.cpp, uav_relay.cpp, etc).
#include <cstdio>
#include <cstdarg>
#include <cstdlib>
struct NativeSerialStub {
    void printf(const char* fmt, ...) {
        va_list args;
        va_start(args, fmt);
        vprintf(fmt, args);
        va_end(args);
    }
};
static NativeSerialStub Serial;
#define MALLOC_CAP_SPIRAM 0
static inline void* heap_caps_malloc(size_t size, uint32_t caps) {
    (void)caps;
    return malloc(size);
}
static inline void heap_caps_free(void* ptr) {
    free(ptr);
}
#endif // ARDUINO

static const char* TAG = "OTA";

// Double buffer slots in PSRAM
static uint8_t* s_slot_a = nullptr;
static uint8_t* s_slot_b = nullptr;
static uint8_t  s_active_slot = 0;  // 0 = A, 1 = B

// Reception state
static uint8_t* s_recv_buf = nullptr;    // points to inactive slot
static size_t   s_recv_size = 0;
static uint8_t  s_total_chunks = 0;
static uint8_t  s_received_mask[32];     // bitfield: up to 256 chunks
static uint8_t  s_chunks_received = 0;
static uint16_t s_staged_version = 0;
static uint16_t s_active_version = 0;
static bool     s_model_complete = false;

// ---------------------------------------------------------------------------
// CRC-16/CCITT for chunk verification
// ---------------------------------------------------------------------------
static uint16_t crc16_ccitt(const uint8_t* data, size_t len) {
    uint16_t crc = 0xFFFF;
    for (size_t i = 0; i < len; i++) {
        crc ^= (uint16_t)data[i] << 8;
        for (int j = 0; j < 8; j++) {
            if (crc & 0x8000)
                crc = (crc << 1) ^ 0x1021;
            else
                crc <<= 1;
        }
    }
    return crc;
}

// ---------------------------------------------------------------------------
bool ota_init() {
    s_slot_a = (uint8_t*)heap_caps_malloc(OTA_MODEL_SLOT_SIZE, MALLOC_CAP_SPIRAM);
    s_slot_b = (uint8_t*)heap_caps_malloc(OTA_MODEL_SLOT_SIZE, MALLOC_CAP_SPIRAM);

    if (!s_slot_a || !s_slot_b) {
        Serial.printf("[%s] Failed to allocate double-buffer slots in PSRAM\n", TAG);
        return false;
    }

    memset(s_slot_a, 0, OTA_MODEL_SLOT_SIZE);
    memset(s_slot_b, 0, OTA_MODEL_SLOT_SIZE);

    // A fresh init means OTA's own bookkeeping has no known staged or
    // active model -- without this, a deinit()+init() cycle (currently
    // unreachable in production since ota_deinit() is never called from
    // main.cpp, but exercised by tests and a plausible future recovery
    // path) would silently carry over version numbers from before the
    // buffers were freed and reallocated.
    s_staged_version = 0;
    s_active_version = 0;

    s_active_slot = 0;
    s_recv_buf = s_slot_b;  // receive into inactive slot
    ota_reset();

    Serial.printf("[%s] Double-buffer initialized. %d KB per slot\n",
                  TAG, OTA_MODEL_SLOT_SIZE / 1024);
    return true;
}

// ---------------------------------------------------------------------------
bool ota_handle_chunk(const uint8_t* data, size_t len) {
    // Data layout after the 33-byte detection event header:
    // We expect the raw OTA payload starting at data[0] to include
    // the chunk sub-header.
    // Minimum: DETECTION_PACKET_SIZE + OTA_CHUNK_HEADER_SIZE + 1 byte payload
    if (len < DETECTION_PACKET_SIZE + OTA_CHUNK_HEADER_SIZE + 1) {
        Serial.printf("[%s] Chunk packet too short: %u bytes\n", TAG, len);
        return false;
    }

    size_t offset = DETECTION_PACKET_SIZE;  // skip detection header
    uint8_t chunk_index   = data[offset + 0];
    uint8_t total_chunks  = data[offset + 1];
    uint16_t chunk_crc    = ((uint16_t)data[offset + 2] << 8) | data[offset + 3];
    uint16_t model_ver    = ((uint16_t)data[offset + 4] << 8) | data[offset + 5];
    // data[offset + 6] is reserved

    const uint8_t* payload = &data[offset + OTA_CHUNK_HEADER_SIZE];
    size_t payload_len = len - offset - OTA_CHUNK_HEADER_SIZE;

    // Verify CRC
    uint16_t computed_crc = crc16_ccitt(payload, payload_len);
    if (computed_crc != chunk_crc) {
        Serial.printf("[%s] Chunk %d CRC mismatch: expected 0x%04X got 0x%04X\n",
                      TAG, chunk_index, chunk_crc, computed_crc);
        return false;
    }

    // If new transfer, reset state
    if (model_ver != s_staged_version || total_chunks != s_total_chunks) {
        ota_reset();
        s_staged_version = model_ver;
        s_total_chunks = total_chunks;
        Serial.printf("[%s] New model transfer: version=%d, chunks=%d\n",
                      TAG, model_ver, total_chunks);
    }

    // Check for duplicate
    uint8_t byte_idx = chunk_index / 8;
    uint8_t bit_idx  = chunk_index % 8;
    if (s_received_mask[byte_idx] & (1 << bit_idx)) {
        // Already received this chunk
        return true;
    }

    // Store chunk
    size_t dest_offset = (size_t)chunk_index * OTA_MAX_CHUNK_PAYLOAD;
    if (dest_offset + payload_len > OTA_MODEL_SLOT_SIZE) {
        Serial.printf("[%s] Chunk %d would overflow slot buffer\n", TAG, chunk_index);
        return false;
    }

    memcpy(&s_recv_buf[dest_offset], payload, payload_len);
    s_received_mask[byte_idx] |= (1 << bit_idx);
    s_chunks_received++;

    // Track total size for last chunk
    if (chunk_index == total_chunks - 1) {
        s_recv_size = dest_offset + payload_len;
    }

    Serial.printf("[%s] Chunk %d/%d received (%u bytes)\n",
                  TAG, chunk_index + 1, total_chunks, payload_len);

    // Check completion
    if (s_chunks_received >= s_total_chunks && s_recv_size > 0) {
        s_model_complete = true;
        Serial.printf("[%s] Model transfer complete! %u bytes, version %d\n",
                      TAG, s_recv_size, s_staged_version);
    }

    return true;
}

// ---------------------------------------------------------------------------
bool ota_model_ready() {
    return s_model_complete;
}

// ---------------------------------------------------------------------------
bool ota_apply_model() {
    if (!s_model_complete) {
        Serial.printf("[%s] No complete model to apply\n", TAG);
        return false;
    }

    Serial.printf("[%s] Applying model version %d (%u bytes)\n",
                  TAG, s_staged_version, s_recv_size);

    bool ok = classifier_swap_model(s_recv_buf, s_recv_size);
    if (ok) {
        s_active_version = s_staged_version;

        // Swap active slot
        s_active_slot = (s_active_slot == 0) ? 1 : 0;
        s_recv_buf = (s_active_slot == 0) ? s_slot_b : s_slot_a;

        s_model_complete = false;
        Serial.printf("[%s] Model v%d is now active\n", TAG, s_active_version);
    } else {
        Serial.printf("[%s] Model swap failed!\n", TAG);
    }

    return ok;
}

// ---------------------------------------------------------------------------
uint16_t ota_staged_version() {
    return s_staged_version;
}

uint16_t ota_active_version() {
    return s_active_version;
}

// ---------------------------------------------------------------------------
void ota_reset() {
    s_total_chunks = 0;
    s_chunks_received = 0;
    s_recv_size = 0;
    s_model_complete = false;
    memset(s_received_mask, 0, sizeof(s_received_mask));
}

// ---------------------------------------------------------------------------
void ota_deinit() {
    if (s_slot_a) { heap_caps_free(s_slot_a); s_slot_a = nullptr; }
    if (s_slot_b) { heap_caps_free(s_slot_b); s_slot_b = nullptr; }
    s_recv_buf = nullptr;
    ota_reset();
}
