/*******************************************************************************
 * Copyright (c) 2026, Jan Koester jan.koester@gmx.net
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 * Redistributions of source code must retain the above copyright
 *      notice, this list of conditions and the following disclaimer.
 * Redistributions in binary form must reproduce the above copyright
 *      notice, this list of conditions and the following disclaimer in the
 *      documentation and/or other materials provided with the distribution.
 * Neither the name of the <organization> nor the
 *      names of its contributors may be used to endorse or promote products
 *      derived from this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * DISCLAIMED. IN NO EVENT SHALL <COPYRIGHT HOLDER> BE LIABLE FOR ANY
 * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
 * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *******************************************************************************/

#include <paritypp/parity.h>
#include <paritypp/protocol.h>
#include <paritypp/auth.h>
#include <paritypp/server.h>
#include <paritypp/client.h>

#include "cluster.h"
#include "exception.h"
#include "backends/cluster.h"

#include <netplus/exception.h>
#include <crypto/sha.h>

#include <cstring>
#include <iostream>
#include <algorithm>
#include <map>
#include <set>
#include <unordered_map>
#include <future>

#ifdef _WIN32
#include <winsock2.h>
#include <io.h>
#include <fcntl.h>
#pragma comment(lib, "ws2_32.lib")
#else
#include <unistd.h>
#include <fcntl.h>
#include <ifaddrs.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#endif

namespace authdb {

    Cluster *g_Cluster = nullptr;

    // --- SessionBlock serialization ---

    // Wire format:
    //   [4B magic 'SESB']
    //   [16B sid][16B uid][16B did]
    //   [4B member_count]
    //   [16B * member_count member uuids]
    //   [4B username_len]
    //   [username_len bytes username]
    //   [4B gpo_count]
    //   [for each GPO: 16B GPOId + 1B GPORes]

    static constexpr uint8_t SESSION_MAGIC[4] = {'S', 'E', 'S', 'B'};

    std::vector<uint8_t> SessionBlock::serialize(const SessionData &sess) {
        uuid::uuid sid, uid, did;
        sess.getSid(sid);
        sess.getUid(uid);
        sess.getDid(did);
        const char *username = sess.getUsername();
        size_t username_len = username ? strlen(username) : 0;

        // Collect members
        std::vector<uuid::uuid> members;
        sess.getMembers(members);

        // Collect GPO results
        std::vector<std::pair<uuid::uuid, bool>> gpo_results;
        sess.getGPOResults(gpo_results);

        // Magic(4) + Header: 16+16+16 + 4 + members*16 + 4 + username + 4 + gpos*17 + 8 timestamp
        size_t total = 4 + 48 + 4 + (members.size() * 16) + 4 + username_len
                     + 4 + (gpo_results.size() * 17) + 8;

        std::vector<uint8_t> buf(total);
        size_t pos = 0;

        // Magic
        std::memcpy(buf.data() + pos, SESSION_MAGIC, 4);
        pos += 4;

        // SID
        std::memcpy(buf.data() + pos, sid.value, 16);
        pos += 16;

        // UID
        std::memcpy(buf.data() + pos, uid.value, 16);
        pos += 16;

        // DID
        std::memcpy(buf.data() + pos, did.value, 16);
        pos += 16;

        // Member count and UUIDs
        uint32_t mc = static_cast<uint32_t>(members.size());
        std::memcpy(buf.data() + pos, &mc, 4);
        pos += 4;

        for (const auto &m : members) {
            std::memcpy(buf.data() + pos, m.value, 16);
            pos += 16;
        }

        // Username length
        uint32_t ul = static_cast<uint32_t>(username_len);
        std::memcpy(buf.data() + pos, &ul, 4);
        pos += 4;

        // Username
        if (username_len > 0) {
            std::memcpy(buf.data() + pos, username, username_len);
            pos += username_len;
        }

        // GPO count and results
        uint32_t gc = static_cast<uint32_t>(gpo_results.size());
        std::memcpy(buf.data() + pos, &gc, 4);
        pos += 4;

        for (const auto &gpo : gpo_results) {
            std::memcpy(buf.data() + pos, gpo.first.value, 16);
            pos += 16;
            buf[pos] = gpo.second ? 1 : 0;
            pos += 1;
        }

        // Timestamp (seconds since epoch)
        int64_t now = std::chrono::duration_cast<std::chrono::seconds>(
            std::chrono::system_clock::now().time_since_epoch()).count();
        std::memcpy(buf.data() + pos, &now, 8);
        pos += 8;

        return buf;
    }

    bool SessionBlock::deserialize(const std::vector<uint8_t> &data,
                                    uuid::uuid &sid, uuid::uuid &uid, uuid::uuid &did,
                                    std::vector<uuid::uuid> &members,
                                    std::string &username,
                                    std::vector<std::pair<uuid::uuid, bool>> &gpo_results,
                                    int64_t &created_at) {
        created_at = 0;
        if (data.size() < 60) return false;  // 4 magic + 48 uuids + 4 member_count + 4 username_len

        // Check magic
        if (std::memcmp(data.data(), SESSION_MAGIC, 4) != 0) return false;

        size_t pos = 4;

        std::memcpy(sid.value, data.data() + pos, 16);
        pos += 16;

        std::memcpy(uid.value, data.data() + pos, 16);
        pos += 16;

        std::memcpy(did.value, data.data() + pos, 16);
        pos += 16;

        uint32_t mc = 0;
        std::memcpy(&mc, data.data() + pos, 4);
        pos += 4;

        members.clear();
        for (uint32_t i = 0; i < mc; i++) {
            if (pos + 16 > data.size()) return false;
            uuid::uuid m;
            std::memcpy(m.value, data.data() + pos, 16);
            members.push_back(m);
            pos += 16;
        }

        if (pos + 4 > data.size()) return false;
        uint32_t ul = 0;
        std::memcpy(&ul, data.data() + pos, 4);
        pos += 4;

        if (pos + ul > data.size()) return false;
        username.assign(reinterpret_cast<const char*>(data.data() + pos), ul);
        pos += ul;

        // Parse GPO results (may be absent in older format)
        gpo_results.clear();
        if (pos + 4 <= data.size()) {
            uint32_t gc = 0;
            std::memcpy(&gc, data.data() + pos, 4);
            pos += 4;

            for (uint32_t i = 0; i < gc; i++) {
                if (pos + 17 > data.size()) return false;
                uuid::uuid gpoid;
                std::memcpy(gpoid.value, data.data() + pos, 16);
                pos += 16;
                bool gpores = (data[pos] != 0);
                pos += 1;
                gpo_results.emplace_back(gpoid, gpores);
            }
        }

        // Parse timestamp (may be absent in older format)
        if (pos + 8 <= data.size()) {
            std::memcpy(&created_at, data.data() + pos, 8);
            pos += 8;
        }

        return true;
    }

