#pragma once
// Known-network profile store for Wi-Fi provisioning.
//
// Pure, host-testable logic. Persistence is behind a tiny storage interface so
// the store can be unit-tested without ESP-IDF NVS and so the board layer stays
// free of credential persistence (docs/wifi-provisioning-plan.md section 37).
//
// Security rules (plan sections 7, 28, 29):
//   * A profile password is never logged, rendered, returned by the Portal or
//     sent over USB. Only `ssid` and metadata are ever formatted for humans.
//   * This module owns no device token, gateway token or API key.
#include <array>
#include <cstddef>
#include <cstdint>
#include <cstring>

namespace wearable {

constexpr size_t WIFI_SSID_MAX = 32;      // 802.11 SSID limit
constexpr size_t WIFI_PASSWORD_MAX = 64;  // WPA2 passphrase limit
constexpr size_t WIFI_MAX_PROFILES = 5;   // plan section 15

struct WiFiProfile {
    char ssid[WIFI_SSID_MAX + 1] = {};
    // Stored, never displayed by the UI and never serialized to a human-facing
    // output path. Kept here because reconnection needs it.
    char password[WIFI_PASSWORD_MAX + 1] = {};
    // Monotonic-ish sequence of the last successful connection. 0 means the
    // profile has never connected, so it sorts after every proven network.
    uint32_t last_success = 0;

    bool empty() const { return ssid[0] == '\0'; }
    void set(const char *name, const char *pass) {
        std::strncpy(ssid, name ? name : "", WIFI_SSID_MAX);
        ssid[WIFI_SSID_MAX] = '\0';
        std::strncpy(password, pass ? pass : "", WIFI_PASSWORD_MAX);
        password[WIFI_PASSWORD_MAX] = '\0';
    }
    bool same_ssid(const char *other) const {
        return other && std::strcmp(ssid, other) == 0;
    }
};

// A serialized store image. Fixed-size and versioned so corrupt/garbage NVS
// data is rejected instead of partially applied.
//
// Version 2 adds a durable `initialized` flag. The version bump is deliberate:
// the image is only ever persisted on-device, and v1 images decode as an
// initialized store for backward compatibility (a v1 image could only exist if
// the store had been written at least once).
constexpr uint8_t WIFI_STORE_VERSION = 1;
constexpr uint8_t WIFI_STORE_VERSION_INITIALIZED = 2;
constexpr size_t WIFI_STORE_MAX_BYTES =
    4 + 1 + 1 + 1 + WIFI_MAX_PROFILES * (1 + WIFI_SSID_MAX + 1 + WIFI_PASSWORD_MAX + 4) + 4;

enum class WiFiStoreResult : uint8_t {
    Ok,
    Full,          // no free slot and the SSID is not already known
    EmptySsid,
    SsidTooLong,
    PasswordTooLong,
    InvalidSsid,   // control characters / reserved
    NotFound,
    IoError,
};

struct WiFiStoreImage {
    std::array<uint8_t, WIFI_STORE_MAX_BYTES> bytes{};
    size_t size = 0;
};

// Outcome of reading the persisted image. The three states are semantically
// distinct and must not be collapsed:
//
//   Missing  the key genuinely does not exist (never written)
//   Present  a value was read from storage; it may still fail validation
//   IoError  the storage layer itself could not be read
//
// "Present but rejected by the profile-store decoder" is Corrupt, which is a
// different state from Missing. Only Missing may enable the one-time build-time
// migration.
enum class WiFiStoreLoadState : uint8_t {
    Missing,
    Present,
    IoError,
};

// Persistence boundary. Implemented by NVS on device and by an in-memory fake
// in host tests. Only the platform layer sees this.
class WiFiProfileStorage {
public:
    virtual ~WiFiProfileStorage() = default;
    // Read the raw image. `Missing` means the key was not found; `Present`
    // means bytes were returned (validation happens in the store); `IoError`
    // means the storage layer failed to read.
    virtual WiFiStoreLoadState load(WiFiStoreImage &out) = 0;
    virtual bool save(const WiFiStoreImage &image) = 0;
    virtual bool erase() = 0;
};

// In-memory storage for tests and for a read-only/unsupported platform.
class MemoryWiFiStorage final : public WiFiProfileStorage {
public:
    bool present = false;
    bool io_error = false;
    WiFiStoreImage image{};
    unsigned writes = 0;
    bool fail_save = false;
    WiFiStoreLoadState load(WiFiStoreImage &out) override {
        if (io_error) return WiFiStoreLoadState::IoError;
        if (!present) return WiFiStoreLoadState::Missing;
        out = image;
        return WiFiStoreLoadState::Present;
    }
    bool save(const WiFiStoreImage &value) override {
        ++writes;
        if (fail_save) return false;
        image = value;
        present = true;
        return true;
    }
    bool erase() override {
        present = false;
        image = WiFiStoreImage{};
        return true;
    }
};

// Validate an SSID before accepting it. Rejects empty, over-long and
// control-character SSIDs (which cannot be displayed or encoded in the QR).
inline bool wifi_ssid_valid(const char *ssid) {
    if (!ssid || ssid[0] == '\0') return false;
    const size_t length = std::strlen(ssid);
    if (length == 0 || length > WIFI_SSID_MAX) return false;
    for (size_t i = 0; i < length; ++i) {
        const auto c = static_cast<unsigned char>(ssid[i]);
        if (c < 0x20 || c == 0x7f) return false;
    }
    return true;
}

inline bool wifi_password_valid(const char *password) {
    if (!password) return true;  // open network
    // WPA2 requires either 8..63 ASCII or exactly 64 hex; allow up to 64 here
    // and let the driver reject an invalid length. Empty means open.
    return std::strlen(password) <= WIFI_PASSWORD_MAX;
}

// Pure store. No I/O except through the injected storage object.
class WiFiProfileStore {
public:
    explicit WiFiProfileStore(WiFiProfileStorage &storage) : storage_(storage) {}

