/*  cactus_fullmode.cpp — Full Mode .cactus runtime for Heltec WiFi LoRa V4
 *
 *  Binary format (little-endian, per weights.py):
 *    [0..7]   magic   = "CACTFM01"  (8 bytes)
 *    [8..11]  u32     header_json_len
 *    [12..]   UTF-8 JSON of header_json_len bytes
 *    [...]    weight blobs in arch order:
 *               linear fp32: W (in*out floats), b (out floats)
 *               linear int8: W (in*out int8_t),  b (out floats)
 *               relu / softmax: no blob
 *
 *  All dynamic memory is heap_caps_malloc(MALLOC_CAP_SPIRAM).
 *  No STL in hot path; JSON parsing is a minimal hand-written scanner.
 */

#include "cactus_fullmode.h"
#include "fullmode_kernels.h"

#include <Arduino.h>
#include <LittleFS.h>
#include <esp_heap_caps.h>
#include <string.h>
#include <stdlib.h>
#include <stdio.h>
#include <math.h>

// ── Logging tag ───────────────────────────────────────────────────────────
static const char* TAG = "CactusFM";

// ── Magic ─────────────────────────────────────────────────────────────────
static const uint8_t CACTUS_MAGIC[8] = { 'C','A','C','T','F','M','0','1' };

// ── Layer type enum ───────────────────────────────────────────────────────
enum LayerType : uint8_t {
    LT_LINEAR_FP32 = 0,
    LT_LINEAR_INT8 = 1,
    LT_RELU        = 2,
    LT_SOFTMAX     = 3,
    LT_UNKNOWN     = 0xFF
};

// ── Per-layer descriptor ──────────────────────────────────────────────────
struct LayerDesc {
    LayerType type;
    int       in_dim;
    int       out_dim;
    float     w_scale;   // int8 only; ignored for fp32
    // Pointers into PSRAM pool
    void*     W;         // float* (fp32) or int8_t* (int8)
    float*    b;         // always float*
};

// ── Per-slot runtime ──────────────────────────────────────────────────────
struct CactusSlot {
    bool        used;
    int         n_layers;
    int         input_dim;
    int         output_dim;
    LayerDesc   layers[CACTUS_FM_MAX_LAYERS];
    // PSRAM pool for this slot (one contiguous allocation, sliced internally)
    uint8_t*    psram_pool;
    size_t      psram_used;
    size_t      psram_cap;
    // Scratch activation buffer — max inter-layer dimension, allocated once
    float*      scratch_a;    // size = max_intermediate_dim
    float*      scratch_b;    // second ping-pong buffer
    int         scratch_dim;
};

static CactusSlot   s_slots[CACTUS_FM_MAX_SLOTS];
static bool         s_initialised = false;

// ── Small helpers ─────────────────────────────────────────────────────────

// Read a u32 little-endian from a byte pointer
static inline uint32_t read_u32_le(const uint8_t* p) {
    return (uint32_t)p[0]
         | ((uint32_t)p[1] << 8)
         | ((uint32_t)p[2] << 16)
         | ((uint32_t)p[3] << 24);
}

// Allocate from the slot's PSRAM pool; returns nullptr on OOM
static void* slot_alloc(CactusSlot& s, size_t bytes) {
    // Align to 4 bytes
    size_t aligned = (bytes + 3) & ~(size_t)3;
    if (s.psram_used + aligned > s.psram_cap) return nullptr;
    void* ptr = s.psram_pool + s.psram_used;
    s.psram_used += aligned;
    return ptr;
}

// ── Minimal JSON key scanner ──────────────────────────────────────────────
// These helpers parse just enough JSON to decode the cactus header.
// The header is small (<2 KB); we work directly on the string in SRAM.

// Skip whitespace
static const char* json_skip_ws(const char* p) {
    while (*p == ' ' || *p == '\t' || *p == '\r' || *p == '\n') p++;
    return p;
}

