#include "model_manager.h"

#ifdef ARDUINO
#include <Arduino.h>
#include <LittleFS.h>
#include <esp_heap_caps.h>
#else
// Host-side stubs so this module can be exercised by a native Unity test --
// same pattern used throughout the codebase (see ota_adapter.cpp,
// health_report.cpp, uav_relay.cpp). Only Serial + heap_caps are stubbed
// because those are the only unconditionally-included-header APIs the
// PDL_DISABLE_TFLITE stub branch below actually calls; LittleFS/File are
// only touched by the real (non-stub) branch, which native tests don't
// build (TFLite-Micro isn't natively host-buildable here).
#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

#include <cstring>

// ---------------------------------------------------------------------------
// CRC-32/BZIP2 (poly 0x04C11DB7, init 0xFFFFFFFF, no reflection, xorout
// 0xFFFFFFFF). Matches the variant used by mesh/command packets. Hoisted out
// of the PDL_DISABLE_TFLITE #else block (it used to live next to
// model_manager_commit_slot, TFLite-coupled code below) so it's testable
// natively without needing a real TFLite-Micro build -- the CRC check in
// commit_slot() is the one piece of this file's safety-relevant logic that
// doesn't actually need a live interpreter to verify.
// ---------------------------------------------------------------------------
static uint32_t model_manager_crc32_bzip2(const uint8_t* data, size_t len) {
    uint32_t crc = 0xFFFFFFFFu;
    for (size_t i = 0; i < len; i++) {
        crc ^= ((uint32_t)data[i]) << 24;
        for (int b = 0; b < 8; b++) {
            if (crc & 0x80000000u) crc = (crc << 1) ^ 0x04C11DB7u;
            else                   crc = (crc << 1);
        }
    }
    return crc ^ 0xFFFFFFFFu;
}

#ifdef PDL_DISABLE_TFLITE
// ============================================================================
// model_manager stub mode — TFLite-Micro disabled (toolchain incompat).
// All slot operations no-op; cactus.full_mode runtime carries inference.
// ============================================================================
static const char* TAG = "ModelMgr";
int  model_manager_init()       { Serial.printf("[%s] TFLite disabled\n", TAG); return 0; }
int  model_manager_register(const char*, const char*, uint16_t, uint8_t, ModelPriority) { return -1; }
ModelInferenceResult model_manager_infer(uint8_t, const float*, uint16_t) {
    ModelInferenceResult r{}; r.valid = false; return r;
}
int  model_manager_infer_all_resident(const float*, uint16_t, ModelInferenceResult*, int) { return 0; }
bool model_manager_page_in(uint8_t)  { return false; }
void model_manager_page_out(uint8_t) {}
bool model_manager_is_loaded(uint8_t){ return false; }
bool model_manager_get_stats(uint8_t, ModelStats& s) { std::memset(&s, 0, sizeof(s)); return false; }
uint8_t model_manager_count()        { return 0; }
bool model_manager_commit_slot(uint8_t, uint32_t) { return false; }
uint32_t model_manager_psram_used() { return 0; }
void model_manager_deinit() {}
#else
#include "tensorflow/lite/micro/all_ops_resolver.h"
#include "tensorflow/lite/micro/micro_interpreter.h"
#include "tensorflow/lite/schema/schema_generated.h"

static const char* TAG = "ModelMgr";

// ---------------------------------------------------------------------------
// Per-slot internal state
// ---------------------------------------------------------------------------
struct ModelSlot {
    // Registration info
    char        name[MODEL_NAME_MAX_LEN];
    char        spiffs_path[48];
    uint16_t    input_size;         // number of float32 input elements
    uint8_t     output_classes;     // number of output classes
    ModelPriority priority;

    // Runtime state
    bool        registered;
    bool        loaded;             // model bytes are in PSRAM
    uint8_t*    model_buf;          // PSRAM buffer holding .tflite flatbuffer
    size_t      model_size;         // bytes
    uint8_t*    arena_buf;          // PSRAM tensor arena
    uint32_t    arena_alloc;        // arena allocation size
    uint32_t    arena_used;         // actual used bytes after AllocateTensors

    // TFLite runtime (one resolver shared, interpreter per-slot)
    const tflite::Model*      tfl_model;
    tflite::MicroInterpreter* interpreter;

    // Statistics
    uint32_t    inference_count;
    uint64_t    total_latency_us;   // sum of all inference times
    uint32_t    last_run_ms;        // millis() of last inference
};

