#pragma once

#include <LittleFS.h>
#include <ArduinoJson.h>
#include "config.h"

struct MacroInfo {
    char name[64];
    char labelColor[12];
    int nodeCount;
    bool hasImage;
};

struct SubInfo {
    char name[64];
    int nodeCount;
};

class MacroStorage {
public:
    int macroCount = 0;
    int order[MAX_MACROS];
    MacroInfo macros[MAX_MACROS];

    int subCount = 0;
    SubInfo subs[MAX_SUBROUTINES];

    bool begin() {
        if (!LittleFS.begin(true)) {
            return false;
        }
        memset(subs, 0, sizeof(subs));
        loadIndex();
        return true;
    }

    void loadIndex() {
        macroCount = 0;
        memset(order, 0, sizeof(order));

        if (!LittleFS.exists(CONFIG_PATH)) {
            saveIndex();
            return;
        }

        File f = LittleFS.open(CONFIG_PATH, "r");
        if (!f) return;

        JsonDocument doc;
        if (deserializeJson(doc, f) != DeserializationError::Ok) {
            f.close();
            return;
        }
        f.close();

        macroCount = doc["count"] | 0;
        JsonArray orderArr = doc["order"].as<JsonArray>();
        for (int i = 0; i < macroCount && i < MAX_MACROS; i++) {
            order[i] = orderArr[i] | i;
        }

        for (int i = 0; i < macroCount; i++) {
            loadMacroMeta(order[i]);
        }
    }

    void saveIndex() {
        JsonDocument doc;
        doc["count"] = macroCount;
        JsonArray orderArr = doc["order"].to<JsonArray>();
        for (int i = 0; i < macroCount; i++) {
            orderArr.add(order[i]);
        }

        File f = LittleFS.open(CONFIG_PATH, "w");
        if (f) {
            serializeJson(doc, f);
            f.close();
        }
    }

    void loadMacroMeta(int slot) {
        if (slot < 0 || slot >= MAX_MACROS) return;

        char path[48];
        snprintf(path, sizeof(path), "/m%d/meta.json", slot);

        MacroInfo& info = macros[slot];
        memset(&info, 0, sizeof(MacroInfo));
        strcpy(info.name, "Unnamed");
        strcpy(info.labelColor, "white");

        if (!LittleFS.exists(path)) return;

        File f = LittleFS.open(path, "r");
        if (!f) return;

        JsonDocument doc;
        if (deserializeJson(doc, f) == DeserializationError::Ok) {
            strlcpy(info.name, doc["name"] | "Unnamed", sizeof(info.name));
            strlcpy(info.labelColor, doc["label_color"] | "white", sizeof(info.labelColor));
            info.nodeCount = doc["nodes"] | 0;
        }
        f.close();

        snprintf(path, sizeof(path), "/m%d/icon.raw", slot);
        info.hasImage = LittleFS.exists(path);
    }

    bool beginMacroWrite(int slot, const char* name, int nodeCount, const char* labelColor = "white") {
        if (slot < 0 || slot >= MAX_MACROS) return false;

        char dir[16];
        snprintf(dir, sizeof(dir), "/m%d", slot);
        LittleFS.mkdir(dir);

        char path[48];
        snprintf(path, sizeof(path), "/m%d/meta.json", slot);
        File f = LittleFS.open(path, "w");
        if (!f) return false;

        JsonDocument doc;
        doc["name"] = name;
        doc["label_color"] = labelColor;
        doc["nodes"] = nodeCount;
        serializeJson(doc, f);
        f.close();

        // Clear existing nodes file
        snprintf(path, sizeof(path), "/m%d/nodes.json", slot);
        File nf = LittleFS.open(path, "w");
        if (nf) {
            nf.print("[");
            nf.close();
        }

        strlcpy(macros[slot].name, name, sizeof(macros[slot].name));
        strlcpy(macros[slot].labelColor, labelColor, sizeof(macros[slot].labelColor));
        macros[slot].nodeCount = nodeCount;

        return true;
    }

    bool writeImageData(int slot, uint8_t* data, size_t len) {
        char path[32];
        snprintf(path, sizeof(path), "/m%d/icon.raw", slot);
        File f = LittleFS.open(path, "w");
        if (!f) return false;
        size_t written = f.write(data, len);
        f.close();
        macros[slot].hasImage = (written == len);
        return macros[slot].hasImage;
    }

    bool writeImageChunk(int slot, uint8_t* data, size_t len, bool first) {
        char path[32];
        snprintf(path, sizeof(path), "/m%d/icon.raw", slot);
        File f = LittleFS.open(path, first ? "w" : "a");
        if (!f) return false;
        f.write(data, len);
        f.close();
        return true;
    }

    bool appendNode(int slot, const char* nodeJson, bool last) {
        char path[48];
        snprintf(path, sizeof(path), "/m%d/nodes.json", slot);
        File f = LittleFS.open(path, "a");
        if (!f) return false;
        f.print(nodeJson);
        if (!last) f.print(",");
        else f.print("]");
        f.close();
        return true;
    }

