#ifndef MODEL_MANAGER_H
#define MODEL_MANAGER_H

#include <cstdint>
#include "config.h"

// ============================================================================
// Multi-Model Manager
//
// Manages up to MODEL_MAX_SLOTS TFLite-Micro models simultaneously.
// Triage models (drone classifier, rf_attack) stay always-resident in PSRAM.
// Specialist models (mavlink, lora_mesh, protocol) are paged in on-demand
// when a triage model flags an anomaly above PAGING_CONFIDENCE_THRESHOLD.
// ============================================================================

// Model priority levels — determines residency policy
enum ModelPriority : uint8_t {
    MODEL_PRIORITY_TRIAGE     = 0,  // Always resident (never paged out)
    MODEL_PRIORITY_SPECIALIST = 1,  // Loaded on-demand when triage triggers
    MODEL_PRIORITY_OPTIONAL   = 2,  // Lowest priority, first to evict
};

// Per-model inference result
struct ModelInferenceResult {
    uint8_t  slot_id;
    uint8_t  class_label;
    float    confidence;
    float    probabilities[MODEL_MAX_OUTPUT_CLASSES];
    uint8_t  num_classes;
    uint32_t latency_us;           // microseconds for this inference
    bool     valid;                // true if inference succeeded
};

// Per-model statistics (read-only view)
struct ModelStats {
    const char* name;
    uint8_t     slot_id;
    uint8_t     priority;
    bool        loaded;            // currently in PSRAM
    uint32_t    inference_count;
    uint32_t    avg_latency_us;    // running average
    uint32_t    model_size_bytes;
    uint32_t    arena_used_bytes;
};

// Initialize the model manager. Scans SPIFFS for registered models and
// loads triage models immediately. Returns number of models registered.
int model_manager_init();

// Register a model. Returns slot_id (0..MODEL_MAX_SLOTS-1) or -1 on failure.
int model_manager_register(const char* name,
                           const char* spiffs_path,
                           uint16_t input_size,
                           uint8_t  output_classes,
                           ModelPriority priority);

// Run inference on a specific model slot.
// input_data: float array of input_size elements.
// Returns result with valid=true on success.
ModelInferenceResult model_manager_infer(uint8_t slot_id,
                                         const float* input_data,
                                         uint16_t input_len);

// Run all currently-resident models against the same input spectrum.
// Fills results[] array (caller allocates MODEL_MAX_SLOTS entries).
// Returns number of valid results written.
int model_manager_infer_all_resident(const float* input_data,
                                     uint16_t input_len,
                                     ModelInferenceResult* results,
                                     int max_results);

// Request that a specialist model be paged in (loaded from SPIFFS).
// This may evict a lower-priority model if PSRAM budget is exceeded.
// Returns true if the model is now resident.
bool model_manager_page_in(uint8_t slot_id);

// Evict a specialist model from PSRAM to free memory.
void model_manager_page_out(uint8_t slot_id);

// Check if a slot is currently loaded in PSRAM.
bool model_manager_is_loaded(uint8_t slot_id);

// Get stats for a model slot. Returns false if slot_id is invalid.
bool model_manager_get_stats(uint8_t slot_id, ModelStats& stats);

// Get total PSRAM used by all loaded models + arenas.
uint32_t model_manager_psram_used();

// Get number of registered models.
uint8_t model_manager_count();

// Commit a previously-loaded model slot as the active model for its role.
// Used by the command_handler UPDATE_MODEL_SLOT command: the actual bytes
// are pushed via the OTA path (ota_adapter.cpp); this call validates the
// staged model's CRC against the expected value and marks it active.
// If the slot isn't loaded or the CRC doesn't match, returns false and
// leaves the previously-active model in place.
bool model_manager_commit_slot(uint8_t slot_id, uint32_t expected_crc32);

// Shutdown and free all resources.
void model_manager_deinit();

#endif // MODEL_MANAGER_H