    // --- Deterministic session group_id (cached) ---

    static std::mutex g_sgidCacheMtx;
    static std::unordered_map<uint64_t, uint64_t> g_sgidCache;

    static uint64_t hashUuids(const uint8_t *data, size_t len) {
        auto hash = netplus::sha256_hash(std::vector<uint8_t>(data, data + len));
        uint64_t id = 0;
        for (int i = 0; i < 8; i++)
            id = (id << 8) | hash[i];
        return id;
    }

    static uint64_t simpleCacheKey(const uint8_t *a, size_t alen, const uint8_t *b, size_t blen) {
        uint64_t k = 0xcbf29ce484222325ULL;
        for (size_t i = 0; i < alen; ++i) k = (k ^ a[i]) * 0x100000001b3ULL;
        for (size_t i = 0; i < blen; ++i) k = (k ^ b[i]) * 0x100000001b3ULL;
        return k;
    }

    uint64_t sidGroupId(const uuid::uuid &sid) {
        static const uint8_t prefix = 0xFF;
        uint64_t ckey = simpleCacheKey(&prefix, 1, sid.value, 16);
        {
            std::lock_guard<std::mutex> g(g_sgidCacheMtx);
            auto it = g_sgidCache.find(ckey);
            if (it != g_sgidCache.end())
                return it->second;
        }
        uint8_t input[17];
        input[0] = 0xFF;
        std::memcpy(input + 1, sid.value, 16);
        uint64_t id = hashUuids(input, 17);
        {
            std::lock_guard<std::mutex> g(g_sgidCacheMtx);
            g_sgidCache[ckey] = id;
        }
        return id;
    }

    // --- Cluster ---

    Cluster::Cluster() {}

    Cluster::~Cluster() {
        stop();
    }