static ModelSlot            s_slots[MODEL_MAX_SLOTS];
static uint8_t              s_registered_count = 0;
static tflite::AllOpsResolver s_resolver;     // shared across all models
static bool                 s_spiffs_mounted = false;

// Scratch interpreters — TFLite-Micro MicroInterpreter uses placement,
// so we allocate them in PSRAM as raw bytes and construct in-place.
// Each interpreter is small (~200 bytes) but must persist while in use.
static uint8_t* s_interp_mem[MODEL_MAX_SLOTS] = { nullptr };
#define INTERPRETER_MEM_SIZE 512  // generous for MicroInterpreter object

// ---------------------------------------------------------------------------
// Internal: ensure SPIFFS is mounted
// ---------------------------------------------------------------------------
static bool ensure_spiffs() {
    if (s_spiffs_mounted) return true;
    if (!LittleFS.begin(true)) {
        Serial.printf("[%s] SPIFFS mount failed\n", TAG);
        return false;
    }
    s_spiffs_mounted = true;
    return true;
}

// ---------------------------------------------------------------------------
// Internal: load model bytes from SPIFFS into PSRAM
// ---------------------------------------------------------------------------
static bool load_model_bytes(ModelSlot& slot) {
    if (!ensure_spiffs()) return false;

    File f = LittleFS.open(slot.spiffs_path, "r");
    if (!f) {
        Serial.printf("[%s] Cannot open: %s\n", TAG, slot.spiffs_path);
        return false;
    }

    slot.model_size = f.size();
    if (slot.model_size == 0) {
        Serial.printf("[%s] Empty model file: %s\n", TAG, slot.spiffs_path);
        f.close();
        return false;
    }

    slot.model_buf = (uint8_t*)heap_caps_malloc(slot.model_size, MALLOC_CAP_SPIRAM);
    if (!slot.model_buf) {
        Serial.printf("[%s] PSRAM alloc failed for %s (%u bytes)\n",
                      TAG, slot.name, slot.model_size);
        f.close();
        return false;
    }

    size_t bytes_read = f.read(slot.model_buf, slot.model_size);
    f.close();

    if (bytes_read != slot.model_size) {
        Serial.printf("[%s] Incomplete read: %s (%u/%u)\n",
                      TAG, slot.name, bytes_read, slot.model_size);
        heap_caps_free(slot.model_buf);
        slot.model_buf = nullptr;
        return false;
    }

    Serial.printf("[%s] Loaded %s from %s (%u bytes)\n",
                  TAG, slot.name, slot.spiffs_path, slot.model_size);
    return true;
}

// ---------------------------------------------------------------------------
// Internal: build TFLite interpreter for a loaded model
// ---------------------------------------------------------------------------
static bool build_slot_interpreter(uint8_t slot_id) {
    ModelSlot& slot = s_slots[slot_id];

    // Parse flatbuffer
    slot.tfl_model = tflite::GetModel(slot.model_buf);
    if (slot.tfl_model->version() != TFLITE_SCHEMA_VERSION) {
        Serial.printf("[%s] %s: schema version mismatch %lu vs %d\n",
                      TAG, slot.name, slot.tfl_model->version(), TFLITE_SCHEMA_VERSION);
        return false;
    }

    // Compute arena size: base arena + scaling for larger models
    // Specialist models with fewer classes can use smaller arenas
    slot.arena_alloc = MODEL_ARENA_SIZE_PER_SLOT;

    slot.arena_buf = (uint8_t*)heap_caps_malloc(slot.arena_alloc, MALLOC_CAP_SPIRAM);
    if (!slot.arena_buf) {
        Serial.printf("[%s] %s: arena alloc failed (%u bytes)\n",
                      TAG, slot.name, slot.arena_alloc);
        return false;
    }

    // Allocate interpreter memory in PSRAM
    if (!s_interp_mem[slot_id]) {
        s_interp_mem[slot_id] = (uint8_t*)heap_caps_malloc(
            INTERPRETER_MEM_SIZE, MALLOC_CAP_SPIRAM);
        if (!s_interp_mem[slot_id]) {
            Serial.printf("[%s] %s: interpreter mem alloc failed\n", TAG, slot.name);
            heap_caps_free(slot.arena_buf);
            slot.arena_buf = nullptr;
            return false;
        }
    }

    // Construct interpreter in-place (modern MicroTFLite: 4-arg)
    slot.interpreter = new (s_interp_mem[slot_id]) tflite::MicroInterpreter(
        slot.tfl_model, s_resolver, slot.arena_buf, slot.arena_alloc);

    TfLiteStatus status = slot.interpreter->AllocateTensors();
    if (status != kTfLiteOk) {
        Serial.printf("[%s] %s: AllocateTensors failed\n", TAG, slot.name);
        slot.interpreter = nullptr;
        heap_caps_free(slot.arena_buf);
        slot.arena_buf = nullptr;
        return false;
    }

    slot.arena_used = slot.interpreter->arena_used_bytes();
    slot.loaded = true;

    Serial.printf("[%s] %s: interpreter ready. Arena %u / %u bytes\n",
                  TAG, slot.name, slot.arena_used, slot.arena_alloc);
    return true;
}

