#include "client.h"
#include "server.h"
#include <random.h>
#include <exception.h>

#include <cstring>
#include <stdexcept>
#include <iostream>
#include <thread>
#include <chrono>
#include <future>
#include <mutex>
#include <set>

#ifdef NDEBUG
#define DBG_LOG(x) do {} while(0)
#else
#define DBG_LOG(x) do { std::cerr << x; } while(0)
#endif

// Milliseconds remaining until deadline, clamped to [0, max_poll_ms].
// Returns 0 when deadline has passed.
static int ms_left(const std::chrono::steady_clock::time_point &deadline,
                   int max_poll_ms = 100) {
    auto now = std::chrono::steady_clock::now();
    if (now >= deadline) return 0;
    int ms = static_cast<int>(
        std::chrono::duration_cast<std::chrono::milliseconds>(deadline - now).count());
    return std::min(ms, max_poll_ms);
}

namespace paritypp {

client::client(size_t data_blocks, size_t parity_blocks,
               const std::vector<node_info>& nodes,
               const credentials& creds)
    : k_(data_blocks), m_(parity_blocks), nodes_(nodes),
      creds_(creds), parity_engine_(data_blocks, parity_blocks) {

    if (nodes_.size() < k_ + m_) {
        std::cerr << "paritypp::client: only " << nodes_.size()
                  << " node(s) for " << k_ << "+" << m_
                  << " config; parity store/retrieve disabled, "
                  << "peer replication still available" << std::endl;
    }

    connections_.resize(nodes_.size());
}

client::~client() = default;

void client::connect_to_node(size_t node_index) {
    auto& nc = connections_[node_index];
    // Reset any previous broken connection
    nc.socket.reset();
    nc.connected = false;

    nc.socket = std::make_unique<netplus::quic>();
    nc.socket->connect(nodes_[node_index].address, nodes_[node_index].port, false);

    // Wait for handshake using a local socketwait (thread-safe)
    netplus::socketwait local_wait;
    auto hs_deadline = std::chrono::steady_clock::now() + std::chrono::seconds(2);
    while (int poll_ms = ms_left(hs_deadline)) {
        if (nc.socket->hasPendingWrite()) {
            if (!local_wait.waitWrite(*nc.socket, poll_ms)) continue;
        } else {
            if (!local_wait.waitRead(*nc.socket, poll_ms)) continue;
        }
        nc.socket->pumpNetwork(MSG_DONTWAIT);
        if (nc.socket->getHandshakeDone()) {
            nc.connected = true;
            nc.dead_until = {};  // clear dead status on successful connect
            // Authenticate if credentials provided
            if (!creds_.client_name.empty()) {
                if (!authenticate_node(node_index)) {
                    nc.socket.reset();
                    nc.connected = false;
                    throw std::runtime_error("Authentication failed on node " +
                                              nodes_[node_index].address);
                }
            }
            return;
        }
    }
    nc.socket.reset();
    nc.connected = false;
    nc.dead_until = std::chrono::steady_clock::now() + std::chrono::seconds(10);
    throw std::runtime_error("Failed to connect to node " + nodes_[node_index].address);
}

bool client::authenticate_node(size_t node_index) {
    // Step 1: Send AUTH_REQUEST with client name
    auto req_payload = protocol::make_auth_request(creds_.client_name);
    auto response = send_recv(node_index, msg_type::AUTH_REQUEST, req_payload);

    if (response.size() < protocol::HEADER_SIZE) return false;

    msg_type resp_type;
    uint32_t resp_len;
    protocol::decode_header(response.data(), response.size(), resp_type, resp_len);

    if (resp_type != msg_type::AUTH_CHALLENGE) return false;

    // Step 2: Parse challenge nonce
    std::vector<uint8_t> nonce;
    if (!protocol::parse_auth_challenge(response.data() + protocol::HEADER_SIZE,
                                         resp_len, nonce)) {
        return false;
    }

    // Step 3: Compute HMAC-SHA256(key, nonce) and send AUTH_RESPONSE
    auto hmac = auth::compute_hmac(creds_.key, nonce);
    auto auth_resp = protocol::make_auth_response(hmac);
    auto final_response = send_recv(node_index, msg_type::AUTH_RESPONSE, auth_resp);

    if (final_response.size() < protocol::HEADER_SIZE) return false;

    protocol::decode_header(final_response.data(), final_response.size(),
                            resp_type, resp_len);

    return resp_type == msg_type::AUTH_OK;
}

void client::ensure_connected(size_t node_index) {
    auto& nc = connections_[node_index];
    if (nc.connected && nc.socket) {
        return;
    }
    // Skip node if recently dead (cooldown period)
    if (nc.dead_until > std::chrono::steady_clock::time_point{}) {
        if (std::chrono::steady_clock::now() < nc.dead_until) {
            throw std::runtime_error("node " + nodes_[node_index].address + " dead (cooldown)");
        }
        nc.dead_until = {};
    }
    // Reset broken state before reconnecting
    nc.socket.reset();
    nc.connected = false;
    connect_to_node(node_index);
}

std::vector<uint8_t> client::send_recv(size_t node_index,
                                         msg_type type,
                                         const std::vector<uint8_t>& payload) {
    netplus::socketwait local_wait;  // thread-safe local instance
    for (int attempt = 0; attempt < 2; ++attempt) {
        try {
            ensure_connected(node_index);
        } catch (...) {
            if (attempt == 0) continue;
            return {};
        }
        if (!connections_[node_index].socket) return {};
        auto& conn = *connections_[node_index].socket;

        uint64_t stream = conn.openStream(true);
        auto msg = protocol::encode_message(type, payload);

        auto sr_deadline = std::chrono::steady_clock::now() + std::chrono::seconds(5);

        bool sent = false;
        while (int poll_ms = ms_left(sr_deadline)) {
            if (local_wait.waitWrite(conn, poll_ms)) {
                conn.sendStreamData(stream, msg, true);
                conn.pumpNetwork(MSG_DONTWAIT);
                sent = true;
                break;
            }
            conn.pumpNetwork(MSG_DONTWAIT);
        }

        if (!sent) {
            try { conn.closeStream(stream); } catch (...) {}
            connections_[node_index].socket.reset();
            connections_[node_index].connected = false;
            if (attempt == 0) continue;
            connections_[node_index].dead_until = std::chrono::steady_clock::now() + std::chrono::seconds(10);
            return {};
        }

        // Wait for response
        std::vector<uint8_t> response;
        while (int poll_ms = ms_left(sr_deadline)) {
            if (!local_wait.waitRead(conn, poll_ms)) continue;
            conn.pumpNetwork(MSG_DONTWAIT);
            if (conn.hasStreamData(stream)) {
                uint8_t hdr[protocol::HEADER_SIZE];
                size_t r = conn.recvStreamData(stream, hdr, protocol::HEADER_SIZE);
                if (r < protocol::HEADER_SIZE) continue;

                msg_type resp_type;
                uint32_t resp_len;
                if (!protocol::decode_header(hdr, r, resp_type, resp_len)) continue;

                response.resize(protocol::HEADER_SIZE + resp_len);
                std::memcpy(response.data(), hdr, protocol::HEADER_SIZE);

                size_t total_read = 0;
                while (total_read < resp_len) {
                    int body_ms = ms_left(sr_deadline);
                    if (body_ms <= 0) break;
                    if (!local_wait.waitRead(conn, body_ms)) continue;
                    conn.pumpNetwork(MSG_DONTWAIT);
                    size_t rd = conn.recvStreamData(stream,
                                                     response.data() + protocol::HEADER_SIZE + total_read,
                                                     resp_len - total_read);
                    total_read += rd;
                }
                break;
            }
        }
        conn.closeStream(stream);

        if (!response.empty())
            return response;

        // Empty response — stale connection, retry once with fresh connection
        connections_[node_index].socket.reset();
        connections_[node_index].connected = false;
        if (attempt > 0) {
            connections_[node_index].dead_until = std::chrono::steady_clock::now() + std::chrono::seconds(10);
        }
    }
    return {};
}

bool client::store_block_on_node(size_t node_index, uint64_t group_id,
                                   uint32_t block_index,
                                   const uint8_t* data, size_t len) {
    if (static_cast<int>(node_index) == local_node_index_ && local_store_) {
        return local_store_->store(group_id, block_index, data, len);
    }
    auto payload = protocol::make_store_block(group_id, block_index, data, len);
    auto response = send_recv(node_index, msg_type::STORE_BLOCK, payload);

    if (response.size() >= protocol::HEADER_SIZE) {
        msg_type type;
        uint32_t resp_len;
        protocol::decode_header(response.data(), response.size(), type, resp_len);
        return type == msg_type::OK;
    }
    return false;
}

bool client::fetch_block_from_node(size_t node_index, uint64_t group_id,
                                     uint32_t block_index,
                                     std::vector<uint8_t>& out) {
    if (static_cast<int>(node_index) == local_node_index_ && local_store_) {
        return local_store_->fetch(group_id, block_index, out);
    }
    auto payload = protocol::make_fetch_block(group_id, block_index);
    auto response = send_recv(node_index, msg_type::FETCH_BLOCK, payload);

    if (response.size() >= protocol::HEADER_SIZE) {
        msg_type type;
        uint32_t resp_len;
        protocol::decode_header(response.data(), response.size(), type, resp_len);
        if (type == msg_type::BLOCK_DATA && resp_len >= 12) {
            const uint8_t* data;
            size_t data_len;
            uint64_t gid;
            uint32_t bidx;
            if (protocol::parse_block_data(response.data() + protocol::HEADER_SIZE,
                                            resp_len, gid, bidx, data, data_len)) {
                out.assign(data, data + data_len);
                return true;
            }
        }
    }
    return false;
}

uint64_t client::generate_group_id() {
    // Use random bytes to generate a unique group ID
    uint64_t id = 0;
    uint8_t buf[8];
    // Simple time-based + random fallback
    auto now = std::chrono::steady_clock::now().time_since_epoch();
    id = static_cast<uint64_t>(
        std::chrono::duration_cast<std::chrono::nanoseconds>(now).count());
    return id;
}

uint64_t client::store(const uint8_t* data, size_t data_len) {
    uint64_t group_id = generate_group_id();
    store(group_id, data, data_len);
    return group_id;
}

uint64_t client::store(const std::vector<uint8_t>& data) {
    return store(data.data(), data.size());
}

// ---- Stripe-based store: process data in STRIPE_SIZE chunks ----
// Pipelined: open streams and send all blocks first, then collect responses.
// This reduces per-stripe latency from O(n × RTT) to O(RTT).

void client::store_stripe(uint64_t group_id, uint32_t stripe_index,
                           const uint8_t* stripe_data, size_t stripe_len) {
    // Use actual data size (rounded up to multiple of k) instead of full STRIPE_SIZE
    // for small payloads — avoids sending 4MB of mostly zeros for tiny data.
    size_t effective_size = std::min(stripe_len, STRIPE_SIZE);
    // Round up to multiple of k_ so split() produces equal-size blocks
    size_t block_align = (effective_size + k_ - 1) / k_ * k_;
    if (block_align < k_) block_align = k_; // minimum 1 byte per block

    std::vector<uint8_t> padded(block_align, 0);
    std::memcpy(padded.data(), stripe_data, effective_size);

    auto sub_blocks = parity::split(padded.data(), block_align, k_);
    auto parity_blocks = parity_engine_.encode(sub_blocks);

    size_t n = k_ + m_;

    auto stripe_deadline = std::chrono::steady_clock::now() + std::chrono::seconds(5);

    // Pending stream info per node-slot
    struct Pending {
        size_t node_index;
        uint64_t stream = 0;
        bool local = false;
        bool sent = false;
        bool ok = false;
        std::string fail_reason; // diagnostic: why this node failed
    };
    std::vector<Pending> pending(n);

    auto send_to_node = [&](size_t i) {
        const uint8_t* bdata = (i < k_) ? sub_blocks[i].data() : parity_blocks[i - k_].data();
        size_t blen = (i < k_) ? sub_blocks[i].size() : parity_blocks[i - k_].size();
        uint32_t block_index = stripe_index * static_cast<uint32_t>(n) + static_cast<uint32_t>(i);

        pending[i].node_index = i;
        pending[i].sent = false;

        if (static_cast<int>(i) == local_node_index_ && local_store_) {
            pending[i].local = true;
            pending[i].ok = local_store_->store(group_id, block_index, bdata, blen);
            if (!pending[i].ok) pending[i].fail_reason = "local store failed";
            pending[i].sent = true;
            return;
        }

        netplus::socketwait local_wait;  // thread-local wait instance

        try {
            ensure_connected(i);
            if (!connections_[i].connected || !connections_[i].socket) {
                pending[i].fail_reason = "connect failed (handshake timeout)";
                return;
            }
            auto& conn = *connections_[i].socket;
            auto payload = protocol::make_store_block(group_id, block_index, bdata, blen);
            auto msg = protocol::encode_message(msg_type::STORE_BLOCK, payload);
            uint64_t stream = conn.openStream(true);
            pending[i].stream = stream;

            // Non-blocking send
            while (int poll_ms = ms_left(stripe_deadline)) {
                if (local_wait.waitWrite(conn, poll_ms)) {
                    conn.sendStreamData(stream, msg, true);
                    // Pump immediately to flush queued UDP datagrams
                    conn.pumpNetwork(MSG_DONTWAIT);
                    pending[i].sent = true;
                    break;
                }
                conn.pumpNetwork(MSG_DONTWAIT);
            }
            // Close leaked stream if send failed
            if (!pending[i].sent) {
                pending[i].fail_reason = "send timeout (waitWrite)";
                try { conn.closeStream(stream); } catch (...) {}
            }
        } catch (const netplus::NetException& e) {
            pending[i].fail_reason = std::string("net: ") + e.what();
            connections_[i].socket.reset();
            connections_[i].connected = false;
            connections_[i].dead_until = std::chrono::steady_clock::now() + std::chrono::seconds(10);
        } catch (const std::exception& e) {
            pending[i].fail_reason = std::string("exception: ") + e.what();
            connections_[i].socket.reset();
            connections_[i].connected = false;
            connections_[i].dead_until = std::chrono::steady_clock::now() + std::chrono::seconds(10);
        } catch (...) {
            pending[i].fail_reason = "unknown exception";
            connections_[i].socket.reset();
            connections_[i].connected = false;
            connections_[i].dead_until = std::chrono::steady_clock::now() + std::chrono::seconds(10);
        }
    };

    // Phase 1: store local blocks + send to remote nodes in parallel
    {
        // Local node first (no thread needed)
        for (size_t i = 0; i < n; i++) {
            if (static_cast<int>(i) == local_node_index_ && local_store_) {
                send_to_node(i);
            }
        }
        // Remote nodes in parallel
        std::vector<std::future<void>> futures;
        for (size_t i = 0; i < n; i++) {
            if (static_cast<int>(i) == local_node_index_ && local_store_) continue;
            futures.push_back(std::async(std::launch::async, send_to_node, i));
        }
        for (auto& f : futures) f.get();
    }

    // Phase 1b: retry failed remote nodes once with a fresh connection
    {
        std::vector<std::future<void>> futures;
        for (size_t i = 0; i < n; i++) {
            if (!pending[i].sent && !pending[i].local) {
                pending[i].fail_reason.clear();
                connections_[i].dead_until = {};
                connections_[i].socket.reset();
                connections_[i].connected = false;
                futures.push_back(std::async(std::launch::async, send_to_node, i));
            }
        }
        for (auto& f : futures) f.get();
    }

    // Phase 2: collect responses using waitReadMulti
    size_t remaining = 0;
    for (auto& p : pending) {
        if (p.sent && !p.local) ++remaining;
    }
    // Pump all connections once — data may already be buffered from Phase 1
    for (auto& p : pending) {
        if (!p.sent || p.local || p.ok) continue;
        if (!connections_[p.node_index].socket) continue;
        try { connections_[p.node_index].socket->pumpNetwork(MSG_DONTWAIT); } catch (...) {}
    }

    for (;;) {
        int poll_ms = ms_left(stripe_deadline);
        if (poll_ms <= 0 || remaining == 0) break;
        // Check buffered data first before blocking
        for (auto& p : pending) {
            if (!p.sent || p.local || p.ok) continue;
            if (!connections_[p.node_index].socket) {
                --remaining;
                p.sent = false;
                if (p.fail_reason.empty()) p.fail_reason = "connection lost during Phase 2";
                continue;
            }
            auto& conn = *connections_[p.node_index].socket;
            if (!conn.hasStreamData(p.stream)) continue;
            try {
                uint8_t hdr[protocol::HEADER_SIZE];
                size_t r = conn.recvStreamData(p.stream, hdr, protocol::HEADER_SIZE);
                if (r >= protocol::HEADER_SIZE) {
                    msg_type resp_type;
                    uint32_t resp_len;
                    if (protocol::decode_header(hdr, r, resp_type, resp_len)) {
                        if (resp_len > 0) {
                            std::vector<uint8_t> discard(resp_len);
                            size_t total = 0;
                            while (total < resp_len) {
                                int body_ms = ms_left(stripe_deadline);
                                if (body_ms <= 0) break;
                                if (!wait_.waitRead(conn, body_ms)) continue;
                                conn.pumpNetwork(MSG_DONTWAIT);
                                total += conn.recvStreamData(p.stream,
                                    discard.data() + total, resp_len - total);
                            }
                        }
                        p.ok = (resp_type == msg_type::OK);
                        if (!p.ok) p.fail_reason = "server returned non-OK response";
                        conn.closeStream(p.stream);
                        --remaining;
                    }
                }
            } catch (const netplus::NetException& e) {
                p.fail_reason = std::string("Phase 2 net: ") + e.what();
                connections_[p.node_index].socket.reset();
                connections_[p.node_index].connected = false;
                --remaining;
            } catch (const std::exception& e) {
                p.fail_reason = std::string("Phase 2 exception: ") + e.what();
                connections_[p.node_index].socket.reset();
                connections_[p.node_index].connected = false;
                --remaining;
            } catch (...) {
                p.fail_reason = "Phase 2 unknown exception";
                connections_[p.node_index].socket.reset();
                connections_[p.node_index].connected = false;
                --remaining;
            }
        }
        if (remaining == 0) break;

        // Build socket list for waitReadMulti
        std::vector<netplus::socket*> wait_socks;
        std::vector<size_t> wait_idx; // maps wait_socks position -> pending index
        for (size_t i = 0; i < pending.size(); ++i) {
            auto& p = pending[i];
            if (!p.sent || p.local || p.ok) continue;
            if (!connections_[p.node_index].socket) continue;
            wait_socks.push_back(connections_[p.node_index].socket.get());
            wait_idx.push_back(i);
        }
        if (wait_socks.empty()) break;

        std::vector<bool> ready;
        wait_.waitReadMulti(wait_socks, ready, std::min(poll_ms, 50));

        // Pump only sockets that are actually ready
        for (size_t wi = 0; wi < wait_socks.size(); ++wi) {
            if (!ready[wi]) continue;
            auto& p = pending[wait_idx[wi]];
            if (!connections_[p.node_index].socket) continue;
            try {
                connections_[p.node_index].socket->pumpNetwork(MSG_DONTWAIT);
            } catch (...) {}
        }
    }

    // Mark response timeout for nodes that never responded
    for (auto& p : pending) {
        if (p.sent && !p.local && !p.ok && p.fail_reason.empty()) {
            p.fail_reason = "response timeout";
        }
    }

    // Close any streams still open
    for (auto& p : pending) {
        if (p.sent && !p.local && !p.ok && connections_[p.node_index].socket) {
            try { connections_[p.node_index].socket->closeStream(p.stream); } catch (...) {}
        }
    }

    // Count failures and build diagnostic
    size_t failed = 0;
    std::string diag;
    for (auto& p : pending) {
        if (!p.ok) {
            ++failed;
            if (!diag.empty()) diag += "; ";
            diag += "node[" + std::to_string(p.node_index) + "]";
            if (p.node_index < nodes_.size()) {
                diag += "(" + nodes_[p.node_index].address + ":" +
                        std::to_string(nodes_[p.node_index].port) + ")";
            }
            diag += ": " + (p.fail_reason.empty() ? "unknown" : p.fail_reason);
        }
    }

    if (failed > m_) {
        DBG_LOG("[PARITY] store_stripe FAILED: " << diag << "\n");
        throw std::runtime_error("store_stripe: too many node failures ("
                                 + std::to_string(failed) + " > " + std::to_string(m_)
                                 + "): " + diag);
    }
}

bool client::retrieve_stripe(std::unique_lock<std::mutex>& lk,
                              uint64_t group_id, uint32_t stripe_index,
                              std::vector<uint8_t>& out) {
    // Retry with backoff when the cluster is rebuilding after a full restart.
    // Without this, retrieve returns false immediately if too many nodes
    // haven't synced the block yet, even though the data exists on peers.
    static constexpr int MAX_RETRIES = 6;           // up to ~30s total
    static constexpr int INITIAL_WAIT_MS = 500;     // first retry after 500ms

    for (int attempt = 0; attempt <= MAX_RETRIES; ++attempt) {
        if (attempt > 0) {
            // Release the global mutex during the backoff sleep so that other
            // client operations (retrieves for other groups, stores, etc.)
            // are not blocked.
            int wait_ms = INITIAL_WAIT_MS * (1 << (attempt - 1));
            lk.unlock();
            std::this_thread::sleep_for(std::chrono::milliseconds(wait_ms));
            lk.lock();
        }

        bool should_retry = false;
        if (!retrieve_stripe_once(group_id, stripe_index, out, should_retry))  {
            if (!should_retry || attempt == MAX_RETRIES)
                return false;
            // Some blocks were found but not enough — cluster may be rebuilding
            continue;
        }
        return true;
    }
    return false;
}

bool client::retrieve_stripe_once(uint64_t group_id, uint32_t stripe_index,
                                   std::vector<uint8_t>& out, bool& should_retry) {
    should_retry = false;
    size_t n = k_ + m_;

    auto retr_deadline = std::chrono::steady_clock::now() + std::chrono::seconds(5);

    // Pipelined fetch: send requests to all nodes in parallel,
    // then collect responses using waitReadMulti.
    struct Pending {
        size_t node_index;
        uint64_t stream = 0;
        bool local = false;
        bool sent = false;
        bool ok = false;
        std::vector<uint8_t> block;
    };
    std::vector<Pending> pending(n);
    size_t local_ok_count = 0;

    // Phase 1: send FETCH_BLOCK to all nodes (local first, remote in parallel)
    auto fetch_from_node = [&](size_t i) {
        uint32_t bidx = stripe_index * static_cast<uint32_t>(n) + static_cast<uint32_t>(i);
        pending[i].node_index = i;

        if (static_cast<int>(i) == local_node_index_ && local_store_) {
            pending[i].local = true;
            pending[i].ok = local_store_->fetch(group_id, bidx, pending[i].block)
                            && !pending[i].block.empty();
            return;
        }

        netplus::socketwait local_wait;  // thread-local wait instance

        try {
            ensure_connected(i);
            if (!connections_[i].connected || !connections_[i].socket) return;
            auto& conn = *connections_[i].socket;
            auto payload = protocol::make_fetch_block(group_id, bidx);
            auto msg = protocol::encode_message(msg_type::FETCH_BLOCK, payload);
            uint64_t stream = conn.openStream(true);
            pending[i].stream = stream;

            while (int poll_ms = ms_left(retr_deadline)) {
                if (local_wait.waitWrite(conn, poll_ms)) {
                    conn.sendStreamData(stream, msg, true);
                    conn.pumpNetwork(MSG_DONTWAIT);
                    pending[i].sent = true;
                    break;
                }
                conn.pumpNetwork(MSG_DONTWAIT);
            }
            if (!pending[i].sent) {
                try { conn.closeStream(stream); } catch (...) {}
            }
        } catch (...) {
            connections_[i].socket.reset();
            connections_[i].connected = false;
            connections_[i].dead_until = std::chrono::steady_clock::now() + std::chrono::seconds(10);
        }
    };

    // Local node first
    for (size_t i = 0; i < n; ++i) {
        if (static_cast<int>(i) == local_node_index_ && local_store_) {
            fetch_from_node(i);
            if (pending[i].ok) ++local_ok_count;
        }
    }
    // Remote nodes in parallel
    {
        std::vector<std::future<void>> futures;
        for (size_t i = 0; i < n; ++i) {
            if (static_cast<int>(i) == local_node_index_ && local_store_) continue;
            futures.push_back(std::async(std::launch::async, fetch_from_node, i));
        }
        for (auto& f : futures) f.get();
    }

    // Phase 2: collect responses using waitReadMulti
    size_t remote_sent = 0;
    for (auto& p : pending) if (p.sent && !p.local) ++remote_sent;

    // Pump all connections once
    for (auto& p : pending) {
        if (!p.sent || p.local || p.ok) continue;
        if (!connections_[p.node_index].socket) continue;
        try { connections_[p.node_index].socket->pumpNetwork(MSG_DONTWAIT); } catch (...) {}
    }

    size_t remaining = remote_sent;
    size_t total_ok = local_ok_count;

    for (;;) {
        int poll_ms = ms_left(retr_deadline);
        if (poll_ms <= 0 || remaining == 0 || total_ok >= k_) break;
        // Check buffered data
        for (auto& p : pending) {
            if (!p.sent || p.local || p.ok) continue;
            if (!connections_[p.node_index].socket) {
                --remaining;
                p.sent = false;
                continue;
            }
            auto& conn = *connections_[p.node_index].socket;
            if (!conn.hasStreamData(p.stream)) continue;
            try {
                uint8_t hdr[protocol::HEADER_SIZE];
                size_t r = conn.recvStreamData(p.stream, hdr, protocol::HEADER_SIZE);
                if (r >= protocol::HEADER_SIZE) {
                    msg_type resp_type;
                    uint32_t resp_len;
                    if (protocol::decode_header(hdr, r, resp_type, resp_len)) {
                        if (resp_type == msg_type::BLOCK_DATA && resp_len >= 12) {
                            std::vector<uint8_t> body(resp_len);
                            size_t total_read = 0;
                            while (total_read < resp_len) {
                                int body_ms = ms_left(retr_deadline);
                                if (body_ms <= 0) break;
                                if (!wait_.waitRead(conn, body_ms)) continue;
                                conn.pumpNetwork(MSG_DONTWAIT);
                                total_read += conn.recvStreamData(p.stream,
                                    body.data() + total_read, resp_len - total_read);
                            }
                            if (total_read >= resp_len) {
                                const uint8_t* data;
                                size_t data_len;
                                uint64_t gid;
                                uint32_t bidx;
                                if (protocol::parse_block_data(body.data(), resp_len,
                                                                gid, bidx, data, data_len)) {
                                    p.block.assign(data, data + data_len);
                                    p.ok = true;
                                    ++total_ok;
                                }
                            }
                        } else if (resp_len > 0) {
                            // Drain non-BLOCK_DATA response body
                            std::vector<uint8_t> discard(resp_len);
                            size_t total_read = 0;
                            while (total_read < resp_len) {
                                int body_ms = ms_left(retr_deadline);
                                if (body_ms <= 0) break;
                                if (!wait_.waitRead(conn, body_ms)) continue;
                                conn.pumpNetwork(MSG_DONTWAIT);
                                total_read += conn.recvStreamData(p.stream,
                                    discard.data() + total_read, resp_len - total_read);
                            }
                        }
                    }
                    conn.closeStream(p.stream);
                    --remaining;
                }
            } catch (...) {
                connections_[p.node_index].socket.reset();
                connections_[p.node_index].connected = false;
                --remaining;
            }
        }
        if (remaining == 0 || total_ok >= k_) break;

        // Wait on all pending sockets
        std::vector<netplus::socket*> wait_socks;
        std::vector<size_t> wait_idx;
        for (size_t i = 0; i < pending.size(); ++i) {
            auto& p = pending[i];
            if (!p.sent || p.local || p.ok) continue;
            if (!connections_[p.node_index].socket) continue;
            wait_socks.push_back(connections_[p.node_index].socket.get());
            wait_idx.push_back(i);
        }
        if (wait_socks.empty()) break;

        std::vector<bool> ready;
        wait_.waitReadMulti(wait_socks, ready, std::min(poll_ms, 50));

        for (size_t wi = 0; wi < wait_socks.size(); ++wi) {
            if (!ready[wi]) continue;
            auto& p = pending[wait_idx[wi]];
            if (!connections_[p.node_index].socket) continue;
            try { connections_[p.node_index].socket->pumpNetwork(MSG_DONTWAIT); } catch (...) {}
        }
    }

    // Close remaining streams
    for (auto& p : pending) {
        if (p.sent && !p.local && !p.ok && connections_[p.node_index].socket) {
            try { connections_[p.node_index].socket->closeStream(p.stream); } catch (...) {}
        }
    }

    // Collect successful blocks
    std::vector<std::vector<uint8_t>> fetched;
    std::vector<size_t> fetch_indices;
    for (auto& p : pending) {
        if (p.ok && !p.block.empty()) {
            fetch_indices.push_back(p.node_index);
            fetched.push_back(std::move(p.block));
        }
    }

    if (fetched.empty()) return false;
    if (fetched.size() < k_) {
        // Some blocks found but not enough — the cluster might be rebuilding
        should_retry = true;
        return false;
    }

    // Check if all data blocks present (first k nodes)
    bool all_data = true;
    for (size_t i = 0; i < k_; i++) {
        bool found = false;
        for (auto idx : fetch_indices) { if (idx == i) { found = true; break; } }
        if (!found) { all_data = false; break; }
    }

    std::vector<std::vector<uint8_t>> data_blocks;
    if (all_data) {
        data_blocks.resize(k_);
        for (size_t j = 0; j < fetch_indices.size(); j++) {
            if (fetch_indices[j] < k_)
                data_blocks[fetch_indices[j]] = std::move(fetched[j]);
        }
    } else {
        size_t block_size = fetched[0].size();
        data_blocks = parity_engine_.decode(fetched, fetch_indices, block_size);
    }

    out = parity::join(data_blocks, 0);
    return true;
}

void client::store(uint64_t group_id, const uint8_t* data, size_t data_len) {
    std::lock_guard<std::mutex> guard(mutex_);
    if (nodes_.size() < k_ + m_) {
        throw std::runtime_error("Not enough nodes for parity store ("
                                 + std::to_string(nodes_.size()) + " < "
                                 + std::to_string(k_ + m_) + ")");
    }

    // Prepend 8-byte length header
    std::vector<uint8_t> header(8);
    protocol::encode_u64(header.data(), static_cast<uint64_t>(data_len));

    // Combine header + data into a logical stream, process in stripes
    size_t framed_len = 8 + data_len;
    uint32_t total_stripes = static_cast<uint32_t>((framed_len + STRIPE_SIZE - 1) / STRIPE_SIZE);
    DBG_LOG("[PARITY] store gid=" << group_id << " data_len=" << data_len
              << " stripes=" << total_stripes << "\n");

    // store_stripe already handles per-node failures and retries internally.
    // Only retry at this level if the stripe itself throws (> m nodes failed).
    for (int attempt = 0; attempt < 2; ++attempt) {
        try {
            uint32_t stripe_idx = 0;
            size_t offset = 0;

            while (offset < framed_len) {
                size_t chunk = std::min<size_t>(STRIPE_SIZE, framed_len - offset);
                std::vector<uint8_t> stripe(chunk, 0);
                size_t filled = 0;

                // Fill from header bytes if still in range
                if (offset < 8) {
                    size_t hdr_copy = std::min<size_t>(8 - offset, chunk);
                    std::memcpy(stripe.data(), header.data() + offset, hdr_copy);
                    filled = hdr_copy;
                }
                // Fill from payload
                if (filled < chunk && offset + filled >= 8) {
                    size_t data_off = offset + filled - 8;
                    size_t data_copy = std::min<size_t>(chunk - filled, data_len - data_off);
                    if (data_copy > 0)
                        std::memcpy(stripe.data() + filled, data + data_off, data_copy);
                    filled += data_copy;
                }

                store_stripe(group_id, stripe_idx, stripe.data(), filled);
                if (stripe_idx % 100 == 0 || stripe_idx + 1 == total_stripes) {
                    DBG_LOG("[PARITY]   stripe " << (stripe_idx+1) << "/" << total_stripes << "\n");
                }
                ++stripe_idx;
                offset += STRIPE_SIZE;
            }
            DBG_LOG("[PARITY] store gid=" << group_id << " complete\n");
            return; // success
        } catch (const std::exception& e) {
            if (attempt == 0) {
                DBG_LOG("[PARITY] store attempt 1 failed: " << e.what()
                          << ", clearing dead marks and retrying\n");
                // Clear dead_until on all nodes so the retry actually reconnects
                for (auto& nc : connections_) {
                    nc.dead_until = {};
                    nc.socket.reset();
                    nc.connected = false;
                }
            } else {
                throw; // second attempt failed, propagate
            }
        }
    }
}

void client::store(uint64_t group_id, const std::vector<uint8_t>& data) {
    store(group_id, data.data(), data.size());
}

// ---- Incremental store session (stripe-based, ~16KB peak RAM) ----

class client_store_session : public store_session {
public:
    client_store_session(client& cli, uint64_t group_id, size_t total_size)
        : cli_(cli), group_id_(group_id) {
        // Write 8-byte length header into the stripe buffer first
        uint8_t hdr[8];
        protocol::encode_u64(hdr, static_cast<uint64_t>(total_size));
        stripe_buf_.insert(stripe_buf_.end(), hdr, hdr + 8);
    }