    void Cluster::init(const ClusterConfig &cfg) {
        cfg_ = cfg;

        // Setup local block store — file-backed if path configured, otherwise memory
        if (!cfg_.store_path.empty()) {
            std::cerr << "Cluster: using file_block_store at " << cfg_.store_path << std::endl;
            store_ = std::make_shared<paritypp::file_block_store>(cfg_.store_path, 0);
        } else {
            std::cerr << "Cluster: using memory_block_store (no PATH configured)" << std::endl;
            store_ = std::make_shared<paritypp::memory_block_store>();
        }

        // Dedicated memory-only store for sessions (store_id=1)
        session_store_ = std::make_shared<paritypp::memory_block_store>();

        // Setup auth
        auth_ = std::make_shared<paritypp::auth_store>();
        // Add self as client (peers use same credentials)
        if (!cfg_.client_name.empty()) {
            auth_->add_client(cfg_.client_name, cfg_.client_key);
        }

        // Setup server config
        paritypp::server::config scfg;
        scfg.bind_address = cfg_.bind_address;
        scfg.port = cfg_.port;
        scfg.cert_file = cfg_.cert_file;
        scfg.key_file = cfg_.key_file;
        scfg.max_connections = 64;
        scfg.require_auth = !cfg_.client_name.empty();

        // Domain data on store_id=0 (default), sessions on store_id=1
        server_ = std::make_unique<paritypp::server>(scfg, store_, auth_);
        server_->add_store(SESSION_STORE_ID, session_store_);

        // Setup client for talking to peers — filter out self
        size_t peer_count = 0;
        if (!cfg_.peers.empty()) {
            // Collect all local identities: hostname + all local IP addresses
            std::vector<std::string> local_addrs;
            {
                char hostname[256] = {};
                gethostname(hostname, sizeof(hostname));
                local_addrs.emplace_back(hostname);
            }
#ifndef _WIN32
            {
                struct ifaddrs *ifaddr = nullptr;
                if (getifaddrs(&ifaddr) == 0) {
                    for (struct ifaddrs *ifa = ifaddr; ifa; ifa = ifa->ifa_next) {
                        if (!ifa->ifa_addr) continue;
                        char buf[INET6_ADDRSTRLEN] = {};
                        if (ifa->ifa_addr->sa_family == AF_INET) {
                            inet_ntop(AF_INET,
                                      &reinterpret_cast<struct sockaddr_in*>(ifa->ifa_addr)->sin_addr,
                                      buf, sizeof(buf));
                            local_addrs.emplace_back(buf);
                        } else if (ifa->ifa_addr->sa_family == AF_INET6) {
                            inet_ntop(AF_INET6,
                                      &reinterpret_cast<struct sockaddr_in6*>(ifa->ifa_addr)->sin6_addr,
                                      buf, sizeof(buf));
                            local_addrs.emplace_back(buf);
                        }
                    }
                    freeifaddrs(ifaddr);
                }
            }
#endif

            std::vector<paritypp::client::node_info> nodes;
            size_t self_idx = SIZE_MAX;
            for (size_t pi = 0; pi < cfg_.peers.size(); ++pi) {
                const auto &peer = cfg_.peers[pi];
                bool is_self = false;
                if (peer.port == cfg_.port) {
                    for (const auto &addr : local_addrs) {
                        if (peer.address == addr) {
                            is_self = true;
                            break;
                        }
                    }
                }
                if (is_self) {
                    self_idx = nodes.size();
                    std::cerr << "Cluster: self node index " << pi << " (" << peer.address
                              << ":" << peer.port << ")" << std::endl;
                }
                nodes.push_back({peer.address, peer.port});
            }

            peer_count = nodes.size();

            paritypp::client::credentials pcreds;
            if (!cfg_.client_name.empty()) {
                pcreds.client_name = cfg_.client_name;
                pcreds.key.assign(cfg_.client_key.begin(), cfg_.client_key.end());
            }

            if (!nodes.empty()) {
                pclient_ = std::make_unique<paritypp::client>(
                    cfg_.data_blocks, cfg_.parity_blocks, nodes, pcreds);
                if (self_idx != SIZE_MAX) {
                    pclient_->set_local_node(self_idx, store_);
                }
                // Dedicated read client — separate QUIC connections, never blocked
                // by push_worker's store() holding pclient_->mutex_
                read_client_ = std::make_unique<paritypp::client>(
                    cfg_.data_blocks, cfg_.parity_blocks, nodes, pcreds);
                if (self_idx != SIZE_MAX) {
                    read_client_->set_local_node(self_idx, store_);
                }
                // Dedicated session write client — used only by push_worker
                session_client_ = std::make_unique<paritypp::client>(
                    cfg_.data_blocks, cfg_.parity_blocks, nodes, pcreds);
                session_client_->set_store_id(SESSION_STORE_ID);
                if (self_idx != SIZE_MAX) {
                    session_client_->set_local_node(self_idx, session_store_);
                }
                // Dedicated session read client — used by request threads
                session_read_client_ = std::make_unique<paritypp::client>(
                    cfg_.data_blocks, cfg_.parity_blocks, nodes, pcreds);
                session_read_client_->set_store_id(SESSION_STORE_ID);
                if (self_idx != SIZE_MAX) {
                    session_read_client_->set_local_node(self_idx, session_store_);
                }
                // Dedicated scrub/rebalance client — avoids blocking pclient_
                // (push_worker + health monitor) during long scrub operations
                scrub_client_ = std::make_unique<paritypp::client>(
                    cfg_.data_blocks, cfg_.parity_blocks, nodes, pcreds);
                if (self_idx != SIZE_MAX) {
                    scrub_client_->set_local_node(self_idx, store_);
                }

                // Pre-establish QUIC connections on all clients to avoid
                // first-request handshake + auth latency
                pclient_->warmup();
                read_client_->warmup();
                session_client_->warmup();
                session_read_client_->warmup();
                scrub_client_->warmup();
            } else {
                std::cerr << "Cluster: no peers configured" << std::endl;
            }
        }

        // Register peer status callback so GET_STATUS reports peer counts.
        // This callback runs on the server thread when a remote node queries us.
        // peer_count is based on actual filtered node list (after self-exclusion).
        server_->set_peer_status_callback([peer_count]() -> std::pair<uint32_t, uint32_t> {
            return {static_cast<uint32_t>(peer_count), static_cast<uint32_t>(peer_count)};
        });
    }

    void Cluster::start() {
        if (running_) return;
        running_ = true;

        // Purge old session blocks from file store and vacuum to reclaim space.
        // Previous versions stored session shards in blocks.bin; identify them
        // by checking if block 0 of each group starts with the SESB magic
        // (after the 8-byte parity length header).
        if (store_) {
            auto all_groups = store_->list_groups();
            int purged = 0;
            for (uint64_t gid : all_groups) {
                std::vector<uint8_t> block0;
                if (store_->fetch(gid, 0, block0) && block0.size() >= 12) {
                    // The parity client prepends an 8-byte length header.
                    // After that, session data starts with 'SESB' magic.
                    if (block0[8] == 'S' && block0[9] == 'E' &&
                        block0[10] == 'S' && block0[11] == 'B') {
                        store_->remove_group(gid);
                        ++purged;
                    }
                }
            }
            if (purged > 0) {
                std::cerr << "Cluster: purged " << purged
                          << " old session groups from file store" << std::endl;
            }
            if (store_->vacuum()) {
                std::cerr << "Cluster: startup vacuum completed" << std::endl;
            }
        }

        server_thread_ = std::thread(&Cluster::server_loop, this);
        health_monitor_thread_ = std::thread(&Cluster::health_monitor_loop, this);
        push_worker_thread_ = std::thread(&Cluster::push_worker_loop, this);
        session_cleanup_thread_ = std::thread(&Cluster::session_cleanup_loop, this);
    }

    size_t Cluster::waitForPeers(int timeout_s) {
        if (!pclient_) return 0;
        size_t k = cfg_.data_blocks;
        size_t n = cfg_.data_blocks + cfg_.parity_blocks;
        size_t best_online = 0;
        for (int elapsed = 0; elapsed < timeout_s; ++elapsed) {
            try {
                auto health = pclient_->get_cluster_status();
                if (health.nodes_online > best_online)
                    best_online = health.nodes_online;
                std::cerr << "Cluster: " << health.nodes_online << "/"
                          << health.nodes_total << " nodes online ("
                          << elapsed << "s/" << timeout_s << "s)" << std::endl;

                // Also check session_client_ health
                if (session_client_) {
                    try {
                        auto shealth = session_client_->get_cluster_status();
                        std::cerr << "Cluster: session_client " << shealth.nodes_online << "/"
                                  << shealth.nodes_total << " nodes online" << std::endl;
                    } catch (...) {}
                }

                if (health.nodes_online >= n) {
                    std::cerr << "Cluster: all nodes online" << std::endl;
                    critical_ = false;
                    degraded_ = false;
                    return health.nodes_online;
                }
                if (health.nodes_online >= k) {
                    std::cerr << "Cluster: DEGRADED — " << health.nodes_online
                              << "/" << n << " nodes online, continuing" << std::endl;
                    critical_ = false;
                    degraded_ = true;
                    return health.nodes_online;
                }
            } catch (const std::exception &e) {
                std::cerr << "Cluster: peer probe failed: " << e.what() << std::endl;
            }
            std::this_thread::sleep_for(std::chrono::seconds(1));
        }
        std::cerr << "Cluster: FAILED — only " << best_online << "/" << n
                  << " nodes reachable after " << timeout_s << "s (need " << k
                  << ")" << std::endl;
        return best_online;
    }