// ---------------------------------------------------------------------------
// Internal: free a slot's runtime resources (but keep registration)
// ---------------------------------------------------------------------------
static void unload_slot(uint8_t slot_id) {
    ModelSlot& slot = s_slots[slot_id];
    if (!slot.loaded) return;

    slot.interpreter = nullptr;
    slot.tfl_model = nullptr;

    if (slot.arena_buf) {
        heap_caps_free(slot.arena_buf);
        slot.arena_buf = nullptr;
    }
    if (slot.model_buf) {
        heap_caps_free(slot.model_buf);
        slot.model_buf = nullptr;
    }

    slot.loaded = false;
    slot.arena_used = 0;

    Serial.printf("[%s] Unloaded slot %d (%s)\n", TAG, slot_id, slot.name);
}

// ---------------------------------------------------------------------------
// Internal: find best candidate to evict (lowest priority, oldest run)
// ---------------------------------------------------------------------------
static int find_eviction_candidate() {
    int best = -1;
    uint8_t worst_priority = MODEL_PRIORITY_TRIAGE;
    uint32_t oldest_run = UINT32_MAX;

    for (int i = 0; i < MODEL_MAX_SLOTS; i++) {
        ModelSlot& slot = s_slots[i];
        if (!slot.registered || !slot.loaded) continue;
        if (slot.priority == MODEL_PRIORITY_TRIAGE) continue; // never evict triage

        if (slot.priority > worst_priority ||
            (slot.priority == worst_priority && slot.last_run_ms < oldest_run)) {
            worst_priority = slot.priority;
            oldest_run = slot.last_run_ms;
            best = i;
        }
    }
    return best;
}

// ===========================================================================
// Public API
// ===========================================================================

int model_manager_init() {
    memset(s_slots, 0, sizeof(s_slots));
    s_registered_count = 0;

    for (int i = 0; i < MODEL_MAX_SLOTS; i++) {
        s_interp_mem[i] = nullptr;
    }

    Serial.printf("[%s] Initialized. Max %d slots, PSRAM budget %u KB\n",
                  TAG, MODEL_MAX_SLOTS, MODEL_PSRAM_BUDGET / 1024);

    return 0;
}

// ---------------------------------------------------------------------------
int model_manager_register(const char* name,
                           const char* spiffs_path,
                           uint16_t input_size,
                           uint8_t  output_classes,
                           ModelPriority priority) {
    if (s_registered_count >= MODEL_MAX_SLOTS) {
        Serial.printf("[%s] Registry full (%d slots)\n", TAG, MODEL_MAX_SLOTS);
        return -1;
    }

    // Find first empty slot
    int slot_id = -1;
    for (int i = 0; i < MODEL_MAX_SLOTS; i++) {
        if (!s_slots[i].registered) {
            slot_id = i;
            break;
        }
    }
    if (slot_id < 0) return -1;

    ModelSlot& slot = s_slots[slot_id];
    memset(&slot, 0, sizeof(ModelSlot));

    strncpy(slot.name, name, MODEL_NAME_MAX_LEN - 1);
    slot.name[MODEL_NAME_MAX_LEN - 1] = '\0';
    strncpy(slot.spiffs_path, spiffs_path, sizeof(slot.spiffs_path) - 1);
    slot.spiffs_path[sizeof(slot.spiffs_path) - 1] = '\0';
    slot.input_size = input_size;
    slot.output_classes = output_classes;
    slot.priority = priority;
    slot.registered = true;
    slot.loaded = false;
    slot.inference_count = 0;
    slot.total_latency_us = 0;
    slot.last_run_ms = 0;

    s_registered_count++;

    Serial.printf("[%s] Registered slot %d: %s (%s) priority=%d in=%d out=%d\n",
                  TAG, slot_id, name, spiffs_path, priority, input_size, output_classes);

    // Auto-load triage models immediately
    if (priority == MODEL_PRIORITY_TRIAGE) {
        if (load_model_bytes(slot) && build_slot_interpreter((uint8_t)slot_id)) {
            Serial.printf("[%s] Triage model %s loaded at init\n", TAG, name);
        } else {
            Serial.printf("[%s] WARNING: Triage model %s failed to load\n", TAG, name);
        }
    }

    return slot_id;
}