    // Load and validate persisted profiles.
    //
    // The three load states are kept distinct:
    //   Missing  -> no store was ever written; initialized() is false and the
    //               one-time build-time migration may run.
    //   Valid    -> a valid image (possibly with zero profiles) was read;
    //               initialized() is true and migration must never run.
    //   Corrupt  -> bytes existed but failed validation (bad CRC, unsupported
    //               version, malformed length, invalid profile). No profiles are
    //               exposed and initialized() is true, so a damaged store can
    //               never be mistaken for a first boot and silently re-import a
    //               build-time credential the user may have removed.
    void load() {
        count_ = 0;
        present_ = false;
        corrupt_ = false;
        load_state_ = WiFiStoreLoadState::Missing;
        loaded_ = false;
        WiFiStoreImage image{};
        const auto storage_state = storage_.load(image);
        switch (storage_state) {
        case WiFiStoreLoadState::Missing:
            // Never written: allow a first-run migration.
            load_state_ = WiFiStoreLoadState::Missing;
            break;
        case WiFiStoreLoadState::IoError:
            // The store could not be read at all. Treat it as present-but-not-
            // missing so a transient read failure cannot trigger a migration
            // that would overwrite data we failed to read.
            load_state_ = WiFiStoreLoadState::IoError;
            present_ = true;
            break;
        case WiFiStoreLoadState::Present:
            if (decode(image)) {
                load_state_ = WiFiStoreLoadState::Present;
                present_ = true;
            } else {
                // Present but rejected: Corrupt. Keep initialized() true so no
                // migration is attempted; expose zero usable profiles.
                load_state_ = WiFiStoreLoadState::IoError;
                present_ = true;
                corrupt_ = true;
                count_ = 0;
            }
            break;
        }
        loaded_ = true;
    }

    size_t count() const { return count_; }
    const WiFiProfile &at(size_t index) const { return profiles_[index]; }