    void Cluster::stop() {
        if (!running_) return;
        running_ = false;
        monitor_cv_.notify_all();
        push_cv_.notify_all();
        if (server_) server_->stop();
        if (health_monitor_thread_.joinable()) {
            health_monitor_thread_.join();
        }
        if (push_worker_thread_.joinable()) {
            push_worker_thread_.join();
        }
        if (session_cleanup_thread_.joinable()) {
            session_cleanup_thread_.join();
        }
        if (server_thread_.joinable()) {
            server_thread_.join();
        }
    }

    void Cluster::server_loop() {
        try {
            server_->run();
        } catch (const netplus::NetException &e) {
            std::cerr << "Cluster server error: " << e.what() << std::endl;
        } catch (const std::exception &e) {
            std::cerr << "Cluster server error: " << e.what() << std::endl;
        }
        running_ = false;
    }

    void Cluster::health_monitor_loop() {
        size_t k = cfg_.data_blocks;
        size_t n = cfg_.data_blocks + cfg_.parity_blocks;
        bool first_run = true;
        int healthy_cycles = 0;             // counts consecutive healthy probes
        static constexpr int REBALANCE_INTERVAL = 30; // rebalance every ~5min (30 × 10s)
        while (running_) {
            if (!first_run) {
                // Probe faster when degraded/critical (3s) vs normal (10s)
                auto interval = (degraded_ || critical_)
                    ? std::chrono::seconds(3)
                    : std::chrono::seconds(10);
                std::unique_lock<std::mutex> lk(monitor_mutex_);
                monitor_cv_.wait_for(lk, interval,
                                     [this]{ return !running_.load(); });
            }
            first_run = false;
            if (!running_ || !pclient_) continue;
            try {
                auto health = pclient_->get_cluster_status();

                // Propagate dead node info to session clients and read clients
                if (session_client_) {
                    std::set<size_t> online_nodes;
                    for (const auto &ns : health.node_statuses)
                        online_nodes.insert(ns.first);
                    for (size_t i = 0; i < n; ++i) {
                        if (online_nodes.find(i) == online_nodes.end()) {
                            session_client_->mark_node_dead(i);
                            if (session_read_client_)
                                session_read_client_->mark_node_dead(i);
                            if (read_client_)
                                read_client_->mark_node_dead(i);
                            if (scrub_client_)
                                scrub_client_->mark_node_dead(i);
                        }
                    }
                }

                // Verify local store is writable without polluting blocks.bin.
                // Previous approach wrote a sentinel record every health cycle
                // which accumulated thousands of dead records in the data file.
                bool store_ok = true;
                if (!cfg_.store_path.empty()) {
                    try {
                        std::string probe_path = cfg_.store_path + "/.health_probe";
#ifdef _WIN32
                        int probe_fd = ::_open(probe_path.c_str(),
                                               _O_WRONLY | _O_CREAT | _O_TRUNC,
                                               _S_IREAD | _S_IWRITE);
                        if (probe_fd >= 0) {
                            store_ok = (::_write(probe_fd, "OK", 2) == 2);
                            ::_close(probe_fd);
                            ::_unlink(probe_path.c_str());
                        } else {
                            store_ok = false;
                        }
#else
                        int probe_fd = ::open(probe_path.c_str(),
                                              O_WRONLY | O_CREAT | O_TRUNC, 0644);
                        if (probe_fd >= 0) {
                            store_ok = (::write(probe_fd, "OK", 2) == 2);
                            ::close(probe_fd);
                            ::unlink(probe_path.c_str());
                        } else {
                            store_ok = false;
                        }
#endif
                    } catch (...) {
                        store_ok = false;
                    }
                    if (!store_ok && !critical_) {
                        std::cerr << "Cluster: CRITICAL — local store not writable"
                                  << std::endl;
                    }
                }

                bool was_critical = critical_;
                
                if (!store_ok) {
                    degraded_ = true;
                    critical_ = true;
                } else if (health.nodes_online >= n) {
                    if (degraded_) {
                        std::cerr << "Cluster: recovered — all "
                                  << health.nodes_online << "/" << n
                                  << " nodes online";
                        // Do NOT bump recovery_epoch_ here — wait until
                        // scrub has repaired/rebalanced data so that
                        // backends don't attempt fetches before blocks
                        // are available on all nodes.
                        if (domains_ready_.load()) {
                            std::cerr << ", starting scrub" << std::endl;
                            std::thread([this](){
                                try { scrub(); } catch (...) {}
                            }).detach();
                        } else {
                            std::cerr << " (scrub deferred — domains not ready)" << std::endl;
                            // No scrub possible yet, bump epoch so backends
                            // at least try once when domains become ready.
                            recovery_epoch_.fetch_add(1);
                        }
                        healthy_cycles = 0;
                    }
                    degraded_ = false;
                    critical_ = false;

                    // Periodic rebalance while healthy: fixes misplaced blocks
                    // from transient store_stripe timeouts that didn't trigger
                    // a full degraded→recovered transition.
                    if (++healthy_cycles >= REBALANCE_INTERVAL) {
                        healthy_cycles = 0;
                        std::thread([this](){
                            std::unique_lock<std::mutex> lk(scrub_mutex_, std::try_to_lock);
                            if (!lk.owns_lock()) return;
                            if (!scrub_client_ || !scrub_client_->try_scrub_lock()) return;
                            scrub_running_ = true;
                            if (server_) server_->set_rebuilding(true);
                            try {
                                auto rb = scrub_client_->rebalance();
                                if (rb.rebalanced > 0) {
                                    std::cerr << "Cluster: periodic rebalance moved "
                                              << rb.rebalanced << " group(s)" << std::endl;
                                    scrub_client_->vacuum_all_nodes();
                                }
                            } catch (...) {}
                            scrub_running_ = false;
                            if (server_) server_->set_rebuilding(false);
                            scrub_client_->release_scrub_lock();
                        }).detach();
                    }
                } else if (health.nodes_online >= k) {
                    if (was_critical) {
                        std::cerr << "Cluster: recovering — " << health.nodes_online
                                  << "/" << n << " nodes online (was CRITICAL)" << std::endl;
                    } else if (!degraded_) {
                        std::cerr << "Cluster: DEGRADED — " << health.nodes_online
                                  << "/" << n << " nodes online" << std::endl;
                    }
                    degraded_ = true;
                    critical_ = false;
                    healthy_cycles = 0;
                } else {
                    if (!critical_) {
                        std::cerr << "Cluster: CRITICAL — only " << health.nodes_online
                                  << "/" << n << " nodes online (need "
                                  << k << ")" << std::endl;
                    }
                    degraded_ = true;
                    critical_ = true;
                }
            } catch (const std::exception &e) {
                if (!critical_) {
                    std::cerr << "Cluster: health probe failed: " << e.what() << std::endl;
                }
                degraded_ = true;
                critical_ = true;
            }
        }
    }