// ---------------------------------------------------------------------------
ModelInferenceResult model_manager_infer(uint8_t slot_id,
                                         const float* input_data,
                                         uint16_t input_len) {
    ModelInferenceResult result;
    memset(&result, 0, sizeof(result));
    result.slot_id = slot_id;
    result.valid = false;

    if (slot_id >= MODEL_MAX_SLOTS || !s_slots[slot_id].registered) {
        return result;
    }

    ModelSlot& slot = s_slots[slot_id];

    // Page-in if not loaded
    if (!slot.loaded) {
        if (!model_manager_page_in(slot_id)) {
            Serial.printf("[%s] Cannot page in %s for inference\n", TAG, slot.name);
            return result;
        }
    }

    if (!slot.interpreter) {
        return result;
    }

    // Validate input length
    if (input_len != slot.input_size) {
        Serial.printf("[%s] %s: input size mismatch %d vs %d\n",
                      TAG, slot.name, input_len, slot.input_size);
        return result;
    }

    uint32_t start_us = micros();

    // Copy input data into input tensor
    TfLiteTensor* input_tensor = slot.interpreter->input(0);
    memcpy(input_tensor->data.f, input_data, input_len * sizeof(float));

    // Run inference
    TfLiteStatus invoke_status = slot.interpreter->Invoke();
    if (invoke_status != kTfLiteOk) {
        Serial.printf("[%s] %s: Invoke() failed\n", TAG, slot.name);
        return result;
    }

    uint32_t elapsed_us = micros() - start_us;

    // Read output
    TfLiteTensor* output_tensor = slot.interpreter->output(0);
    float* probs = output_tensor->data.f;

    float max_prob = -1.0f;
    int max_idx = 0;
    uint8_t num_out = slot.output_classes;
    if (num_out > MODEL_MAX_OUTPUT_CLASSES) num_out = MODEL_MAX_OUTPUT_CLASSES;

    for (int i = 0; i < num_out; i++) {
        result.probabilities[i] = probs[i];
        if (probs[i] > max_prob) {
            max_prob = probs[i];
            max_idx = i;
        }
    }

    result.class_label = (uint8_t)max_idx;
    result.confidence = max_prob;
    result.num_classes = num_out;
    result.latency_us = elapsed_us;
    result.valid = true;

    // Update stats
    slot.inference_count++;
    slot.total_latency_us += elapsed_us;
    slot.last_run_ms = millis();

    return result;
}

// ---------------------------------------------------------------------------
int model_manager_infer_all_resident(const float* input_data,
                                     uint16_t input_len,
                                     ModelInferenceResult* results,
                                     int max_results) {
    int count = 0;
    for (int i = 0; i < MODEL_MAX_SLOTS && count < max_results; i++) {
        if (!s_slots[i].registered || !s_slots[i].loaded) continue;

        results[count] = model_manager_infer((uint8_t)i, input_data, input_len);
        if (results[count].valid) {
            count++;
        }
    }
    return count;
}

// ---------------------------------------------------------------------------
bool model_manager_page_in(uint8_t slot_id) {
    if (slot_id >= MODEL_MAX_SLOTS || !s_slots[slot_id].registered) return false;

    ModelSlot& slot = s_slots[slot_id];
    if (slot.loaded) return true; // already loaded

    // Check if we have PSRAM budget
    uint32_t current_used = model_manager_psram_used();
    // Estimate needed: model file + arena
    uint32_t estimated_need = MODEL_ARENA_SIZE_PER_SLOT + OTA_MODEL_SLOT_SIZE;

    if (current_used + estimated_need > MODEL_PSRAM_BUDGET) {
        // Try eviction
        int victim = find_eviction_candidate();
        if (victim >= 0 && victim != slot_id) {
            Serial.printf("[%s] Evicting slot %d (%s) to make room for %s\n",
                          TAG, victim, s_slots[victim].name, slot.name);
            unload_slot((uint8_t)victim);
        } else {
            Serial.printf("[%s] Cannot page in %s — PSRAM budget exhausted\n",
                          TAG, slot.name);
            return false;
        }
    }

    // Load from SPIFFS
    if (!load_model_bytes(slot)) return false;
    if (!build_slot_interpreter(slot_id)) {
        // Failed to build, free the loaded bytes
        if (slot.model_buf) {
            heap_caps_free(slot.model_buf);
            slot.model_buf = nullptr;
        }
        return false;
    }

    return true;
}