    // Raw storage outcome of the last load. Use `store_was_corrupt()` and
    // `store_missing()` for the migration decision; this is the low-level state.
    WiFiStoreLoadState load_state() const { return load_state_; }
    // True when the key existed but the image failed validation.
    bool store_was_corrupt() const { return corrupt_; }
    // True only when no store has ever been persisted.
    bool store_missing() const { return load_state_ == WiFiStoreLoadState::Missing; }

    // True when a valid store image has ever been persisted, even if it now
    // holds zero profiles. A one-time build-time migration must only run when
    // this is false AND the store was genuinely Missing (not Corrupt).
    bool initialized() const { return present_; }

    // Persist an initialized-but-empty store if none exists yet. Used after a
    // first-run migration so the migration is never repeated. Writing when a
    // store already exists is a no-op.
    WiFiStoreResult mark_initialized() {
        if (present_) return WiFiStoreResult::Ok;
        return persist_or_rollback() ? WiFiStoreResult::Ok : WiFiStoreResult::IoError;
    }

    // Find an exact SSID, or nullptr.
    const WiFiProfile *find(const char *ssid) const {
        for (size_t i = 0; i < count_; ++i)
            if (profiles_[i].same_ssid(ssid)) return &profiles_[i];
        return nullptr;
    }

    WiFiStoreResult add_or_replace(const char *ssid, const char *password) {
        if (!ssid || ssid[0] == '\0') return WiFiStoreResult::EmptySsid;
        if (std::strlen(ssid) > WIFI_SSID_MAX) return WiFiStoreResult::SsidTooLong;
        if (!wifi_ssid_valid(ssid)) return WiFiStoreResult::InvalidSsid;
        if (!wifi_password_valid(password)) return WiFiStoreResult::PasswordTooLong;
        for (size_t i = 0; i < count_; ++i) {
            if (profiles_[i].same_ssid(ssid)) {
                // Preserve last_success across a password change.
                profiles_[i].set(ssid, password);
                return persist_or_rollback() ? WiFiStoreResult::Ok : WiFiStoreResult::IoError;
            }
        }
        if (count_ >= WIFI_MAX_PROFILES) return WiFiStoreResult::Full;
        profiles_[count_].set(ssid, password);
        profiles_[count_].last_success = 0;
        ++count_;
        return persist_or_rollback() ? WiFiStoreResult::Ok : WiFiStoreResult::IoError;
    }

    WiFiStoreResult remove(const char *ssid) {
        for (size_t i = 0; i < count_; ++i) {
            if (profiles_[i].same_ssid(ssid)) {
                for (size_t j = i + 1; j < count_; ++j) profiles_[j - 1] = profiles_[j];
                profiles_[--count_] = WiFiProfile{};
                return persist_or_rollback() ? WiFiStoreResult::Ok : WiFiStoreResult::IoError;
            }
        }
        return WiFiStoreResult::NotFound;
    }

    // Remove every profile. Callers must confirm with the user first
    // (plan section 20); this function itself does not prompt.
    //
    // The store remains initialized: clearing is a deliberate user action and
    // must not re-enable the one-time build-time migration on the next boot.
    WiFiStoreResult clear() {
        if (count_ == 0 && present_) return WiFiStoreResult::Ok;
        const size_t previous = count_;
        const bool previous_present = present_;
        count_ = 0;
        present_ = true;
        for (auto &profile : profiles_) profile = WiFiProfile{};
        if (!persist_or_rollback()) {
            count_ = previous;
            present_ = previous_present;
            return WiFiStoreResult::IoError;
        }
        return WiFiStoreResult::Ok;
    }

    // Record a successful connection. `sequence` is a caller-supplied counter
    // (boot count or seconds) that orders "most recently proven" networks.
    WiFiStoreResult mark_success(const char *ssid, uint32_t sequence) {
        for (size_t i = 0; i < count_; ++i) {
            if (profiles_[i].same_ssid(ssid)) {
                profiles_[i].last_success = sequence ? sequence : 1;
                return persist_or_rollback() ? WiFiStoreResult::Ok : WiFiStoreResult::IoError;
            }
        }
        return WiFiStoreResult::NotFound;
    }