    void Cluster::pushSession(const SessionData &sess) {
        if (!session_client_) return;

        // Queue asynchronously — the session is already cached locally so
        // getSession() will find it immediately via findLocal().  The push
        // worker thread will replicate to the cluster without blocking the
        // caller (login request).
        queuePush(sess);
    }

    void Cluster::queuePush(const SessionData &sess) {
        if (!session_client_) return;
        auto data = SessionBlock::serialize(sess);
        uuid::uuid sid;
        sess.getSid(sid);
        uint64_t sid_gid = sidGroupId(sid);
        PendingPush pp{data, sid_gid, 0};
        std::lock_guard<std::mutex> pguard(push_mutex_);
        push_queue_.push_back(std::move(pp));
        push_cv_.notify_one();
    }

    void Cluster::queueDataPush(uint64_t dgid, std::vector<uint8_t> buf) {
        std::lock_guard<std::mutex> pguard(push_mutex_);
        // Replace any pending push/remove for the same group (only latest matters)
        for (auto &dp : data_push_queue_) {
            if (dp.dgid == dgid) {
                dp.buf = std::move(buf);
                dp.is_remove = false;
                push_cv_.notify_one();
                return;
            }
        }
        data_push_queue_.push_back({dgid, std::move(buf), false});
        push_cv_.notify_one();
    }

    void Cluster::queueDataRemove(uint64_t dgid) {
        std::lock_guard<std::mutex> pguard(push_mutex_);
        // Replace any pending push for this group with a remove
        for (auto &dp : data_push_queue_) {
            if (dp.dgid == dgid) {
                dp.buf.clear();
                dp.is_remove = true;
                push_cv_.notify_one();
                return;
            }
        }
        data_push_queue_.push_back({dgid, {}, true});
        push_cv_.notify_one();
    }

    void Cluster::push_worker_loop() {
        while (running_) {
            std::unique_lock<std::mutex> lk(push_mutex_);
            push_cv_.wait_for(lk, std::chrono::seconds(5),
                              [this]{ return !push_queue_.empty() || !data_push_queue_.empty() || !running_.load(); });
            if (!running_) break;

            // Don't try to push while critical
            if (critical_) continue;

            // Drain ALL pending session pushes FIRST (more time-critical)
            std::vector<PendingPush> sess_batch;
            while (!push_queue_.empty()) {
                sess_batch.push_back(std::move(push_queue_.front()));
                push_queue_.pop_front();
            }

            // Drain ALL pending data pushes
            std::vector<PendingDataPush> data_batch;
            while (!data_push_queue_.empty()) {
                data_batch.push_back(std::move(data_push_queue_.front()));
                data_push_queue_.pop_front();
            }

            lk.unlock();

            // Push sessions first
            for (auto &item : sess_batch) {
                try {
                    session_client_->store(item.sid_gid, item.data.data(), item.data.size());
                } catch (const std::exception &e) {
                    item.retries++;
                    if (item.retries < 5) {
                        std::lock_guard<std::mutex> pguard(push_mutex_);
                        push_queue_.push_back(std::move(item));
                    } else {
                        std::cerr << "[CLUSTER] push_worker: dropped session after 5 retries: "
                                  << e.what() << std::endl;
                    }
                }
            }

            // Push data items (manifests) — drop after 3 retries to avoid infinite loop
            for (auto &dp : data_batch) {
                try {
                    if (dp.is_remove)
                        pclient_->remove(dp.dgid);
                    else
                        pclient_->store(dp.dgid, dp.buf.data(), dp.buf.size());
                } catch (const std::exception &e) {
                    dp.retries++;
                    if (!dp.is_remove && dp.retries < 3) {
                        std::lock_guard<std::mutex> pguard(push_mutex_);
                        data_push_queue_.push_back(std::move(dp));
                    } else {
                        std::cerr << "[CLUSTER] push_worker: dropped data push after "
                                  << dp.retries << " retries: " << e.what() << std::endl;
                    }
                }
            }
        }
    }

    void Cluster::removeSessionBySid(const uuid::uuid &session_id) {
        uint64_t sid_gid = sidGroupId(session_id);
        if (session_client_) {
            try {
                session_client_->remove(sid_gid);
            } catch (...) {}
        }
    }