// Match a literal string; returns pointer after it, or nullptr
static const char* json_expect(const char* p, const char* lit) {
    size_t n = strlen(lit);
    if (strncmp(p, lit, n) == 0) return p + n;
    return nullptr;
}

// Parse a JSON string value (returns pointer to first char after closing ")
// Writes up to buf_size-1 chars into buf.
static const char* json_parse_string(const char* p, char* buf, size_t buf_size) {
    p = json_skip_ws(p);
    if (*p != '"') return nullptr;
    p++;
    size_t i = 0;
    while (*p && *p != '"') {
        if (*p == '\\') { p++; if (!*p) return nullptr; }
        if (i + 1 < buf_size) buf[i++] = *p;
        p++;
    }
    buf[i] = '\0';
    if (*p == '"') p++;
    return p;
}

// Parse a JSON integer (or float, truncated)
static const char* json_parse_int(const char* p, int* out) {
    p = json_skip_ws(p);
    char* end;
    long v = strtol(p, &end, 10);
    if (end == p) return nullptr;
    *out = (int)v;
    return end;
}

// Parse a JSON float
static const char* json_parse_float(const char* p, float* out) {
    p = json_skip_ws(p);
    char* end;
    float v = strtof(p, &end);
    if (end == p) return nullptr;
    *out = v;
    return end;
}

// Find the first occurrence of a key in the current JSON object level.
// Returns pointer to the value (after the colon), or nullptr.
// Caution: this is a simple linear scan; duplicate keys return the first hit.
static const char* json_find_key(const char* obj_start, const char* key) {
    char needle[64];
    snprintf(needle, sizeof(needle), "\"%s\"", key);
    const char* p = strstr(obj_start, needle);
    if (!p) return nullptr;
    p += strlen(needle);
    p = json_skip_ws(p);
    if (*p != ':') return nullptr;
    p++;
    return json_skip_ws(p);
}

// ── SPIFFS mount guard ────────────────────────────────────────────────────
static bool s_spiffs_ok = false;
static bool ensure_spiffs() {
    if (s_spiffs_ok) return true;
    if (!LittleFS.begin(true)) {
        Serial.printf("[%s] SPIFFS mount failed\n", TAG);
        return false;
    }
    s_spiffs_ok = true;
    return true;
}

// ── Public: cactus_fm_init ────────────────────────────────────────────────
bool cactus_fm_init(void) {
    if (s_initialised) return true;
    memset(s_slots, 0, sizeof(s_slots));
    s_initialised = true;
    Serial.printf("[%s] init complete (max %d slots, %u KB each)\n",
                  TAG, CACTUS_FM_MAX_SLOTS,
                  (unsigned)(CACTUS_FM_PSRAM_PER_SLOT / 1024));
    return true;
}