    // Selection order: last successful first (highest sequence), then the rest
    // in insertion order. Ties keep insertion order, so without proof we never
    // reorder just because another AP looks stronger (plan section 17).
    size_t selection_order(std::array<const WiFiProfile *, WIFI_MAX_PROFILES> &out) const {
        size_t n = 0;
        std::array<bool, WIFI_MAX_PROFILES> used{};
        for (;;) {
            size_t best = WIFI_MAX_PROFILES;
            uint32_t best_sequence = 0;
            for (size_t i = 0; i < count_; ++i) {
                if (used[i] || profiles_[i].last_success == 0) continue;
                if (best == WIFI_MAX_PROFILES || profiles_[i].last_success > best_sequence) {
                    best = i;
                    best_sequence = profiles_[i].last_success;
                }
            }
            if (best == WIFI_MAX_PROFILES) break;
            used[best] = true;
            out[n++] = &profiles_[best];
        }
        for (size_t i = 0; i < count_ && n < WIFI_MAX_PROFILES; ++i)
            if (!used[i]) out[n++] = &profiles_[i];
        return n;
    }

private:
    bool persist_or_rollback() {
        WiFiStoreImage image{};
        encode(image);
        if (!storage_.save(image)) {
            // Reload what is actually durable so RAM cannot drift from NVS.
            load();
            return false;
        }
        // A successful write always leaves an initialized store behind.
        present_ = true;
        return true;
    }

    void encode(WiFiStoreImage &image) const {
        size_t offset = 0;
        image.bytes[offset++] = 'W';
        image.bytes[offset++] = 'F';
        image.bytes[offset++] = 'P';
        image.bytes[offset++] = 'S';
        image.bytes[offset++] = WIFI_STORE_VERSION_INITIALIZED;
        image.bytes[offset++] = static_cast<uint8_t>(count_);
        // Any image written from RAM is initialized by definition.
        image.bytes[offset++] = 1;
        for (size_t i = 0; i < count_; ++i) {
            const auto &profile = profiles_[i];
            const auto ssid_length = static_cast<uint8_t>(std::strlen(profile.ssid));
            const auto password_length = static_cast<uint8_t>(std::strlen(profile.password));
            image.bytes[offset++] = ssid_length;
            std::memcpy(image.bytes.data() + offset, profile.ssid, ssid_length);
            offset += ssid_length;
            image.bytes[offset++] = password_length;
            std::memcpy(image.bytes.data() + offset, profile.password, password_length);
            offset += password_length;
            for (unsigned b = 0; b < 4; ++b)
                image.bytes[offset++] = uint8_t(profile.last_success >> (b * 8));
        }
        const uint32_t crc = crc32(image.bytes.data(), offset);
        for (unsigned b = 0; b < 4; ++b) image.bytes[offset++] = uint8_t(crc >> (b * 8));
        image.size = offset;
    }