    void Cluster::session_cleanup_loop() {
        while (running_) {
            // Sleep 5 minutes between cleanup runs
            for (int i = 0; i < 300 && running_; ++i)
                std::this_thread::sleep_for(std::chrono::seconds(1));
            if (!running_) break;
            if (critical_) continue;
            if (!session_read_client_) continue;

            try {
                int64_t now = std::chrono::duration_cast<std::chrono::seconds>(
                    std::chrono::system_clock::now().time_since_epoch()).count();

                std::set<uint64_t> all_groups;
                size_t n = session_read_client_->get_cluster_status().nodes_total;
                for (size_t i = 0; i < n; i++) {
                    std::vector<uint64_t> node_groups;
                    if (session_read_client_->list_groups_on_node(i, node_groups))
                        all_groups.insert(node_groups.begin(), node_groups.end());
                }

                int removed = 0;
                for (uint64_t gid : all_groups) {
                    try {
                        auto data = session_read_client_->retrieve(gid);
                        if (data.empty()) continue;
                        uuid::uuid sid, uid, did;
                        std::vector<uuid::uuid> members;
                        std::string username;
                        std::vector<std::pair<uuid::uuid, bool>> gpo_results;
                        int64_t created_at = 0;
                        if (SessionBlock::deserialize(data, sid, uid, did, members, username,
                                                      gpo_results, created_at)) {
                            if (created_at > 0 && (now - created_at) > SESSION_TTL_SECONDS) {
                                session_client_->remove(gid);
                                removed++;
                            }
                        }
                    } catch (...) {}
                }
                if (removed > 0) {
                    std::cerr << "[CLUSTER-SESSION] cleanup: removed " << removed
                              << " expired sessions" << std::endl;
                }
            } catch (const std::exception &e) {
                std::cerr << "[CLUSTER-SESSION] cleanup error: " << e.what() << std::endl;
            }
        }
    }

    bool Cluster::fetchSessionLocalBySid(const uuid::uuid &session_id,
                                          uuid::uuid &sid, uuid::uuid &uid, uuid::uuid &did,
                                          std::vector<uuid::uuid> &members,
                                          std::string &username,
                                          std::vector<std::pair<uuid::uuid, bool>> &gpo_results) {
        // Retrieve via dedicated read client — never blocked by push_worker stores
        uint64_t sid_gid = sidGroupId(session_id);
        if (session_read_client_) {
            auto data = session_read_client_->retrieve(sid_gid);
            if (!data.empty()) {
                int64_t created_at;
                return SessionBlock::deserialize(data, sid, uid, did, members, username, gpo_results, created_at);
            }
        }
        return false;
    }

    bool Cluster::hasSessionBySid(const uuid::uuid &session_id) {
        uint64_t sid_gid = sidGroupId(session_id);

        if (session_read_client_) {
            try {
                auto data = session_read_client_->retrieve(sid_gid);
                if (!data.empty()) return true;
            } catch (...) {}
        }

        return false;
    }

    bool Cluster::fetchSessionBySid(const uuid::uuid &session_id,
                                     uuid::uuid &sid, uuid::uuid &uid, uuid::uuid &did,
                                     std::vector<uuid::uuid> &members,
                                     std::string &username,
                                     std::vector<std::pair<uuid::uuid, bool>> &gpo_results) {
        uint64_t sid_gid = sidGroupId(session_id);

        if (session_read_client_) {
            try {
                auto data = session_read_client_->retrieve(sid_gid);
                if (!data.empty()) {
                    int64_t created_at;
                    if (SessionBlock::deserialize(data, sid, uid, did, members, username, gpo_results, created_at)) {
                        return true;
                    }
                }
            } catch (const std::exception &e) {
                return false;
            } catch (...) {
                return false;
            }
        }

        return false;
    }

    void Cluster::listAllSessions(std::vector<SessionInfo> &out) {
        out.clear();
        if (!session_read_client_) return;
        std::set<std::string> seen_sids;

        std::set<uint64_t> all_groups;
        try {
            size_t n = session_read_client_->get_cluster_status().nodes_total;
            for (size_t i = 0; i < n; i++) {
                std::vector<uint64_t> node_groups;
                if (session_read_client_->list_groups_on_node(i, node_groups)) {
                    all_groups.insert(node_groups.begin(), node_groups.end());
                }
            }

            // Try to retrieve and deserialize each group as a session
            for (uint64_t gid : all_groups) {
                try {
                    auto data = session_read_client_->retrieve(gid);
                    if (data.empty()) continue;
                    SessionInfo info;
                    if (SessionBlock::deserialize(data, info.sid, info.uid, info.did,
                                                  info.members, info.username, info.gpo_results,
                                                  info.created_at)) {
                        std::string key = info.sid.c_str();
                        if (!info.sid.empty() && seen_sids.find(key) == seen_sids.end()) {
                            seen_sids.insert(key);
                            out.push_back(std::move(info));
                        }
                    }
                } catch (...) {}
            }
        } catch (...) {}
    }

    // --- DataOpBlock serialization ---