// ── Public: cactus_fm_load ────────────────────────────────────────────────
int cactus_fm_load(const char* spiffs_path) {
    if (!s_initialised) {
        Serial.printf("[%s] cactus_fm_init() not called\n", TAG);
        return -1;
    }

    // Find a free slot
    int sid = -1;
    for (int i = 0; i < CACTUS_FM_MAX_SLOTS; i++) {
        if (!s_slots[i].used) { sid = i; break; }
    }
    if (sid < 0) {
        Serial.printf("[%s] All slots full\n", TAG);
        return -1;
    }

    CactusSlot& slot = s_slots[sid];
    memset(&slot, 0, sizeof(CactusSlot));

    // ---- Allocate PSRAM pool ----
    slot.psram_pool = (uint8_t*)heap_caps_malloc(CACTUS_FM_PSRAM_PER_SLOT,
                                                  MALLOC_CAP_SPIRAM);
    if (!slot.psram_pool) {
        Serial.printf("[%s] PSRAM alloc failed (%u KB) for %s\n",
                      TAG, (unsigned)(CACTUS_FM_PSRAM_PER_SLOT / 1024), spiffs_path);
        return -1;
    }
    slot.psram_cap  = CACTUS_FM_PSRAM_PER_SLOT;
    slot.psram_used = 0;

    // ---- Open SPIFFS file ----
    if (!ensure_spiffs()) goto fail;

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

        // ---- Read and validate magic ----
        uint8_t magic_buf[8];
        if (f.read(magic_buf, 8) != 8 ||
            memcmp(magic_buf, CACTUS_MAGIC, 8) != 0) {
            Serial.printf("[%s] %s: bad magic\n", TAG, spiffs_path);
            f.close();
            goto fail;
        }

        // ---- Read header JSON length ----
        uint8_t hlen_buf[4];
        if (f.read(hlen_buf, 4) != 4) {
            Serial.printf("[%s] %s: truncated header len\n", TAG, spiffs_path);
            f.close();
            goto fail;
        }
        uint32_t hlen = read_u32_le(hlen_buf);
        if (hlen >= CACTUS_FM_MAX_HEADER_JSON) {
            Serial.printf("[%s] %s: header JSON too large (%lu)\n",
                          TAG, spiffs_path, (unsigned long)hlen);
            f.close();
            goto fail;
        }

        // ---- Read header JSON into SRAM stack buffer ----
        char hdr[CACTUS_FM_MAX_HEADER_JSON];
        if (f.read((uint8_t*)hdr, hlen) != (int)hlen) {
            Serial.printf("[%s] %s: short header read\n", TAG, spiffs_path);
            f.close();
            goto fail;
        }
        hdr[hlen] = '\0';

        // ---- Parse input_shape and output_shape ----
        {
            const char* p = json_find_key(hdr, "input_shape");
            if (!p) goto bad_header;
            // value is [N]
            if (*p == '[') { p++; p = json_skip_ws(p); }
            int idim = 0;
            p = json_parse_int(p, &idim);
            if (!p) goto bad_header;
            slot.input_dim = idim;

            p = json_find_key(hdr, "output_shape");
            if (!p) goto bad_header;
            if (*p == '[') { p++; p = json_skip_ws(p); }
            int odim = 0;
            p = json_parse_int(p, &odim);
            if (!p) goto bad_header;
            slot.output_dim = odim;
        }

        // ---- Parse arch array ----
        // Locate "arch":[ and walk through each layer object
        {
            const char* arch_start = strstr(hdr, "\"arch\"");
            if (!arch_start) goto bad_header;
            arch_start = strchr(arch_start, '[');
            if (!arch_start) goto bad_header;
            arch_start++; // past '['

            const char* p = arch_start;
            int lcount = 0;
            int max_intermediate = 0;

            while (*p && lcount < CACTUS_FM_MAX_LAYERS) {
                p = json_skip_ws(p);
                if (*p == ']') break; // end of arch array
                if (*p != '{') { p++; continue; } // skip commas, etc.

                // Find matching closing brace
                const char* obj_start = p;
                int depth = 0;
                const char* obj_end = p;
                while (*obj_end) {
                    if (*obj_end == '{') depth++;
                    else if (*obj_end == '}') { depth--; if (depth == 0) { obj_end++; break; } }
                    obj_end++;
                }
                // obj_start..obj_end is the layer object

                // Copy to temp buffer for key scanning
                char obj[512];
                size_t olen = (size_t)(obj_end - obj_start);
                if (olen >= sizeof(obj)) olen = sizeof(obj) - 1;
                memcpy(obj, obj_start, olen);
                obj[olen] = '\0';

                // Parse "type"
                char type_str[16] = "";
                const char* tp = json_find_key(obj, "type");
                if (tp) {
                    char tmp[16];
                    json_parse_string(tp - 1, tmp, sizeof(tmp)); // tp points past ":"
                    // tp is already past colon; parse_string needs to start at '"'
                    // Re-approach: find "type" value directly
                    const char* tv = json_find_key(obj, "type");
                    if (tv && *tv == '"') {
                        size_t i = 0;
                        tv++; // skip opening "
                        while (*tv && *tv != '"' && i < sizeof(type_str)-1)
                            type_str[i++] = *tv++;
                        type_str[i] = '\0';
                    }
                }

                LayerDesc& ld = slot.layers[lcount];
                memset(&ld, 0, sizeof(LayerDesc));

                if (strcmp(type_str, "relu") == 0) {
                    ld.type = LT_RELU;
                    lcount++;
                    p = obj_end;
                    continue;
                }
                if (strcmp(type_str, "softmax") == 0) {
                    ld.type = LT_SOFTMAX;
                    lcount++;
                    p = obj_end;
                    continue;
                }
                if (strcmp(type_str, "linear") == 0) {
                    int in_d = 0, out_d = 0;
                    char dtype[8] = "fp32";
                    float wscale = 1.0f;

                    const char* vp;
                    vp = json_find_key(obj, "in");
                    if (vp) json_parse_int(vp, &in_d);
                    vp = json_find_key(obj, "out");
                    if (vp) json_parse_int(vp, &out_d);

                    // dtype
                    vp = json_find_key(obj, "dtype");
                    if (vp && *vp == '"') {
                        size_t i = 0; vp++;
                        while (*vp && *vp != '"' && i < sizeof(dtype)-1)
                            dtype[i++] = *vp++;
                        dtype[i] = '\0';
                    }

                    bool is_int8 = (strcmp(dtype, "int8") == 0);

                    if (is_int8) {
                        vp = json_find_key(obj, "w_scale");
                        if (vp) json_parse_float(vp, &wscale);
                    }

                    // Allocate W and b in PSRAM pool
                    size_t w_bytes = is_int8
                        ? (size_t)(in_d * out_d)         // int8
                        : (size_t)(in_d * out_d) * 4;    // float32
                    size_t b_bytes = (size_t)out_d * 4;  // always float32

                    ld.W = slot_alloc(slot, w_bytes);
                    ld.b = (float*)slot_alloc(slot, b_bytes);

                    if (!ld.W || !ld.b) {
                        Serial.printf("[%s] PSRAM pool OOM at layer %d\n", TAG, lcount);
                        f.close();
                        goto fail;
                    }

                    // Read weight blob from file
                    if ((size_t)f.read((uint8_t*)ld.W, w_bytes) != w_bytes) {
                        Serial.printf("[%s] Short W read layer %d\n", TAG, lcount);
                        f.close();
                        goto fail;
                    }
                    if ((size_t)f.read((uint8_t*)ld.b, b_bytes) != b_bytes) {
                        Serial.printf("[%s] Short b read layer %d\n", TAG, lcount);
                        f.close();
                        goto fail;
                    }

                    ld.type    = is_int8 ? LT_LINEAR_INT8 : LT_LINEAR_FP32;
                    ld.in_dim  = in_d;
                    ld.out_dim = out_d;
                    ld.w_scale = wscale;

                    if (out_d > max_intermediate) max_intermediate = out_d;
                    lcount++;
                    p = obj_end;
                    continue;
                }

                // Unknown layer type — skip
                Serial.printf("[%s] Unknown layer type '%s' — skipped\n", TAG, type_str);
                p = obj_end;
            }

            slot.n_layers = lcount;

            // ---- Allocate scratch activation ping-pong buffers in PSRAM ----
            // We need max(input_dim, max_intermediate) for ping-pong to be safe
            int scratch_needed = slot.input_dim;
            if (max_intermediate > scratch_needed) scratch_needed = max_intermediate;
            if (slot.output_dim  > scratch_needed) scratch_needed = slot.output_dim;
            slot.scratch_dim = scratch_needed;

            slot.scratch_a = (float*)slot_alloc(slot, (size_t)scratch_needed * 4);
            slot.scratch_b = (float*)slot_alloc(slot, (size_t)scratch_needed * 4);

            if (!slot.scratch_a || !slot.scratch_b) {
                Serial.printf("[%s] PSRAM OOM for scratch buffers\n", TAG);
                f.close();
                goto fail;
            }
        }

        f.close();
        slot.used = true;

        Serial.printf("[%s] Loaded slot %d: %s  in=%d out=%d layers=%d  PSRAM %u / %u KB\n",
                      TAG, sid, spiffs_path,
                      slot.input_dim, slot.output_dim, slot.n_layers,
                      (unsigned)(slot.psram_used / 1024),
                      (unsigned)(slot.psram_cap / 1024));
        return sid;

bad_header:
        Serial.printf("[%s] %s: malformed header JSON\n", TAG, spiffs_path);
        f.close();
        goto fail;
    }

