#include "rf_classifier.h"

#include <Arduino.h>
#include <LittleFS.h>
#include <esp_heap_caps.h>

#ifdef PDL_DISABLE_TFLITE
// ============================================================================
// TFLite stub mode — TFLite-Micro lib disabled (toolchain incompat).
// Classification responsibility shifts to cactus.full_mode runtime.
// ============================================================================
static const char* TAG = "RFClassifier";
static float s_thr[9] = {
    CONFIDENCE_THRESHOLD, CONFIDENCE_THRESHOLD, CONFIDENCE_THRESHOLD,
    CONFIDENCE_THRESHOLD, CONFIDENCE_THRESHOLD, CONFIDENCE_THRESHOLD,
    CONFIDENCE_THRESHOLD, CONFIDENCE_THRESHOLD, CONFIDENCE_THRESHOLD};
bool  classifier_init() { Serial.printf("[%s] TFLite disabled — passthrough\n", TAG); return false; }
bool  classifier_infer(const float*, ClassificationResult&) { return false; }
bool  classifier_swap_model(const uint8_t*, size_t) { return false; }
bool  classifier_set_confidence_threshold(uint8_t id, float t) {
    if (id >= 9 || t < 0.0f || t > 1.0f) return false;
    s_thr[id] = t; return true;
}
float classifier_get_confidence_threshold(uint8_t id) { return id < 9 ? s_thr[id] : CONFIDENCE_THRESHOLD; }
void  classifier_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 = "RFClassifier";

#ifndef PDL_DISABLE_TFLITE
// TFLite-Micro state
static const tflite::Model*          s_model       = nullptr;
static tflite::MicroInterpreter*     s_interpreter = nullptr;
static tflite::AllOpsResolver        s_resolver;
static uint8_t*                      s_arena       = nullptr;
static uint8_t*                      s_model_buf   = nullptr;
static size_t                        s_model_size  = 0;
#endif

// Runtime-tunable confidence thresholds. Index 0 = legacy classifier.
// Higher indices mirror model_manager slot IDs (not enforced here; stored
// only for command_handler SET_THRESHOLD acknowledgement).
#define RF_CLASSIFIER_MAX_THRESHOLDS 9  // slot 0 (legacy) + up to 8 model_mgr slots
static float s_confidence_thresholds[RF_CLASSIFIER_MAX_THRESHOLDS] = {
    CONFIDENCE_THRESHOLD, CONFIDENCE_THRESHOLD, CONFIDENCE_THRESHOLD,
    CONFIDENCE_THRESHOLD, CONFIDENCE_THRESHOLD, CONFIDENCE_THRESHOLD,
    CONFIDENCE_THRESHOLD, CONFIDENCE_THRESHOLD, CONFIDENCE_THRESHOLD
};

// ---------------------------------------------------------------------------
// Load .tflite model file from SPIFFS into PSRAM
// ---------------------------------------------------------------------------
static bool load_model_from_spiffs(const char* path) {
    if (!LittleFS.begin(true)) {
        Serial.printf("[%s] SPIFFS mount failed\n", TAG);
        return false;
    }

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

    s_model_size = f.size();
    if (s_model_size == 0) {
        Serial.printf("[%s] Model file is empty\n", TAG);
        f.close();
        return false;
    }

    // Allocate model buffer in PSRAM
    s_model_buf = (uint8_t*)heap_caps_malloc(s_model_size, MALLOC_CAP_SPIRAM);
    if (!s_model_buf) {
        Serial.printf("[%s] Failed to allocate %u bytes in PSRAM for model\n", TAG, s_model_size);
        f.close();
        return false;
    }

    size_t bytes_read = f.read(s_model_buf, s_model_size);
    f.close();

    if (bytes_read != s_model_size) {
        Serial.printf("[%s] Model read incomplete: %u / %u\n", TAG, bytes_read, s_model_size);
        heap_caps_free(s_model_buf);
        s_model_buf = nullptr;
        return false;
    }

    Serial.printf("[%s] Model loaded from %s (%u bytes)\n", TAG, path, s_model_size);
    return true;
}