    bool finalizeMacro(int slot) {
        // Add to index if not already present
        bool found = false;
        for (int i = 0; i < macroCount; i++) {
            if (order[i] == slot) { found = true; break; }
        }
        if (!found && macroCount < MAX_MACROS) {
            order[macroCount] = slot;
            macroCount++;
        }
        loadMacroMeta(slot);
        saveIndex();
        return true;
    }

    bool deleteMacro(int slot) {
        char path[48];
        snprintf(path, sizeof(path), "/m%d/meta.json", slot);
        LittleFS.remove(path);
        snprintf(path, sizeof(path), "/m%d/nodes.json", slot);
        LittleFS.remove(path);
        snprintf(path, sizeof(path), "/m%d/icon.raw", slot);
        LittleFS.remove(path);
        snprintf(path, sizeof(path), "/m%d", slot);
        LittleFS.rmdir(path);

        // Remove from order
        int idx = -1;
        for (int i = 0; i < macroCount; i++) {
            if (order[i] == slot) { idx = i; break; }
        }
        if (idx >= 0) {
            for (int i = idx; i < macroCount - 1; i++) {
                order[i] = order[i + 1];
            }
            macroCount--;
        }
        saveIndex();
        return true;
    }

    bool loadNodes(int slot, JsonDocument& doc) {
        char path[48];
        snprintf(path, sizeof(path), "/m%d/nodes.json", slot);
        File f = LittleFS.open(path, "r");
        if (!f) return false;
        DeserializationError err = deserializeJson(doc, f);
        f.close();
        return err == DeserializationError::Ok;
    }

    void reorder(int* newOrder, int count) {
        macroCount = count;
        for (int i = 0; i < count && i < MAX_MACROS; i++) {
            order[i] = newOrder[i];
        }
        saveIndex();
    }

    size_t getFreeSpace() {
        return LittleFS.totalBytes() - LittleFS.usedBytes();
    }

    // --- Sub-routine storage ---
    //
    // Upload protocol on the wire is "sub_begin → 0..N sub_node → sub_end".
    // Storage writes go to ``/sub/s{N}/nodes.tmp`` during the upload and only
    // get renamed to the live ``/sub/s{N}/nodes.json`` once sub_end fires,
    // confirming we have all the expected nodes AND that the assembled text
    // parses as valid JSON. Three failure modes are now handled atomically:
    //
    //   1. Upload aborts mid-stream (USB unplug, host crash): tmp file exists
    //      but nodes.json is untouched, so loadSubNodes continues to see the
    //      LAST GOOD version (or returns false if it never existed).
    //   2. nodeCount=0 — sub-routine that flattens to nothing. We skip tmp
    //      entirely and write "[]" straight to nodes.json so the file is
    //      immediately valid.
    //   3. Malformed JSON (corrupt host send): finalizeSubWrite re-parses
    //      the tmp before promoting it. If parse fails, the bad tmp is
    //      removed and the live file is left as-is.
    //
    // Engine-side, loadSubNodes is unchanged (just reads nodes.json), so
    // the engine never sees a partial / malformed file on this path.

    bool beginSubWrite(int slot, const char* name, int nodeCount) {
        if (slot < 0 || slot >= MAX_SUBROUTINES) return false;

        char dir[32];
        snprintf(dir, sizeof(dir), "/sub/s%d", slot);
        LittleFS.mkdir("/sub");
        LittleFS.mkdir(dir);

        char path[48];
        snprintf(path, sizeof(path), "/sub/s%d/meta.json", slot);
        File f = LittleFS.open(path, "w");
        if (!f) return false;
        JsonDocument doc;
        doc["name"] = name;
        doc["nodes"] = nodeCount;
        serializeJson(doc, f);
        f.close();

        // Sweep any leftover tmp from a previous interrupted upload so the
        // append path starts from a known-empty state.
        snprintf(path, sizeof(path), "/sub/s%d/nodes.tmp", slot);
        if (LittleFS.exists(path)) LittleFS.remove(path);

        if (nodeCount <= 0) {
            // Empty sub-routine — no append phase will follow, so commit
            // the valid empty array straight to the live file. No tmp dance
            // needed.
            snprintf(path, sizeof(path), "/sub/s%d/nodes.json", slot);
            File nf = LittleFS.open(path, "w");
            if (nf) { nf.print("[]"); nf.close(); }
        } else {
            // Open tmp with the opening bracket. appendSubNode will fill
            // it in; finalizeSubWrite will rename it to nodes.json on
            // success.
            snprintf(path, sizeof(path), "/sub/s%d/nodes.tmp", slot);
            File nf = LittleFS.open(path, "w");
            if (nf) { nf.print("["); nf.close(); }
        }

        strlcpy(subs[slot].name, name, sizeof(subs[slot].name));
        subs[slot].nodeCount = nodeCount;

        if (slot >= subCount) subCount = slot + 1;
        return true;
    }