fail:
    if (slot.psram_pool) {
        heap_caps_free(slot.psram_pool);
        slot.psram_pool = nullptr;
    }
    memset(&slot, 0, sizeof(CactusSlot));
    return -1;
}

// ── Public: cactus_fm_infer ───────────────────────────────────────────────
bool cactus_fm_infer(int slot_id,
                     const float* in,  size_t in_len,
                     float*       out, size_t out_len)
{
    if (slot_id < 0 || slot_id >= CACTUS_FM_MAX_SLOTS) return false;
    CactusSlot& slot = s_slots[slot_id];
    if (!slot.used) return false;

    if ((int)in_len != slot.input_dim) {
        Serial.printf("[%s] slot %d: in_len %u != %d\n",
                      TAG, slot_id, (unsigned)in_len, slot.input_dim);
        return false;
    }
    if ((int)out_len < slot.output_dim) {
        Serial.printf("[%s] slot %d: out_len %u < %d\n",
                      TAG, slot_id, (unsigned)out_len, slot.output_dim);
        return false;
    }

    // Copy input into scratch_a (ping)
    memcpy(slot.scratch_a, in, (size_t)slot.input_dim * sizeof(float));

    // Walk layers: ping = scratch_a, pong = scratch_b, swap after each linear
    float* ping = slot.scratch_a;
    float* pong = slot.scratch_b;
    // Track current row width
    int cur_dim = slot.input_dim;

    for (int l = 0; l < slot.n_layers; l++) {
        LayerDesc& ld = slot.layers[l];
        switch (ld.type) {
            case LT_LINEAR_FP32:
                cactus_fullmode_gemm_fp32(
                    ping, (const float*)ld.W, ld.b,
                    pong, /*N=*/1, ld.in_dim, ld.out_dim);
                cur_dim = ld.out_dim;
                { float* tmp = ping; ping = pong; pong = tmp; }
                break;

            case LT_LINEAR_INT8:
                cactus_fullmode_gemm_int8(
                    ping, (const int8_t*)ld.W, ld.w_scale, ld.b,
                    pong, /*N=*/1, ld.in_dim, ld.out_dim);
                cur_dim = ld.out_dim;
                { float* tmp = ping; ping = pong; pong = tmp; }
                break;

            case LT_RELU:
                cactus_fullmode_relu_fp32(ping, cur_dim);
                break;

            case LT_SOFTMAX:
                cactus_fullmode_softmax_fp32(ping, pong, /*N=*/1, cur_dim);
                { float* tmp = ping; ping = pong; pong = tmp; }
                break;

            default:
                break;
        }
    }

    // Copy result to caller's output buffer
    memcpy(out, ping, (size_t)slot.output_dim * sizeof(float));
    return true;
}

// ── Public: cactus_fm_unload ──────────────────────────────────────────────
bool cactus_fm_unload(int slot_id) {
    if (slot_id < 0 || slot_id >= CACTUS_FM_MAX_SLOTS) return false;
    CactusSlot& slot = s_slots[slot_id];
    if (!slot.used) return false;

    heap_caps_free(slot.psram_pool);
    memset(&slot, 0, sizeof(CactusSlot));

    Serial.printf("[%s] slot %d unloaded\n", TAG, slot_id);
    return true;
}

// ── Public: cactus_fm_psram_used ─────────────────────────────────────────
size_t cactus_fm_psram_used(void) {
    size_t total = 0;
    for (int i = 0; i < CACTUS_FM_MAX_SLOTS; i++) {
        if (s_slots[i].used) total += s_slots[i].psram_used;
    }
    return total;
}