// ---------------------------------------------------------------------------
void model_manager_page_out(uint8_t slot_id) {
    if (slot_id >= MODEL_MAX_SLOTS) return;
    ModelSlot& slot = s_slots[slot_id];
    if (slot.priority == MODEL_PRIORITY_TRIAGE) {
        Serial.printf("[%s] Refusing to page out triage model %s\n", TAG, slot.name);
        return;
    }
    unload_slot(slot_id);
}

// ---------------------------------------------------------------------------
bool model_manager_is_loaded(uint8_t slot_id) {
    if (slot_id >= MODEL_MAX_SLOTS) return false;
    return s_slots[slot_id].registered && s_slots[slot_id].loaded;
}

// ---------------------------------------------------------------------------
bool model_manager_get_stats(uint8_t slot_id, ModelStats& stats) {
    if (slot_id >= MODEL_MAX_SLOTS || !s_slots[slot_id].registered) return false;

    const ModelSlot& slot = s_slots[slot_id];
    stats.name = slot.name;
    stats.slot_id = slot_id;
    stats.priority = slot.priority;
    stats.loaded = slot.loaded;
    stats.inference_count = slot.inference_count;
    stats.avg_latency_us = (slot.inference_count > 0)
        ? (uint32_t)(slot.total_latency_us / slot.inference_count)
        : 0;
    stats.model_size_bytes = (uint32_t)slot.model_size;
    stats.arena_used_bytes = slot.arena_used;
    return true;
}

// ---------------------------------------------------------------------------
uint32_t model_manager_psram_used() {
    uint32_t total = 0;
    for (int i = 0; i < MODEL_MAX_SLOTS; i++) {
        if (!s_slots[i].registered || !s_slots[i].loaded) continue;
        total += (uint32_t)s_slots[i].model_size;
        total += s_slots[i].arena_alloc;
        if (s_interp_mem[i]) total += INTERPRETER_MEM_SIZE;
    }
    return total;
}

// ---------------------------------------------------------------------------
uint8_t model_manager_count() {
    return s_registered_count;
}

bool model_manager_commit_slot(uint8_t slot_id, uint32_t expected_crc32) {
    if (slot_id >= MODEL_MAX_SLOTS) {
        Serial.printf("[%s] commit_slot: invalid slot %u\n", TAG, (unsigned)slot_id);
        return false;
    }
    ModelSlot& slot = s_slots[slot_id];
    if (!slot.registered) {
        Serial.printf("[%s] commit_slot: slot %u not registered\n", TAG, (unsigned)slot_id);
        return false;
    }
    if (!slot.loaded || !slot.model_buf || slot.model_size == 0) {
        Serial.printf("[%s] commit_slot: slot %u has no loaded bytes\n",
                      TAG, (unsigned)slot_id);
        return false;
    }

    // If caller passes 0 as expected_crc32 we skip the CRC check (useful during
    // bring-up or when the controller doesn't know the CRC yet).
    if (expected_crc32 != 0) {
        uint32_t actual = model_manager_crc32_bzip2(slot.model_buf, slot.model_size);
        if (actual != expected_crc32) {
            Serial.printf("[%s] commit_slot: CRC mismatch slot %u want=0x%08lX got=0x%08lX\n",
                          TAG, (unsigned)slot_id,
                          (unsigned long)expected_crc32, (unsigned long)actual);
            return false;
        }
    }

    // Rebuild interpreter on the freshly-staged bytes. If rebuild succeeds,
    // the slot is committed active; otherwise the slot is marked not loaded
    // so the next infer call will attempt to page it in cleanly.
    if (!build_slot_interpreter(slot_id)) {
        Serial.printf("[%s] commit_slot: interpreter rebuild failed for slot %u\n",
                      TAG, (unsigned)slot_id);
        slot.loaded = false;
        return false;
    }

    Serial.printf("[%s] commit_slot: slot %u committed (%u bytes, CRC 0x%08lX)\n",
                  TAG, (unsigned)slot_id, (unsigned)slot.model_size,
                  (unsigned long)expected_crc32);
    return true;
}

// ---------------------------------------------------------------------------
void model_manager_deinit() {
    for (int i = 0; i < MODEL_MAX_SLOTS; i++) {
        unload_slot((uint8_t)i);
        if (s_interp_mem[i]) {
            heap_caps_free(s_interp_mem[i]);
            s_interp_mem[i] = nullptr;
        }
    }
    memset(s_slots, 0, sizeof(s_slots));
    s_registered_count = 0;
}
#endif // PDL_DISABLE_TFLITE