    bool appendSubNode(int slot, const char* nodeJson, bool last) {
        char path[48];
        snprintf(path, sizeof(path), "/sub/s%d/nodes.tmp", slot);
        File f = LittleFS.open(path, "a");
        if (!f) return false;
        f.print(nodeJson);
        if (!last) f.print(",");
        else f.print("]");
        f.close();
        return true;
    }

    // Promote the just-written tmp to the live nodes.json — only if it
    // parses as valid JSON. Returns true on successful swap. If the tmp
    // file doesn't exist (because beginSubWrite already committed the
    // empty-sub case directly to nodes.json), this is a no-op success.
    // If the tmp file is malformed, the live nodes.json is left as-is
    // (preserving the previous good version) and the bad tmp is removed.
    bool finalizeSubWrite(int slot) {
        if (slot < 0 || slot >= MAX_SUBROUTINES) return false;

        char tmpPath[48], livePath[48];
        snprintf(tmpPath,  sizeof(tmpPath),  "/sub/s%d/nodes.tmp",  slot);
        snprintf(livePath, sizeof(livePath), "/sub/s%d/nodes.json", slot);

        if (!LittleFS.exists(tmpPath)) {
            // beginSubWrite handled the empty-sub case directly. Nothing
            // to promote, but make sure nodes.json exists with at least
            // an empty array so loadSubNodes never returns false here.
            if (!LittleFS.exists(livePath)) {
                File nf = LittleFS.open(livePath, "w");
                if (nf) { nf.print("[]"); nf.close(); }
            }
            return true;
        }

        // Validate the tmp before promoting. If it doesn't parse, the
        // previous live file (if any) is left untouched — the device will
        // keep using the last known-good version of this sub.
        //
        // NOTE: no Serial.printf in this function. It's called from
        // cmdSubEnd during profile upload, which shares the USB CDC pipe
        // with the JSON command/response stream. Any text emitted here
        // would corrupt the host's readline() on the next response and
        // tear down the serial connection. Failure is communicated up
        // through the bool return value; the caller turns that into a
        // sendError(...) JSON payload.
        {
            File vf = LittleFS.open(tmpPath, "r");
            if (!vf) { LittleFS.remove(tmpPath); return false; }
            JsonDocument vdoc;
            DeserializationError err = deserializeJson(vdoc, vf);
            vf.close();
            if (err != DeserializationError::Ok) {
                LittleFS.remove(tmpPath);
                return false;
            }
        }

        // Atomically swap tmp -> live. Some LittleFS versions don't
        // overwrite on rename, so remove the live file first; the window
        // between remove and rename is tiny (microseconds) compared to
        // the full upload, so accepting it here is fine.
        if (LittleFS.exists(livePath)) LittleFS.remove(livePath);
        if (!LittleFS.rename(tmpPath, livePath)) {
            LittleFS.remove(tmpPath);
            return false;
        }
        return true;
    }

    bool loadSubNodes(int slot, JsonDocument& doc) {
        // Stays silent (no Serial.printf) on failure — the storage layer
        // can be exercised from inside the protocol handler in edge
        // cases (e.g. a re-upload while the engine just finished using
        // the sub), and any text emitted on the USB CDC pipe corrupts
        // the host's JSON response stream. The engine's caller handles
        // the failure case with its own diagnostic line.
        char path[48];
        snprintf(path, sizeof(path), "/sub/s%d/nodes.json", slot);
        File f = LittleFS.open(path, "r");
        if (!f) return false;
        DeserializationError err = deserializeJson(doc, f);
        f.close();
        return err == DeserializationError::Ok;
    }

    int findSubByName(const char* name) {
        for (int i = 0; i < subCount; i++) {
            if (strcmp(subs[i].name, name) == 0) return i;
        }
        return -1;
    }

    void clearAllSubs() {
        for (int i = 0; i < subCount; i++) {
            char path[48];
            snprintf(path, sizeof(path), "/sub/s%d/meta.json", i);
            LittleFS.remove(path);
            snprintf(path, sizeof(path), "/sub/s%d/nodes.json", i);
            LittleFS.remove(path);
            // Sweep any stray tmp left over from an interrupted upload.
            snprintf(path, sizeof(path), "/sub/s%d/nodes.tmp", i);
            if (LittleFS.exists(path)) LittleFS.remove(path);
            snprintf(path, sizeof(path), "/sub/s%d", i);
            LittleFS.rmdir(path);
        }
        subCount = 0;
    }

    void loadSubIndex() {
        subCount = 0;
        for (int i = 0; i < MAX_SUBROUTINES; i++) {
            char path[48];
            snprintf(path, sizeof(path), "/sub/s%d/meta.json", i);
            if (!LittleFS.exists(path)) break;

            File f = LittleFS.open(path, "r");
            if (!f) break;
            JsonDocument doc;
            if (deserializeJson(doc, f) == DeserializationError::Ok) {
                strlcpy(subs[i].name, doc["name"] | "Unnamed", sizeof(subs[i].name));
                subs[i].nodeCount = doc["nodes"] | 0;
                subCount = i + 1;
            }
            f.close();
        }
    }
};