    bool decode(const WiFiStoreImage &image) {
        const auto &b = image.bytes;
        if (image.size < 10 || image.size > WIFI_STORE_MAX_BYTES) return false;
        if (b[0] != 'W' || b[1] != 'F' || b[2] != 'P' || b[3] != 'S') return false;
        if (b[4] != WIFI_STORE_VERSION && b[4] != WIFI_STORE_VERSION_INITIALIZED) return false;
        // A v1 image predates the explicit flag and could only have been written
        // by an initialized store, so it decodes as initialized.
        const bool had_flag = b[4] >= WIFI_STORE_VERSION_INITIALIZED;
        const uint8_t declared = b[5];
        if (declared > WIFI_MAX_PROFILES) return false;
        // Validate the whole image before mutating any member, so a rejected
        // image cannot leave a half-applied `present_`/`count_` behind.
        if (get32(b.data() + image.size - 4) != crc32(b.data(), image.size - 4)) return false;
        size_t offset = 6;
        if (had_flag) {
            if (image.size < 11) return false;
            ++offset;
        }
        WiFiProfile parsed[WIFI_MAX_PROFILES];
        for (uint8_t i = 0; i < declared; ++i) {
            if (offset >= image.size - 4) return false;
            const uint8_t ssid_length = b[offset++];
            if (ssid_length == 0 || ssid_length > WIFI_SSID_MAX) return false;
            if (offset + ssid_length >= image.size) return false;
            std::memcpy(parsed[i].ssid, b.data() + offset, ssid_length);
            parsed[i].ssid[ssid_length] = '\0';
            offset += ssid_length;
            if (offset >= image.size - 4) return false;
            const uint8_t password_length = b[offset++];
            if (password_length > WIFI_PASSWORD_MAX) return false;
            if (offset + password_length > image.size - 4) return false;
            std::memcpy(parsed[i].password, b.data() + offset, password_length);
            parsed[i].password[password_length] = '\0';
            offset += password_length;
            if (offset + 4 > image.size - 4) return false;
            parsed[i].last_success = get32(b.data() + offset);
            offset += 4;
            if (!wifi_ssid_valid(parsed[i].ssid)) return false;
        }
        if (offset != image.size - 4) return false;
        for (size_t i = 0; i < declared; ++i) profiles_[i] = parsed[i];
        for (size_t i = declared; i < WIFI_MAX_PROFILES; ++i) profiles_[i] = WiFiProfile{};
        count_ = declared;
        // A successfully decoded image is always an initialized store. The
        // explicit version-2 flag cannot make a valid image look uninitialized,
        // which would re-enable the one-time migration.
        present_ = true;
        return true;
    }

    static uint32_t crc32(const uint8_t *p, size_t n) {
        uint32_t c = 0xffffffffU;
        for (size_t i = 0; i < n; ++i) {
            c ^= p[i];
            for (int j = 0; j < 8; ++j) c = (c >> 1) ^ (0xedb88320U & (0U - (c & 1)));
        }
        return ~c;
    }
    static uint32_t get32(const uint8_t *p) {
        uint32_t v = 0;
        for (unsigned i = 0; i < 4; ++i) v |= uint32_t(p[i]) << (i * 8);
        return v;
    }

    WiFiProfileStorage &storage_;
    std::array<WiFiProfile, WIFI_MAX_PROFILES> profiles_{};
    size_t count_ = 0;
    bool loaded_ = false;
    // A store image has been persisted (valid, corrupt, or unreadable).
    bool present_ = false;
    // Storage-level load outcome, kept so callers can distinguish Missing from
    // Corrupt rather than relying on initialized() alone.
    WiFiStoreLoadState load_state_ = WiFiStoreLoadState::Missing;
    // True when the key existed but failed validation.
    bool corrupt_ = false;
};

// Pure build-time migration policy.
//
// The build-time `CONFIG_DEVICE_LAB_WIFI_SSID` may be imported exactly once, and
// only on a genuine first boot. It must never be used to repair a store that was
// deliberately emptied or that was found damaged:
//
//   Missing  + build-time SSID  -> migrate (one time only)
//   Valid    (any profile count)-> never migrate
//   Corrupt  (present, invalid) -> never migrate
//   IoError  (unreadable)       -> never migrate
//
// `initialized` is accepted so callers that only know the initialization mark
// still get the safe answer, but a Missing store is the only state that enables
// migration.
inline bool should_migrate_build_wifi(WiFiStoreLoadState load_state, bool initialized,
                                      bool store_corrupt, bool config_present) {
    if (!config_present) return false;
    if (store_corrupt) return false;      // damaged store is not a first boot
    if (initialized) return false;        // valid-but-empty is already initialized
    // Only a genuinely absent store may adopt the build-time credential.
    return load_state == WiFiStoreLoadState::Missing;
}

}  // namespace wearable