    std::vector<uint8_t> DataOpBlock::serialize(OpType op, const std::string &domain,
                                                 const uuid::uuid &uid, int dataType,
                                                 const AuthData *records) {
        std::vector<uint8_t> buf;

        // Op type
        buf.push_back(static_cast<uint8_t>(op));

        // Domain
        uint32_t dlen = static_cast<uint32_t>(domain.size());
        buf.insert(buf.end(), reinterpret_cast<const uint8_t*>(&dlen), reinterpret_cast<const uint8_t*>(&dlen) + 4);
        buf.insert(buf.end(), domain.begin(), domain.end());

        // UID
        buf.insert(buf.end(), uid.value, uid.value + 16);

        // Data type
        uint32_t dt = static_cast<uint32_t>(dataType);
        buf.insert(buf.end(), reinterpret_cast<const uint8_t*>(&dt), reinterpret_cast<const uint8_t*>(&dt) + 4);

        // Records (only for Create/Edit)
        if (op == OpType::Delete || !records) {
            uint32_t rc = 0;
            buf.insert(buf.end(), reinterpret_cast<const uint8_t*>(&rc), reinterpret_cast<const uint8_t*>(&rc) + 4);
        } else {
            // Count records
            uint32_t rc = 0;
            for (const AuthData *cur = records; cur; cur = cur->next()) rc++;

            buf.insert(buf.end(), reinterpret_cast<const uint8_t*>(&rc), reinterpret_cast<const uint8_t*>(&rc) + 4);

            for (const AuthData *cur = records; cur; cur = cur->next()) {
                // fieldname (255 bytes fixed)
                buf.insert(buf.end(), reinterpret_cast<const uint8_t*>(cur->Data.fieldname),
                           reinterpret_cast<const uint8_t*>(cur->Data.fieldname) + 255);
                // storage type
                uint32_t st = static_cast<uint32_t>(cur->Data.storage);
                buf.insert(buf.end(), reinterpret_cast<const uint8_t*>(&st), reinterpret_cast<const uint8_t*>(&st) + 4);
                // data size
                uint32_t ds = static_cast<uint32_t>(cur->Data.datasize);
                buf.insert(buf.end(), reinterpret_cast<const uint8_t*>(&ds), reinterpret_cast<const uint8_t*>(&ds) + 4);
                // data
                if (ds > 0 && cur->Data.data) {
                    buf.insert(buf.end(), reinterpret_cast<const uint8_t*>(cur->Data.data),
                               reinterpret_cast<const uint8_t*>(cur->Data.data) + ds);
                }
            }
        }

        return buf;
    }

    bool DataOpBlock::deserialize(const std::vector<uint8_t> &data,
                                   OpType &op, std::string &domain,
                                   uuid::uuid &uid, int &dataType,
                                   std::vector<AuthData::Record> &records) {
        if (data.size() < 1 + 4) return false;
        size_t pos = 0;

        // Op type
        op = static_cast<OpType>(data[pos++]);

        // Domain
        uint32_t dlen = 0;
        std::memcpy(&dlen, data.data() + pos, 4); pos += 4;
        if (pos + dlen > data.size()) return false;
        domain.assign(reinterpret_cast<const char*>(data.data() + pos), dlen);
        pos += dlen;

        // UID
        if (pos + 16 > data.size()) return false;
        std::memcpy(uid.value, data.data() + pos, 16); pos += 16;

        // Data type
        if (pos + 4 > data.size()) return false;
        uint32_t dt = 0;
        std::memcpy(&dt, data.data() + pos, 4); pos += 4;
        dataType = static_cast<int>(dt);

        // Record count
        if (pos + 4 > data.size()) return false;
        uint32_t rc = 0;
        std::memcpy(&rc, data.data() + pos, 4); pos += 4;

        records.clear();
        for (uint32_t i = 0; i < rc; i++) {
            if (pos + 255 + 4 + 4 > data.size()) return false;
            AuthData::Record rec;
            std::memcpy(rec.fieldname, data.data() + pos, 255); pos += 255;

            uint32_t st = 0;
            std::memcpy(&st, data.data() + pos, 4); pos += 4;
            rec.storage = static_cast<int>(st);

            uint32_t ds = 0;
            std::memcpy(&ds, data.data() + pos, 4); pos += 4;
            rec.datasize = ds;

            if (ds > 0) {
                if (pos + ds > data.size()) return false;
                rec.data = new char[ds];
                std::memcpy(rec.data, data.data() + pos, ds);
                pos += ds;
            }

            std::memcpy(rec.ruid, uid.value, 16);
            rec.type = dataType;
            records.push_back(rec);
        }

        return true;
    }

    // --- Vacuum: compact block stores on all nodes (cluster-locked) ---

    size_t Cluster::vacuumAllNodes() {
        if (!scrub_client_) return 0;

        std::unique_lock<std::mutex> scrub_lk(scrub_mutex_, std::try_to_lock);
        if (!scrub_lk.owns_lock()) {
            std::cerr << "[VACUUM] skipped — scrub/rebalance is running" << std::endl;
            return 0;
        }

        if (!scrub_client_->try_scrub_lock()) {
            std::cerr << "[VACUUM] skipped — another node holds the scrub lock" << std::endl;
            return 0;
        }

        size_t result = scrub_client_->vacuum_all_nodes();

        scrub_client_->release_scrub_lock();
        return result;
    }

    // --- Scrub: verify and repair all data groups ---