    ~client_store_session() override = default;

    void feed(const uint8_t* data, size_t len) override {
        if (!ok_ || finalized_) return;
        const uint8_t* p = data;
        size_t remaining = len;

        while (remaining > 0) {
            size_t space = STRIPE_SIZE - stripe_buf_.size();
            size_t n = std::min(remaining, space);
            stripe_buf_.insert(stripe_buf_.end(), p, p + n);
            p += n;
            remaining -= n;

            if (stripe_buf_.size() >= STRIPE_SIZE) {
                flush_stripe();
                if (!ok_) return;
            }
        }
    }

    void finalize() override {
        if (!ok_ || finalized_) return;
        finalized_ = true;
        // Flush last partial stripe (no padding needed — store_stripe handles it)
        if (!stripe_buf_.empty()) {
            flush_stripe();
        }
    }

    bool ok() const override { return ok_; }
    std::string error_message() const override { return error_msg_; }

private:
    void flush_stripe() {
        std::lock_guard<std::mutex> guard(cli_.mutex_);
        try {
            cli_.store_stripe(group_id_, stripe_idx_,
                              stripe_buf_.data(), stripe_buf_.size());
        } catch (const std::exception& e) {
            error_msg_ = "stripe " + std::to_string(stripe_idx_) + " failed: " + e.what();
            ok_ = false;
        }
        stripe_buf_.clear();
        ++stripe_idx_;
    }