// ---------------------------------------------------------------------------
// Build interpreter from model buffer already in memory
// ---------------------------------------------------------------------------
static bool build_interpreter() {
    s_model = tflite::GetModel(s_model_buf);
    if (s_model->version() != TFLITE_SCHEMA_VERSION) {
        Serial.printf("[%s] Model schema version mismatch: %lu vs %d\n",
                      TAG, s_model->version(), TFLITE_SCHEMA_VERSION);
        return false;
    }

    // Allocate tensor arena in PSRAM
    if (!s_arena) {
        s_arena = (uint8_t*)heap_caps_malloc(TFLITE_ARENA_SIZE, MALLOC_CAP_SPIRAM);
        if (!s_arena) {
            Serial.printf("[%s] Failed to allocate %d bytes arena in PSRAM\n", TAG, TFLITE_ARENA_SIZE);
            return false;
        }
    }

    // Create interpreter (modern MicroTFLite: 4-arg, no ErrorReporter)
    static tflite::MicroInterpreter interpreter(
        s_model, s_resolver, s_arena, TFLITE_ARENA_SIZE);
    s_interpreter = &interpreter;

    TfLiteStatus alloc_status = s_interpreter->AllocateTensors();
    if (alloc_status != kTfLiteOk) {
        Serial.printf("[%s] AllocateTensors() failed\n", TAG);
        return false;
    }

    // Validate input tensor shape
    TfLiteTensor* input = s_interpreter->input(0);
    if (input->dims->size != 2 || input->dims->data[1] != SPECTRUM_BINS) {
        Serial.printf("[%s] Unexpected input shape: dims=%d, bins=%d\n",
                      TAG, input->dims->size, input->dims->data[1]);
        return false;
    }

    // Validate output tensor shape
    TfLiteTensor* output = s_interpreter->output(0);
    if (output->dims->data[output->dims->size - 1] != NUM_CLASSES) {
        Serial.printf("[%s] Unexpected output classes: %d\n",
                      TAG, output->dims->data[output->dims->size - 1]);
        return false;
    }

    Serial.printf("[%s] Interpreter ready. Arena used: %u / %d bytes\n",
                  TAG, s_interpreter->arena_used_bytes(), TFLITE_ARENA_SIZE);
    return true;
}

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

bool classifier_init() {
    if (!load_model_from_spiffs(MODEL_PATH)) {
        return false;
    }
    return build_interpreter();
}

bool classifier_infer(const float* spectrum, ClassificationResult& result) {
    if (!s_interpreter) {
        Serial.printf("[%s] Interpreter not initialized\n", TAG);
        return false;
    }

    // Copy spectrum into input tensor
    TfLiteTensor* input = s_interpreter->input(0);
    float* input_data = input->data.f;
    memcpy(input_data, spectrum, SPECTRUM_BINS * sizeof(float));

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

    // Read output (softmax probabilities)
    TfLiteTensor* output = s_interpreter->output(0);
    float* probs = output->data.f;

    // Find argmax
    float max_prob = -1.0f;
    int max_idx = 0;
    for (int i = 0; i < NUM_CLASSES; 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;

    return true;
}

bool classifier_swap_model(const uint8_t* model_data, size_t model_size) {
    Serial.printf("[%s] Hot-swapping model (%u bytes)\n", TAG, model_size);

    // Allocate new buffer in PSRAM
    uint8_t* new_buf = (uint8_t*)heap_caps_malloc(model_size, MALLOC_CAP_SPIRAM);
    if (!new_buf) {
        Serial.printf("[%s] Failed to allocate PSRAM for new model\n", TAG);
        return false;
    }
    memcpy(new_buf, model_data, model_size);

    // Free old model buffer
    if (s_model_buf) {
        heap_caps_free(s_model_buf);
    }
    s_model_buf = new_buf;
    s_model_size = model_size;

    // Rebuild interpreter with new model
    return build_interpreter();
}

bool classifier_set_confidence_threshold(uint8_t model_id, float threshold) {
    if (threshold < 0.0f || threshold > 1.0f) return false;
    if (model_id >= RF_CLASSIFIER_MAX_THRESHOLDS) return false;
    s_confidence_thresholds[model_id] = threshold;
    Serial.printf("[%s] Confidence threshold[%u] set to %.3f\n",
                  TAG, (unsigned)model_id, threshold);
    return true;
}

float classifier_get_confidence_threshold(uint8_t model_id) {
    if (model_id >= RF_CLASSIFIER_MAX_THRESHOLDS) return CONFIDENCE_THRESHOLD;
    return s_confidence_thresholds[model_id];
}

void classifier_deinit() {
    s_interpreter = nullptr;
    s_model = nullptr;

    if (s_arena) {
        heap_caps_free(s_arena);
        s_arena = nullptr;
    }
    if (s_model_buf) {
        heap_caps_free(s_model_buf);
        s_model_buf = nullptr;
    }
    s_model_size = 0;
}
#endif // PDL_DISABLE_TFLITE