    Cluster::ScrubResult Cluster::scrub() {
        ScrubResult result;
        if (!scrub_client_) return result;

        // Only one scrub/rebalance at a time
        std::unique_lock<std::mutex> scrub_lk(scrub_mutex_, std::try_to_lock);
        if (!scrub_lk.owns_lock()) {
            std::cerr << "[SCRUB] skipped — another scrub/rebalance is already running" << std::endl;
            return result;
        }
        scrub_running_ = true;
        if (server_) server_->set_rebuilding(true);

        // Cluster-wide lock: prevent concurrent scrub across nodes
        if (!scrub_client_->try_scrub_lock()) {
            std::cerr << "[SCRUB] skipped — another node is scrubbing" << std::endl;
            scrub_running_ = false;
            if (server_) server_->set_rebuilding(false);
            return result;
        }

        size_t n = cfg_.data_blocks + cfg_.parity_blocks;

        // With multi-store support, session data lives in store_id=1 and
        // scrub_client_ uses store_id=0 (default), so list_groups_on_node()
        // only returns domain groups — no session filtering needed.

        // 1. Collect all group IDs from every node
        uint64_t lock_gid = paritypp::client::scrub_lock_group_id();
        auto known_domain_gids = ClusterBackend::knownDomainGroupIds();
        std::map<uint64_t, std::set<size_t>> group_nodes;  // group_id -> set of node indices that have it
        std::set<uint64_t> orphaned_groups;
        for (size_t ni = 0; ni < n; ++ni) {
            std::vector<uint64_t> node_groups;
            if (scrub_client_->list_groups_on_node(ni, node_groups)) {
                for (uint64_t gid : node_groups) {
                    if (gid == lock_gid) continue;  // skip lock group
                    if (known_domain_gids.count(gid)) {
                        group_nodes[gid].insert(ni);
                    } else {
                        orphaned_groups.insert(gid);
                    }
                }
            }
        }

        // Remove orphaned groups (old per-entity groups or stale data)
        // Safety: never delete orphaned groups before domain backends are
        // loaded — knownDomainGroupIds() would return an incomplete set and
        // valid domain data would be wiped out.
        if (!orphaned_groups.empty() && domains_ready_.load() && !known_domain_gids.empty()) {
            std::cerr << "[SCRUB] removing " << orphaned_groups.size()
                      << " orphaned groups" << std::endl;
            for (uint64_t gid : orphaned_groups) {
                try {
                    scrub_client_->remove(gid);
                } catch (...) {}
            }
        } else if (!orphaned_groups.empty()) {
            std::cerr << "[SCRUB] skipping " << orphaned_groups.size()
                      << " orphaned groups — domains not ready yet"
                      << std::endl;
        }

        std::cerr << "[SCRUB] found " << group_nodes.size()
                  << " data groups across " << n << " nodes"
                  << " (removed " << orphaned_groups.size() << " orphaned)" << std::endl;

        // 2. Repair under-replicated AND rebalance misplaced in one pass
        for (auto &[gid, present_nodes] : group_nodes) {
            result.groups_checked++;

            if (present_nodes.size() >= n) {
                // All nodes have this group — OK
                continue;
            }

            // Group is under-replicated or misplaced: retrieve and re-store
            std::cerr << "[SCRUB] group " << gid << " present on "
                      << present_nodes.size() << "/" << n
                      << " nodes — repairing" << std::endl;

            try {
                auto data = scrub_client_->retrieve(gid);
                if (data.empty()) {
                    std::cerr << "[SCRUB] group " << gid
                              << " — retrieve returned empty, skipping (NOT removing)" << std::endl;
                    result.groups_failed++;
                    continue;
                }

                // Validate retrieved data before removing the existing blocks.
                // A manifest must have at least a valid AuthHeader + magic + count.
                if (data.size() < sizeof(AuthHeader) + sizeof(uint32_t) + sizeof(uint32_t)) {
                    std::cerr << "[SCRUB] group " << gid
                              << " — retrieved data too small (" << data.size()
                              << " bytes), skipping (NOT removing)" << std::endl;
                    result.groups_failed++;
                    continue;
                }

                // Verify AuthHeader is present
                AuthHeader head;
                std::memcpy(&head, data.data(), sizeof(AuthHeader));
                if (head.start != 0xFE) {
                    std::cerr << "[SCRUB] group " << gid
                              << " — retrieved data has invalid header marker"
                              << ", skipping (NOT removing)" << std::endl;
                    result.groups_failed++;
                    continue;
                }

                // Data is valid — re-store onto the canonical nodes
                // (block_index % n == node_index). store() alone fixes
                // both under-replication and misplacement — it writes
                // every block to its correct node. Do NOT remove() the
                // group first: paritypp::client::remove() is best-effort
                // across all nodes and never throws on partial failure,
                // so a remove-then-restore here could wipe the only good
                // copy of this auth data if the following store() then
                // hit a transient node failure (>m nodes unreachable) —
                // turning a merely-incomplete group into a permanently
                // empty one. Any stale blocks left on nodes that no
                // longer canonically own this group are a space concern
                // only, reclaimed by vacuum_all_nodes() below. (Same bug
                // and same fix as mediadb's Cluster::scrub() and
                // libparitypp's client::rebalance() — this was a third,
                // independent copy of the same logic.)
                scrub_client_->store(gid, data.data(), data.size());
                result.groups_repaired++;

                std::cerr << "[SCRUB] group " << gid << " — repaired" << std::endl;
            } catch (const std::exception &e) {
                std::cerr << "[SCRUB] group " << gid
                          << " — repair failed: " << e.what() << std::endl;
                result.groups_failed++;
            }
        }

        // 3. Vacuum all nodes to reclaim space from overwritten / deleted blocks
        if (result.groups_repaired > 0) {
            scrub_client_->vacuum_all_nodes();
        }

        std::cerr << "[SCRUB] complete: checked=" << result.groups_checked
                  << " repaired=" << result.groups_repaired
                  << " failed=" << result.groups_failed << std::endl;

        scrub_running_ = false;
        if (server_) server_->set_rebuilding(false);
        scrub_client_->release_scrub_lock();

        // Signal backends to re-fetch now that data is repaired/rebalanced
        recovery_epoch_.fetch_add(1);

        return result;
    }

    // --- Data replication ---

    void Cluster::replicateOp(OpType op, const std::string &domain,
                               const uuid::uuid &uid, int dataType,
                               const AuthData *records) {
        if (!pclient_) return;

        auto data = DataOpBlock::serialize(op, domain, uid, dataType, records);

        auto hash = netplus::sha256_hash(data);
        uint64_t localId = 0;
        for (int i = 0; i < 8; i++) localId = (localId << 8) | hash[i];

        // Distribute erasure-coded shards to peer nodes
        try {
            pclient_->store(localId, data.data(), data.size());
        } catch (const netplus::NetException &e) {
            std::cerr << "Cluster replicate failed: " << e.what() << std::endl;
        } catch (const std::exception &e) {
            std::cerr << "Cluster replicate failed: " << e.what() << std::endl;
        }
    }



};