    client& cli_;
    uint64_t group_id_;
    std::vector<uint8_t> stripe_buf_;
    uint32_t stripe_idx_ = 0;
    bool finalized_ = false;
    bool ok_ = true;
    std::string error_msg_;
};

std::unique_ptr<store_session> client::begin_store(uint64_t group_id, size_t total_size) {
    return std::make_unique<client_store_session>(*this, group_id, total_size);
}

// ---- Stripe-based retrieve ----

std::vector<uint8_t> client::retrieve(uint64_t group_id, size_t original_size) {
    std::unique_lock<std::mutex> guard(mutex_);
    std::vector<uint8_t> result;
    for (uint32_t s = 0; ; ++s) {
        std::vector<uint8_t> stripe;
        if (!retrieve_stripe(guard, group_id, s, stripe)) break;
        result.insert(result.end(), stripe.begin(), stripe.end());
    }

    if (original_size > 0 && original_size <= result.size())
        result.resize(original_size);
    return result;
}

std::vector<uint8_t> client::retrieve(uint64_t group_id) {
    std::unique_lock<std::mutex> guard(mutex_);
    std::vector<uint8_t> result;
    result.reserve(STRIPE_SIZE * 4);

    for (uint32_t s = 0; ; ++s) {
        std::vector<uint8_t> stripe;
        if (!retrieve_stripe(guard, group_id, s, stripe)) break;
        result.insert(result.end(), stripe.begin(), stripe.end());
    }

    // No stripes found → key does not exist, return empty
    if (result.empty())
        return {};

    // Strip 8-byte length header
    if (result.size() < 8)
        throw std::runtime_error("retrieve: reconstructed data too small for length header");
    uint64_t total_len = protocol::decode_u64(result.data());
    if (total_len > result.size() - 8)
        throw std::runtime_error("retrieve: invalid length header ("
                                 + std::to_string(total_len) + " > "
                                 + std::to_string(result.size() - 8) + ")");
    return std::vector<uint8_t>(result.begin() + 8, result.begin() + 8 + total_len);
}

std::vector<uint8_t> client::retrieve_range(uint64_t group_id,
                                              uint64_t offset, uint64_t length) {
    std::unique_lock<std::mutex> guard(mutex_);
    if (length == 0) return {};

    // The stored framed data is: [8-byte length header][raw data]
    // User offset 0 = byte 8 in framed stream.
    uint64_t framed_start = offset + 8;
    uint64_t framed_end = framed_start + length;

    // Compute which stripes we need
    uint32_t first_stripe = static_cast<uint32_t>(framed_start / STRIPE_SIZE);
    uint32_t last_stripe = static_cast<uint32_t>((framed_end - 1) / STRIPE_SIZE);

    // Fetch only the needed stripes
    std::vector<uint8_t> buf;
    buf.reserve(static_cast<size_t>((last_stripe - first_stripe + 1)) * STRIPE_SIZE);

    for (uint32_t s = first_stripe; s <= last_stripe; ++s) {
        std::vector<uint8_t> stripe;
        if (!retrieve_stripe(guard, group_id, s, stripe)) {
            throw std::runtime_error("retrieve_range: stripe " + std::to_string(s) + " not found");
        }
        buf.insert(buf.end(), stripe.begin(), stripe.end());
    }

    // Slice out the requested range within fetched stripes
    uint64_t buf_offset = framed_start - static_cast<uint64_t>(first_stripe) * STRIPE_SIZE;
    if (buf_offset + length > buf.size()) {
        // Clamp to available data
        length = buf.size() > buf_offset ? buf.size() - buf_offset : 0;
    }
    if (length == 0) return {};
    return std::vector<uint8_t>(buf.begin() + static_cast<std::ptrdiff_t>(buf_offset),
                                 buf.begin() + static_cast<std::ptrdiff_t>(buf_offset + length));
}

void client::remove(uint64_t group_id) {
    std::lock_guard<std::mutex> guard(mutex_);
    std::vector<uint8_t> payload(8);
    protocol::encode_u64(payload.data(), group_id);

    for (size_t i = 0; i < nodes_.size(); i++) {
        try {
            send_recv(i, msg_type::DELETE_GROUP, payload);
        } catch (...) {
            // Best-effort deletion
        }
    }
}

std::vector<std::pair<size_t, std::vector<uint32_t>>> client::list(uint64_t group_id) {
    std::lock_guard<std::mutex> guard(mutex_);
    std::vector<std::pair<size_t, std::vector<uint32_t>>> result;
    std::vector<uint8_t> payload(8);
    protocol::encode_u64(payload.data(), group_id);

    for (size_t i = 0; i < nodes_.size(); i++) {
        try {
            auto response = send_recv(i, msg_type::LIST_BLOCKS, payload);
            if (response.size() >= protocol::HEADER_SIZE) {
                msg_type type;
                uint32_t resp_len;
                protocol::decode_header(response.data(), response.size(), type, resp_len);
                if (type == msg_type::BLOCK_LIST && resp_len >= 12) {
                    const uint8_t* rp = response.data() + protocol::HEADER_SIZE;
                    uint32_t count = protocol::decode_u32(rp + 8);
                    std::vector<uint32_t> indices;
                    for (uint32_t j = 0; j < count && (12 + j * 4) < resp_len; j++) {
                        indices.push_back(protocol::decode_u32(rp + 12 + j * 4));
                    }
                    result.emplace_back(i, std::move(indices));
                }
            }
        } catch (...) {}
    }
    return result;
}

size_t client::replicate_to_peers(uint64_t group_id, uint32_t block_index,
                                 const uint8_t* data, size_t len) {
    std::lock_guard<std::mutex> guard(mutex_);
    static constexpr int MAX_RETRIES = 2;
    static constexpr int BASE_DELAY_MS = 100;
    size_t success_count = 0;

    for (size_t i = 0; i < nodes_.size(); i++) {
        bool ok = false;
        for (int attempt = 0; attempt < MAX_RETRIES && !ok; attempt++) {
            try {
                if (store_block_on_node(i, group_id, block_index, data, len)) {
                    ok = true;
                    success_count++;
                } else {
                    std::cerr << "replicate_to_peers: node " << i
                              << " store rejected (attempt " << (attempt + 1)
                              << "/" << MAX_RETRIES << ")" << std::endl;
                    connections_[i].socket.reset();
                    connections_[i].connected = false;
                    if (attempt + 1 < MAX_RETRIES) {
                        int delay = BASE_DELAY_MS * (1 << attempt);
                        std::this_thread::sleep_for(std::chrono::milliseconds(delay));
                    }
                }
            } catch (const netplus::NetException &e) {
                connections_[i].socket.reset();
                connections_[i].connected = false;
                if (attempt + 1 < MAX_RETRIES) {
                    int delay = BASE_DELAY_MS * (1 << attempt);
                    std::cerr << "replicate_to_peers: node " << i
                              << " failed (attempt " << (attempt + 1)
                              << "): " << e.what()
                              << ", retrying in " << delay << "ms" << std::endl;
                    std::this_thread::sleep_for(std::chrono::milliseconds(delay));
                } else {
                    std::cerr << "replicate_to_peers: node " << i
                              << " failed after " << MAX_RETRIES
                              << " attempts: " << e.what() << std::endl;
                }
            } catch (const std::exception &e) {
                connections_[i].socket.reset();
                connections_[i].connected = false;
                if (attempt + 1 < MAX_RETRIES) {
                    int delay = BASE_DELAY_MS * (1 << attempt);
                    std::cerr << "replicate_to_peers: node " << i
                              << " failed (attempt " << (attempt + 1)
                              << "): " << e.what()
                              << ", retrying in " << delay << "ms" << std::endl;
                    std::this_thread::sleep_for(std::chrono::milliseconds(delay));
                } else {
                    std::cerr << "replicate_to_peers: node " << i
                              << " failed after " << MAX_RETRIES
                              << " attempts: " << e.what() << std::endl;
                }
            }
        }
    }
    return success_count;
}

bool client::fetch_from_peers(uint64_t group_id, uint32_t block_index,
                               std::vector<uint8_t>& out) {
    std::lock_guard<std::mutex> guard(mutex_);
    static constexpr int MAX_RETRIES = 2;
    static constexpr int BASE_DELAY_MS = 100;

    for (size_t i = 0; i < nodes_.size(); i++) {
        for (int attempt = 0; attempt < MAX_RETRIES; attempt++) {
            try {
                if (fetch_block_from_node(i, group_id, block_index, out) && !out.empty())
                    return true;
                break; // fetch succeeded but block not found on this node, try next
            } catch (const netplus::NetException &e) {
                connections_[i].socket.reset();
                connections_[i].connected = false;
                if (attempt + 1 < MAX_RETRIES) {
                    int delay = BASE_DELAY_MS * (1 << attempt);
                    std::cerr << "fetch_from_peers: node " << i
                              << " failed (attempt " << (attempt + 1)
                              << "): " << e.what()
                              << ", retrying in " << delay << "ms" << std::endl;
                    std::this_thread::sleep_for(std::chrono::milliseconds(delay));
                } else {
                    std::cerr << "fetch_from_peers: node " << i
                              << " failed after " << MAX_RETRIES
                              << " attempts: " << e.what() << std::endl;
                }
            } catch (const std::exception &e) {
                connections_[i].socket.reset();
                connections_[i].connected = false;
                if (attempt + 1 < MAX_RETRIES) {
                    int delay = BASE_DELAY_MS * (1 << attempt);
                    std::cerr << "fetch_from_peers: node " << i
                              << " failed (attempt " << (attempt + 1)
                              << "): " << e.what()
                              << ", retrying in " << delay << "ms" << std::endl;
                    std::this_thread::sleep_for(std::chrono::milliseconds(delay));
                } else {
                    std::cerr << "fetch_from_peers: node " << i
                              << " failed after " << MAX_RETRIES
                              << " attempts: " << e.what() << std::endl;
                }
            }
        }
    }
    return false;
}

bool client::get_node_status_impl(size_t node_index, status_info& info) {
    if (node_index >= nodes_.size()) return false;
    try {
        auto resp = send_recv(node_index, msg_type::GET_STATUS, {});
        if (resp.size() < protocol::HEADER_SIZE) return false;

        msg_type rtype;
        uint32_t rlen;
        if (!protocol::decode_header(resp.data(), resp.size(), rtype, rlen))
            return false;

        if (rtype != msg_type::STATUS_RESP)
            return false;

        return protocol::parse_status_resp(
            resp.data() + protocol::HEADER_SIZE, rlen, info);
    } catch (...) {
        return false;
    }
}

bool client::get_node_status(size_t node_index, status_info& info) {
    std::lock_guard<std::mutex> guard(mutex_);
    return get_node_status_impl(node_index, info);
}

bool client::list_groups_on_node(size_t node_index, std::vector<uint64_t>& groups) {
    std::lock_guard<std::mutex> guard(mutex_);
    if (node_index >= nodes_.size()) return false;
    if (static_cast<int>(node_index) == local_node_index_ && local_store_) {
        groups = local_store_->list_groups();
        return true;
    }
    try {
        auto resp = send_recv(node_index, msg_type::LIST_GROUPS, {});
        if (resp.size() < protocol::HEADER_SIZE) return false;

        msg_type rtype;
        uint32_t rlen;
        if (!protocol::decode_header(resp.data(), resp.size(), rtype, rlen))
            return false;

        if (rtype != msg_type::GROUP_LIST)
            return false;

        return protocol::parse_group_list(
            resp.data() + protocol::HEADER_SIZE, rlen, groups);
    } catch (...) {
        return false;
    }
}

client::cluster_health client::get_cluster_status() {
    std::lock_guard<std::mutex> guard(mutex_);
    cluster_health health;
    health.nodes_total = static_cast<uint32_t>(nodes_.size());
    health.nodes_online = 0;
    health.overall = node_status::OK;

    // Probe all nodes in parallel — avoid sequential 2s waits per dead node
    struct ProbeResult {
        size_t index;
        bool reachable = false;
        status_info info;
    };
    std::vector<ProbeResult> results(nodes_.size());
    {
        std::vector<std::future<void>> futures;
        for (size_t i = 0; i < nodes_.size(); i++) {
            results[i].index = i;
            futures.push_back(std::async(std::launch::async, [this, i, &results]() {
                results[i].reachable = get_node_status_impl(i, results[i].info);
            }));
        }
        for (auto& f : futures) f.get();
    }

    for (auto& r : results) {
        if (r.reachable) {
            health.nodes_online++;
            health.node_statuses.emplace_back(r.index, r.info);
            if (r.info.status == node_status::DEGRADED ||
                r.info.status == node_status::REBUILD) {
                health.overall = r.info.status;
            }
        }
    }

    if (health.nodes_online == 0) {
        health.overall = node_status::DEGRADED;
    } else if (health.nodes_online < health.nodes_total) {
        if (health.overall == node_status::OK)
            health.overall = node_status::DEGRADED;
    }

    return health;
}

void client::set_local_node(size_t node_index, std::shared_ptr<block_store> store) {
    std::lock_guard<std::mutex> guard(mutex_);
    local_node_index_ = static_cast<int>(node_index);
    local_store_ = std::move(store);
}

void client::warmup() {
    std::lock_guard<std::mutex> guard(mutex_);
    for (size_t i = 0; i < nodes_.size(); ++i) {
        if (static_cast<int>(i) == local_node_index_) continue;
        try {
            ensure_connected(i);
        } catch (...) {
            // Best-effort: node may be offline, will retry lazily later
        }
    }
}

size_t client::vacuum_all_nodes() {
    std::lock_guard<std::mutex> guard(mutex_);
    size_t ok_count = 0;
    for (size_t i = 0; i < nodes_.size(); i++) {
        try {
            if (static_cast<int>(i) == local_node_index_ && local_store_) {
                if (local_store_->vacuum()) ++ok_count;
                continue;
            }
            auto resp = send_recv(i, msg_type::VACUUM, {});
            if (resp.size() >= protocol::HEADER_SIZE) {
                msg_type rtype;
                uint32_t rlen;
                if (protocol::decode_header(resp.data(), resp.size(), rtype, rlen)
                    && rtype == msg_type::OK) {
                    ++ok_count;
                }
            }
        } catch (...) {}
    }
    return ok_count;
}

client::rebalance_result client::rebalance() {
    rebalance_result result;
    size_t n = k_ + m_;
    uint64_t lock_gid = scrub_lock_group_id();

    // Step 1: Collect group presence per node (N round trips total, not N×G)
    std::map<uint64_t, std::set<size_t>> group_nodes; // gid → set of node indices
    for (size_t i = 0; i < nodes_.size(); ++i) {
        std::vector<uint64_t> node_groups;
        if (list_groups_on_node(i, node_groups)) {
            for (uint64_t gid : node_groups) {
                if (gid == lock_gid) continue;
                group_nodes[gid].insert(i);
            }
        }
    }

    if (group_nodes.empty())
        return result;

    std::cerr << "[PARITY] rebalance: " << group_nodes.size() << " group(s) to check" << std::endl;

    // Step 2: Groups present on all n nodes are likely OK — only do the
    // expensive per-block placement check for groups that aren't on all nodes.
    // This skips the costly list(gid) call for the majority of groups.
    std::vector<uint64_t> candidates;
    for (auto& [gid, nodes_set] : group_nodes) {
        if (nodes_set.size() >= n) {
            ++result.already_ok;
        } else {
            candidates.push_back(gid);
        }
    }

    // Step 3: Fix misplaced / under-replicated groups
    for (uint64_t gid : candidates) {
        try {
            auto data = retrieve(gid);
            if (data.empty()) {
                std::cerr << "[PARITY] rebalance: gid=" << gid
                          << " retrieve returned empty, skipping" << std::endl;
                ++result.failed;
                continue;
            }

            remove(gid);
            store(gid, data);

            ++result.rebalanced;
            std::cerr << "[PARITY] rebalance: gid=" << gid
                      << " re-stored (" << data.size() << " bytes)" << std::endl;
        } catch (const std::exception& e) {
            std::cerr << "[PARITY] rebalance: gid=" << gid
                      << " failed: " << e.what() << std::endl;
            ++result.failed;
        }
    }

    // Step 4: Vacuum all nodes to reclaim space from deleted blocks
    if (result.rebalanced > 0)
        vacuum_all_nodes();

    std::cerr << "[PARITY] rebalance done: " << result.rebalanced << " rebalanced, "
              << result.already_ok << " ok, " << result.failed << " failed" << std::endl;

    return result;
}

void client::mark_node_dead(size_t node_index) {
    std::lock_guard<std::mutex> guard(mutex_);
    if (node_index < connections_.size()) {
        connections_[node_index].socket.reset();
        connections_[node_index].connected = false;
        connections_[node_index].dead_until =
            std::chrono::steady_clock::now() + std::chrono::seconds(10);
    }
}

// --- Distributed scrub lock ---

uint64_t client::scrub_lock_group_id() {
    // Deterministic well-known ID from fixed string
    const char key[] = "__paritypp_scrub_lock__";
    uint64_t h = 0xcbf29ce484222325ULL; // FNV-1a
    for (const char *p = key; *p; ++p) {
        h ^= static_cast<uint8_t>(*p);
        h *= 0x100000001b3ULL;
    }
    return h;
}

bool client::try_scrub_lock(int lock_ttl_seconds) {
    uint64_t gid = scrub_lock_group_id();
    // Check existing lock
    try {
        auto existing = retrieve(gid);
        if (existing.size() == sizeof(int64_t)) {
            int64_t stored_ts = 0;
            std::memcpy(&stored_ts, existing.data(), sizeof(int64_t));
            int64_t now = std::chrono::duration_cast<std::chrono::seconds>(
                std::chrono::system_clock::now().time_since_epoch()).count();
            if (now - stored_ts < lock_ttl_seconds) {
                DBG_LOG("[PARITY] scrub lock held (" << (now - stored_ts) << "s ago)\n");
                return false;
            }
        }
    } catch (...) {}
    // Write our timestamp as the lock
    int64_t now = std::chrono::duration_cast<std::chrono::seconds>(
        std::chrono::system_clock::now().time_since_epoch()).count();
    try {
        store(gid, reinterpret_cast<const uint8_t*>(&now), sizeof(int64_t));
    } catch (...) {
        return false;
    }
    return true;
}

void client::release_scrub_lock() {
    try { remove(scrub_lock_group_id()); } catch (...) {}
}

// --- Batch API implementation ---

static constexpr uint8_t BATCH_MAGIC[4] = {'B', 'T', 'C', 'H'};

std::vector<uint8_t> client::encode_batch(const std::vector<batch_item>& items) {
    // Calculate total size: 4B magic + 4B count + per-item (8B key + 4B len + data)
    size_t total = 8;
    for (const auto& item : items)
        total += 12 + item.data.size();

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

    // Magic
    std::memcpy(payload.data() + pos, BATCH_MAGIC, 4);
    pos += 4;

    // Item count
    uint32_t count = static_cast<uint32_t>(items.size());
    protocol::encode_u32(payload.data() + pos, count);
    pos += 4;

    // Items
    for (const auto& item : items) {
        protocol::encode_u64(payload.data() + pos, item.key);
        pos += 8;
        uint32_t dlen = static_cast<uint32_t>(item.data.size());
        protocol::encode_u32(payload.data() + pos, dlen);
        pos += 4;
        if (dlen > 0) {
            std::memcpy(payload.data() + pos, item.data.data(), dlen);
            pos += dlen;
        }
    }
    return payload;
}

bool client::decode_batch(const std::vector<uint8_t>& payload,
                           std::vector<batch_item>& items) {
    items.clear();
    if (payload.size() < 8) return false;

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

    uint32_t count = protocol::decode_u32(payload.data() + 4);
    size_t pos = 8;

    items.reserve(count);
    for (uint32_t i = 0; i < count; ++i) {
        if (pos + 12 > payload.size()) return false;
        batch_item item;
        item.key = protocol::decode_u64(payload.data() + pos);
        pos += 8;
        uint32_t dlen = protocol::decode_u32(payload.data() + pos);
        pos += 4;
        if (pos + dlen > payload.size()) return false;
        item.data.assign(payload.data() + pos, payload.data() + pos + dlen);
        pos += dlen;
        items.push_back(std::move(item));
    }
    return true;
}

void client::store_batch(uint64_t batch_gid, const std::vector<batch_item>& items) {
    auto payload = encode_batch(items);
    store(batch_gid, payload.data(), payload.size());
}

std::vector<client::batch_item> client::retrieve_batch(uint64_t batch_gid) {
    auto payload = retrieve(batch_gid);
    std::vector<batch_item> items;
    if (!decode_batch(payload, items))
        throw std::runtime_error("invalid batch payload");
    return items;
}

bool client::retrieve_batch_item(uint64_t batch_gid, uint64_t item_key,
                                  std::vector<uint8_t>& out) {
    auto payload = retrieve(batch_gid);
    std::vector<batch_item> items;
    if (!decode_batch(payload, items)) return false;
    for (const auto& item : items) {
        if (item.key == item_key) {
            out = item.data;
            return true;
        }
    }
    return false;
}

} // namespace paritypp
