#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 <array>
#include <algorithm>
#include <limits>
#include <cstdlib>
#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);
}

// Enable extended runtime diagnostics when PARITYPP_DIAG is set to a non-zero
// value. This keeps default logs quiet while allowing deep timeout tracing.
static bool parity_diag_enabled() {
    static int cached = -1;
    if (cached >= 0) return cached == 1;
    const char* v = std::getenv("PARITYPP_DIAG");
    if (!v || v[0] == '\0' || v[0] == '0') {
        cached = 0;
    } else {
        cached = 1;
    }
    return cached == 1;
}

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>();
    // nonblock=true: connect() must return immediately after sending the
    // Initial packet, not block internally in completeHandshake() (which
    // has its own hardcoded 5s timeout, unrelated to hs_deadline below and
    // dead-on-failure since it throws before this function's own poll loop
    // ever runs). The loop right below is what's actually supposed to
    // drive the handshake to completion via pumpNetwork().
    nc.socket->connect(nodes_[node_index].address, nodes_[node_index].port, true);

    // Wait for handshake using a local socketwait (thread-safe)
    netplus::socketwait local_wait;
    auto hs_deadline = std::chrono::steady_clock::now() + std::chrono::seconds(5);
    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
            nc.consecutive_failures = 0;
            nc.last_activity = std::chrono::steady_clock::now();
            // Authenticate if credentials provided
            if (!creds_.client_name.empty()) {
                if (!authenticate_node(node_index)) {
                    mark_connection_failed(node_index);
                    throw std::runtime_error("Authentication failed on node " +
                                              nodes_[node_index].address);
                }
            }
            return;
        }
    }
    mark_connection_failed(node_index);
    throw std::runtime_error("Failed to connect to node " + nodes_[node_index].address);
}

void client::mark_connection_failed(size_t node_index) {
    auto& nc = connections_[node_index];
    nc.socket.reset();
    nc.connected = false;
    // If we're still within a previous cooldown, this call is just the
    // fallout of ensure_connected's "still dead" throw on a retry that
    // never actually attempted a new connection — not a fresh failure.
    // Don't extend the cooldown further, or a node that fails once would
    // never see its cooldown expire: every subsequent retry re-triggers
    // this same catch block before the cooldown has a chance to elapse,
    // pushing dead_until forward forever.
    auto now = std::chrono::steady_clock::now();
    if (nc.dead_until > now) return;
    ++nc.consecutive_failures;
    // 1s, 2s, 4s, 8s, capped at the old flat 10s from the 5th failure on.
    int shift = std::min(nc.consecutive_failures - 1, 4);
    auto cooldown = std::chrono::seconds(1 << shift);
    if (cooldown > std::chrono::seconds(10)) cooldown = std::chrono::seconds(10);
    nc.dead_until = now + cooldown;
}

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) {
        // Nothing pumps this connection while it's not actively in use, so
        // the QUIC keepalive PING (socket.h, fires after 15s idle) never
        // gets to run between infrequent calls — past _max_idle_timeout
        // (30s) either side may have silently discarded it already. Treat
        // "idle a while" the same as "known dead" and reconnect now, rather
        // than discover it via a wasted 5s send_recv timeout on first use.
        static constexpr auto IDLE_RECONNECT_THRESHOLD = std::chrono::seconds(20);
        auto now = std::chrono::steady_clock::now();
        if (now - nc.last_activity < IDLE_RECONNECT_THRESHOLD) {
            // Refresh on reuse, not just on (re)connect — otherwise a
            // connection genuinely kept busy by repeated calls spaced under
            // the threshold apart would still look idle-since-first-connect
            // and eventually get force-reconnected despite never actually
            // going idle.
            nc.last_activity = now;
            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, store_id_);

        int send_timeout_s = 5;
        int response_timeout_s = 5;
        if (type == msg_type::STORE_BLOCK) {
            size_t payload_mib = (payload.size() + ((1024 * 1024) - 1)) / (1024 * 1024);
            send_timeout_s = 10 + static_cast<int>(payload_mib) * 2;
            response_timeout_s = 45 + static_cast<int>(payload_mib) * 10;
            if (send_timeout_s > 45) send_timeout_s = 45;
            if (response_timeout_s > 180) response_timeout_s = 180;
        }
        auto send_deadline = std::chrono::steady_clock::now()
                           + std::chrono::seconds(send_timeout_s);
        auto response_deadline = std::chrono::steady_clock::now()
                               + std::chrono::seconds(response_timeout_s);

        bool sent = false;
        while (int poll_ms = ms_left(send_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 (...) {}
            if (attempt == 0) {
                connections_[node_index].socket.reset();
                connections_[node_index].connected = false;
                continue;
            }
            mark_connection_failed(node_index);
            return {};
        }

        // Wait for response. Handle partial header/body reads explicitly to
        // avoid consuming a fragmented header and then timing out.
        std::vector<uint8_t> response;
        std::array<uint8_t, protocol::HEADER_SIZE> hdr{};
        size_t hdr_read = 0;
        uint32_t resp_len = 0;
        size_t body_read = 0;
        bool header_done = false;
        while (int poll_ms = ms_left(response_deadline)) {
            if (!local_wait.waitRead(conn, poll_ms)) continue;
            conn.pumpNetwork(MSG_DONTWAIT);
            if (!conn.hasStreamData(stream)) continue;

            if (!header_done) {
                while (hdr_read < protocol::HEADER_SIZE && conn.hasStreamData(stream)) {
                    size_t rd = conn.recvStreamData(stream,
                                                    hdr.data() + hdr_read,
                                                    protocol::HEADER_SIZE - hdr_read);
                    if (rd == 0) break;
                    hdr_read += rd;
                }
                if (hdr_read < protocol::HEADER_SIZE) continue;

                msg_type resp_type = msg_type::PROTOCOL_ERROR;
                if (!protocol::decode_header(hdr.data(), hdr_read, resp_type, resp_len)) {
                    break;
                }
                response.resize(protocol::HEADER_SIZE + resp_len);
                std::memcpy(response.data(), hdr.data(), protocol::HEADER_SIZE);
                header_done = true;
            }

            while (body_read < resp_len) {
                if (!conn.hasStreamData(stream)) break;
                size_t rd = conn.recvStreamData(stream,
                                                response.data() + protocol::HEADER_SIZE + body_read,
                                                resp_len - body_read);
                if (rd == 0) break;
                body_read += rd;
            }
            if (header_done && body_read >= resp_len) break;
        }
        conn.closeStream(stream);

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

        if (parity_diag_enabled()) {
            std::cerr << "[PARITY-DIAG] send_recv empty response: node=" << node_index
                      << " type=" << static_cast<int>(type)
                      << " header_done=" << header_done
                      << " hdr_read=" << hdr_read
                      << " body_read=" << body_read
                      << " resp_len=" << resp_len << "\n";
        }

        // Empty response — stale connection, retry once with fresh connection
        if (attempt > 0) {
            mark_connection_failed(node_index);
        } else {
            connections_[node_index].socket.reset();
            connections_[node_index].connected = false;
        }
    }
    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) {
    auto stripe_begin = std::chrono::steady_clock::now();
    // 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_;

    // A full stripe write includes stream open, write scheduling, remote block
    // store, and response collection across all nodes. Using a single short
    // deadline for both phases causes false response timeouts under import
    // load: the send phase eats into the time budget needed for remote disk
    // write + ACK. Keep separate deadlines so successful sends still get a
    // generous response window.
    size_t stripe_mib = (block_align + ((1024 * 1024) - 1)) / (1024 * 1024);
    int send_timeout_s = 8 + static_cast<int>(stripe_mib) * 2 + static_cast<int>(m_);
    int response_timeout_s = 30 + static_cast<int>(stripe_mib) * 8 + static_cast<int>(m_) * 2;
    if (send_timeout_s < 12) send_timeout_s = 12;
    if (send_timeout_s > 45) send_timeout_s = 45;
    if (response_timeout_s < 45) response_timeout_s = 45;
    if (response_timeout_s > 180) response_timeout_s = 180;
    auto send_deadline = std::chrono::steady_clock::now()
                       + std::chrono::seconds(send_timeout_s);
    auto response_deadline = std::chrono::steady_clock::now()
                           + std::chrono::seconds(response_timeout_s);

    if (parity_diag_enabled()) {
        std::cerr << "[PARITY-DIAG] store_stripe begin gid=" << group_id
                  << " stripe=" << stripe_index
                  << " stripe_len=" << stripe_len
                  << " block_align=" << block_align
                  << " nodes=" << n
                  << " k=" << k_ << " m=" << m_
                  << " send_timeout_s=" << send_timeout_s
                  << " response_timeout_s=" << response_timeout_s << "\n";
    }

    // 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::array<uint8_t, protocol::HEADER_SIZE> resp_hdr{};
        size_t resp_hdr_read = 0;
        bool resp_header_done = false;
        msg_type resp_type = msg_type::PROTOCOL_ERROR;
        uint32_t resp_len = 0;
        size_t resp_body_read = 0;
    };
    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;
        }

        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, store_id_);
            uint64_t stream = conn.openStream(true);
            pending[i].stream = stream;

            // Non-blocking send
            while (int poll_ms = ms_left(send_deadline)) {
                if (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();
            mark_connection_failed(i);
        } catch (const std::exception& e) {
            pending[i].fail_reason = std::string("exception: ") + e.what();
            mark_connection_failed(i);
        } catch (...) {
            pending[i].fail_reason = "unknown exception";
            mark_connection_failed(i);
        }
    };

    // Phase 1: store local block + send to remote nodes concurrently.
    // All remote sends are interleaved (round-robin waitWrite/pump) so
    // QUIC congestion windows on independent connections grow in parallel
    // instead of serialising the full RTT per node.
    {
        // Local node first (synchronous, no network)
        for (size_t i = 0; i < n; i++) {
            if (static_cast<int>(i) == local_node_index_ && local_store_) {
                send_to_node(i);
            }
        }

        // Prepare remote sends: open streams + build messages, but don't block on waitWrite yet
        struct RemoteSend {
            size_t idx;
            std::vector<uint8_t> msg;
            bool ready = false;   // stream opened + message built
        };
        std::vector<RemoteSend> remotes;
        remotes.reserve(n);
        for (size_t i = 0; i < n; i++) {
            if (static_cast<int>(i) == local_node_index_ && local_store_) continue;
            RemoteSend rs;
            rs.idx = i;
            pending[i].node_index = i;
            pending[i].sent = false;
            try {
                ensure_connected(i);
                if (!connections_[i].connected || !connections_[i].socket) {
                    pending[i].fail_reason = "connect failed (handshake timeout)";
                    remotes.push_back(std::move(rs));
                    continue;
                }
                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);
                auto payload = protocol::make_store_block(group_id, block_index, bdata, blen);
                rs.msg = protocol::encode_message(msg_type::STORE_BLOCK, payload, store_id_);
                auto& conn = *connections_[i].socket;
                uint64_t stream = conn.openStream(true);
                pending[i].stream = stream;
                rs.ready = true;
            } catch (const netplus::NetException& e) {
                pending[i].fail_reason = std::string("net: ") + e.what();
                mark_connection_failed(i);
            } catch (const std::exception& e) {
                pending[i].fail_reason = std::string("exception: ") + e.what();
                mark_connection_failed(i);
            } catch (...) {
                pending[i].fail_reason = "unknown exception";
                mark_connection_failed(i);
            }
            remotes.push_back(std::move(rs));
        }

        // Interleaved send: round-robin across all ready remote nodes
        // until all have sent or deadline expires.
        size_t unsent = 0;
        for (auto& rs : remotes) { if (rs.ready && !pending[rs.idx].sent) ++unsent; }

        while (unsent > 0) {
            int poll_ms = ms_left(send_deadline);
            if (poll_ms <= 0) break;
            bool progress = false;
            for (auto& rs : remotes) {
                if (!rs.ready || pending[rs.idx].sent) continue;
                auto& conn = *connections_[rs.idx].socket;
                try {
                    conn.pumpNetwork(MSG_DONTWAIT);
                    if (wait_.waitWrite(conn, std::min(poll_ms / std::max<int>((int)unsent, 1), 50))) {
                        conn.sendStreamData(pending[rs.idx].stream, rs.msg, true);
                        conn.pumpNetwork(MSG_DONTWAIT);
                        pending[rs.idx].sent = true;
                        --unsent;
                        progress = true;
                    }
                } catch (const netplus::NetException& e) {
                    pending[rs.idx].fail_reason = std::string("send net: ") + e.what();
                    mark_connection_failed(rs.idx);
                    rs.ready = false;
                    --unsent;
                } catch (const std::exception& e) {
                    pending[rs.idx].fail_reason = std::string("send exception: ") + e.what();
                    mark_connection_failed(rs.idx);
                    rs.ready = false;
                    --unsent;
                } catch (...) {
                    pending[rs.idx].fail_reason = "send unknown exception";
                    mark_connection_failed(rs.idx);
                    rs.ready = false;
                    --unsent;
                }
            }
            // Pump all ready connections to advance QUIC state
            if (!progress) {
                for (auto& rs : remotes) {
                    if (!rs.ready || pending[rs.idx].sent) continue;
                    if (!connections_[rs.idx].socket) continue;
                    try { connections_[rs.idx].socket->pumpNetwork(MSG_DONTWAIT); } catch (...) {}
                }
            }
        }
        // Mark timeout for nodes that never sent
        for (auto& rs : remotes) {
            if (rs.ready && !pending[rs.idx].sent) {
                pending[rs.idx].fail_reason = "send timeout (interleaved waitWrite)";
                try { connections_[rs.idx].socket->closeStream(pending[rs.idx].stream); } catch (...) {}
            }
        }
    }

    // Phase 1b: retry failed remote nodes once with a fresh connection
    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;
            send_to_node(i);
        }
    }

    if (parity_diag_enabled()) {
        size_t sent_remote = 0;
        for (auto& p : pending) if (p.sent && !p.local) ++sent_remote;
        std::cerr << "[PARITY-DIAG] store_stripe after send gid=" << group_id
                  << " stripe=" << stripe_index
                  << " remote_sent=" << sent_remote
                  << " remote_total=" << (n - ((local_node_index_ >= 0 && static_cast<size_t>(local_node_index_) < n) ? 1 : 0))
                  << "\n";
    }

    // 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(response_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 {
                if (!p.resp_header_done) {
                    while (p.resp_hdr_read < protocol::HEADER_SIZE && conn.hasStreamData(p.stream)) {
                        size_t rd = conn.recvStreamData(p.stream,
                                                        p.resp_hdr.data() + p.resp_hdr_read,
                                                        protocol::HEADER_SIZE - p.resp_hdr_read);
                        if (rd == 0) break;
                        p.resp_hdr_read += rd;
                    }
                    if (p.resp_hdr_read < protocol::HEADER_SIZE) continue;

                    if (!protocol::decode_header(p.resp_hdr.data(), p.resp_hdr_read,
                                                 p.resp_type, p.resp_len)) {
                        p.fail_reason = "invalid response header";
                        conn.closeStream(p.stream);
                        --remaining;
                        continue;
                    }
                    p.resp_header_done = true;
                }

                while (p.resp_body_read < p.resp_len) {
                    if (!conn.hasStreamData(p.stream)) break;
                    uint8_t discard[1024];
                    size_t chunk = std::min<size_t>(sizeof(discard), p.resp_len - p.resp_body_read);
                    size_t rd = conn.recvStreamData(p.stream, discard, chunk);
                    if (rd == 0) break;
                    p.resp_body_read += rd;
                }

                if (p.resp_body_read < p.resp_len) continue;

                p.ok = (p.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 ALL pending connections — not just ready ones.
        // This ensures PTO (Probe Timeout) retransmissions fire even when
        // the peer hasn't sent any data (tail loss scenario).
        for (size_t wi = 0; wi < wait_socks.size(); ++wi) {
            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";
            if (p.resp_hdr_read > 0 || p.resp_header_done || p.resp_body_read > 0) {
                p.fail_reason += " (hdr=" + std::to_string(p.resp_hdr_read) +
                                 "/" + std::to_string(protocol::HEADER_SIZE) +
                                 ", body=" + std::to_string(p.resp_body_read) +
                                 "/" + std::to_string(p.resp_len) + ")";
            }
        }
    }

    // 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 (...) {}
        }
    }

    // A response-timeout indicates the connection can be half-open/stuck.
    // Force reconnect on next attempt instead of reusing a poisoned session.
    for (auto& p : pending) {
        if (p.local || p.ok) continue;
        if (p.fail_reason.find("response timeout") != std::string::npos) {
            connections_[p.node_index].socket.reset();
            connections_[p.node_index].connected = false;
            connections_[p.node_index].dead_until = {};
        }
    }

    // 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_) {
        if (parity_diag_enabled()) {
            auto elapsed_ms = std::chrono::duration_cast<std::chrono::milliseconds>(
                std::chrono::steady_clock::now() - stripe_begin).count();
            std::cerr << "[PARITY-DIAG] store_stripe fail gid=" << group_id
                      << " stripe=" << stripe_index
                      << " failed=" << failed << "/" << n
                      << " elapsed_ms=" << elapsed_ms
                      << " diag=" << diag << "\n";
        }
        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);
    }
    if (failed > 0) {
        std::cerr << "[PARITY] store_stripe WARNING: " << failed << "/" << n
                  << " node(s) failed for gid=" << group_id << " stripe="
                  << stripe_index << ": " << diag << "\n";
    }
}

bool client::retrieve_stripe(std::unique_lock<std::mutex>& lk,
                              uint64_t group_id, uint32_t stripe_index,
                              std::vector<uint8_t>& out,
                              std::chrono::steady_clock::time_point deadline,
                              bool* definitively_not_found) {
    // 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.
    // Bounded by `deadline` in addition to MAX_RETRIES, so a caller with a
    // tight overall time budget (e.g. Cluster::fetch()) doesn't have this
    // alone eat the full ~30s regardless of how much time it has left.
    static constexpr int MAX_RETRIES = 6;           // up to ~30s total
    static constexpr int INITIAL_WAIT_MS = 500;     // first retry after 500ms

    // Every early return below is a "gave up without a definitive answer"
    // case (deadline/retries exhausted while some node stayed ambiguous) —
    // only the should_retry==false branch inside the loop overwrites this
    // with true (explicit misses leave fewer than k possible shards).
    if (definitively_not_found) *definitively_not_found = false;
    std::vector<bool> confirmed_not_found(k_ + m_, false);

    for (int attempt = 0; attempt <= MAX_RETRIES; ++attempt) {
        if (std::chrono::steady_clock::now() >= deadline) return false;

        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));
            auto remaining_ms = std::chrono::duration_cast<std::chrono::milliseconds>(
                deadline - std::chrono::steady_clock::now()).count();
            if (remaining_ms <= 0) return false;
            wait_ms = static_cast<int>(std::min<long long>(wait_ms, remaining_ms));
            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,
                                  static_cast<size_t>(attempt),
                                  confirmed_not_found, deadline))  {
            if (!should_retry) {
                // Every node gave an explicit not-found answer — genuine,
                // retrying would just get the same answer again.
                if (definitively_not_found) *definitively_not_found = true;
                return false;
            }
            if (attempt == MAX_RETRIES) return false;  // still ambiguous, out of attempts
            // 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,
    size_t candidate_offset,
    std::vector<bool>& confirmed_not_found,
    std::chrono::steady_clock::time_point deadline)
{
    should_retry = false;

    const size_t n = k_ + m_;
    const auto now = std::chrono::steady_clock::now();

    auto retr_deadline = std::min(
        now + std::chrono::seconds(10),
        deadline
    );

    struct Pending {
        size_t node_index = 0;
        uint64_t stream = 0;

        bool local = false;
        bool sent = false;
        bool finished = false;
        bool ok = false;
        bool not_found = false;

        std::vector<uint8_t> block;

        std::array<uint8_t, protocol::HEADER_SIZE> resp_hdr{};
        size_t resp_hdr_read = 0;
        bool resp_header_done = false;

        msg_type resp_type = msg_type::PROTOCOL_ERROR;
        uint32_t resp_len = 0;

        std::vector<uint8_t> resp_body;
        size_t resp_body_read = 0;

        std::string fail_reason;
    };

    std::vector<Pending> pending(n);

    for (size_t i = 0; i < n; ++i) {
        pending[i].node_index = i;
    }

    size_t total_ok = 0;

    /*
     * Read the local shard first.
     *
     * The local block is fetched synchronously because it is always the
     * lowest-latency source. If the local shard is available, it immediately
     * contributes towards the required k blocks.
     */
    if (local_node_index_ >= 0 &&
        static_cast<size_t>(local_node_index_) < n &&
        local_store_)
    {
        const size_t i = static_cast<size_t>(local_node_index_);
        const uint32_t block_index =
            stripe_index * static_cast<uint32_t>(n) +
            static_cast<uint32_t>(i);

        auto& p = pending[i];

        p.local = true;
        p.finished = true;

        try {
            p.ok = local_store_->fetch(
                group_id,
                block_index,
                p.block
            ) && !p.block.empty();

            p.not_found = !p.ok;

            if (p.ok) {
                ++total_ok;
            } else {
                p.fail_reason = "local block not found";
            }
        } catch (const std::exception& e) {
            p.fail_reason =
                std::string("local fetch exception: ") + e.what();
        } catch (...) {
            p.fail_reason = "local fetch unknown exception";
        }
    }

    /*
     * Build the remote candidate list.
     *
     * Priority:
     *   1. Already connected peers
     *   2. Peers that are not connected yet but are available
     *   3. Nodes currently in reconnect cooldown are skipped
     *
     * candidate_offset rotates equally ranked nodes between retries so that
     * repeated attempts distribute load across the cluster.
     */
    std::vector<size_t> remote_order;
    remote_order.reserve(n);

    auto add_candidates = [&](bool connected_only) {
        for (size_t step = 0; step < n; ++step) {
            const size_t i = (candidate_offset + step) % n;

            if (static_cast<int>(i) == local_node_index_ && local_store_) {
                continue;
            }

            const auto& nc = connections_[i];

            const bool connected =
                nc.connected && nc.socket != nullptr;

            const bool cooling_down =
                nc.dead_until > std::chrono::steady_clock::now();

            if (cooling_down) {
                continue;
            }

            if (connected_only != connected) {
                continue;
            }

            remote_order.push_back(i);
        }
    };

    add_candidates(true);
    add_candidates(false);

    /*
     * Mark nodes currently in reconnect cooldown.
     *
     * These nodes are intentionally skipped for this retrieval attempt in
     * order to avoid spending time on connections that recently failed.
     */
    for (size_t i = 0; i < n; ++i) {
        if (static_cast<int>(i) == local_node_index_ && local_store_) {
            continue;
        }

        if (connections_[i].dead_until > std::chrono::steady_clock::now()) {
            pending[i].fail_reason = "node dead (cooldown)";
        }
    }

    /*
     * Open a FETCH_BLOCK stream on a remote node.
     *
     * A connection is established if necessary, a new QUIC stream is opened,
     * and the FETCH_BLOCK request is transmitted. Any connection or transport
     * failure immediately marks the node as failed so future retries can
     * deprioritize it.
     */
    auto start_fetch = [&](size_t i) {
        auto& p = pending[i];

        const uint32_t block_index =
            stripe_index * static_cast<uint32_t>(n) +
            static_cast<uint32_t>(i);

        try {
            ensure_connected(i);

            if (!connections_[i].connected ||
                !connections_[i].socket)
            {
                p.fail_reason = "connect failed";
                p.finished = true;
                return;
            }

            auto& conn = *connections_[i].socket;

            const auto payload =
                protocol::make_fetch_block(
                    group_id,
                    block_index
                );

            const auto message =
                protocol::encode_message(
                    msg_type::FETCH_BLOCK,
                    payload,
                    store_id_
                );

            p.stream = conn.openStream(true);

            netplus::socketwait local_wait;

            while (int poll_ms = ms_left(retr_deadline)) {
                try {
                    conn.pumpNetwork(MSG_DONTWAIT);
                } catch (...) {
                }

                if (!local_wait.waitWrite(conn, poll_ms)) {
                    continue;
                }

                conn.sendStreamData(
                    p.stream,
                    message,
                    true
                );

                conn.pumpNetwork(MSG_DONTWAIT);

                p.sent = true;
                return;
            }

            p.fail_reason = "send timeout";

            try {
                conn.closeStream(p.stream);
            } catch (...) {
            }

            p.finished = true;
        } catch (const netplus::NetException& e) {
            p.fail_reason =
                std::string("net: ") + e.what();
            p.finished = true;
            mark_connection_failed(i);
        } catch (const std::exception& e) {
            p.fail_reason =
                std::string("exception: ") + e.what();
            p.finished = true;
            mark_connection_failed(i);
        } catch (...) {
            p.fail_reason = "unknown exception";
            p.finished = true;
            mark_connection_failed(i);
        }
    };

    /*
     * Send FETCH requests to every available node.
     *
     * Previous versions contacted only the minimum number of nodes required
     * to theoretically obtain k blocks. That strategy caused degraded reads
     * to fail when one selected node was unreachable, even though sufficient
     * parity information existed elsewhere.
     *
     * The current implementation sends requests to every available node,
     * excluding nodes in reconnect cooldown. The receive loop terminates as
     * soon as k valid blocks have been collected.
     */
    for (const size_t i : remote_order) {
        if (std::chrono::steady_clock::now() >= retr_deadline) {
            if (pending[i].fail_reason.empty()) {
                pending[i].fail_reason =
                    "not started: retrieve deadline reached";
            }
            continue;
        }

        start_fetch(i);
    }

    size_t remaining = 0;

    for (auto& p : pending) {
        if (p.sent && !p.local && !p.finished) {
            ++remaining;
        }
    }

    /*
     * Pump every active QUIC connection once.
     *
     * This allows acknowledgements, retransmissions, and PTO processing to
     * progress before waiting for incoming responses.
     */
    for (auto& p : pending) {
        if (!p.sent || p.local || p.finished) {
            continue;
        }

        if (!connections_[p.node_index].socket) {
            continue;
        }

        try {
            connections_[p.node_index].socket
                ->pumpNetwork(MSG_DONTWAIT);
        } catch (...) {
        }
    }

    /*
     * Collect responses until k valid blocks have been received.
     *
     * Responses may arrive in arbitrary order. As soon as k valid shards have
     * been collected, retrieval can stop because Reed-Solomon decoding only
     * requires any k surviving shards.
     */
    while (remaining > 0 && total_ok < k_) {
        const int poll_ms = ms_left(retr_deadline);

        if (poll_ms <= 0) {
            break;
        }

        bool progress = false;

        for (auto& p : pending) {
            if (!p.sent ||
                p.local ||
                p.finished)
            {
                continue;
            }

            const size_t i = p.node_index;

            if (!connections_[i].socket) {
                p.fail_reason =
                    "socket reset before response";
                p.finished = true;
                p.sent = false;

                if (remaining > 0) {
                    --remaining;
                }

                continue;
            }

            auto& conn = *connections_[i].socket;

            try {
                conn.pumpNetwork(MSG_DONTWAIT);

                if (!conn.hasStreamData(p.stream)) {
                    continue;
                }

                progress = true;

                if (!p.resp_header_done) {
                    while (
                        p.resp_hdr_read < protocol::HEADER_SIZE &&
                        conn.hasStreamData(p.stream))
                    {
                        const size_t read =
                            conn.recvStreamData(
                                p.stream,
                                p.resp_hdr.data() + p.resp_hdr_read,
                                protocol::HEADER_SIZE - p.resp_hdr_read
                            );

                        if (read == 0) {
                            break;
                        }

                        p.resp_hdr_read += read;
                    }

                    if (p.resp_hdr_read < protocol::HEADER_SIZE) {
                        continue;
                    }

                    if (!protocol::decode_header(
                            p.resp_hdr.data(),
                            p.resp_hdr_read,
                            p.resp_type,
                            p.resp_len))
                    {
                        p.fail_reason =
                            "invalid response header";
                        p.finished = true;

                        try {
                            conn.closeStream(p.stream);
                        } catch (...) {
                        }

                        if (remaining > 0) {
                            --remaining;
                        }

                        continue;
                    }

                    p.resp_header_done = true;
                    p.resp_body.resize(p.resp_len);
                }

                while (
                    p.resp_body_read < p.resp_len &&
                    conn.hasStreamData(p.stream))
                {
                    const size_t read =
                        conn.recvStreamData(
                            p.stream,
                            p.resp_body.data() + p.resp_body_read,
                            p.resp_len - p.resp_body_read
                        );

                    if (read == 0) {
                        break;
                    }

                    p.resp_body_read += read;
                }

                if (p.resp_body_read < p.resp_len) {
                    continue;
                }

                if (p.resp_type == msg_type::BLOCK_DATA &&
                    p.resp_len >= 12)
                {
                    const uint8_t* block_data = nullptr;
                    size_t block_size = 0;
                    uint64_t response_group_id = 0;
                    uint32_t response_block_index = 0;

                    const bool valid =
                        protocol::parse_block_data(
                            p.resp_body.data(),
                            p.resp_len,
                            response_group_id,
                            response_block_index,
                            block_data,
                            block_size
                        );

                    const uint32_t expected_block_index =
                        stripe_index * static_cast<uint32_t>(n) +
                        static_cast<uint32_t>(i);

                    if (!valid) {
                        p.fail_reason =
                            "invalid BLOCK_DATA response";
                    } else if (
                        response_group_id != group_id ||
                        response_block_index != expected_block_index)
                    {
                        p.fail_reason =
                            "BLOCK_DATA identity mismatch";
                    } else if (block_size == 0) {
                        p.fail_reason =
                            "empty BLOCK_DATA response";
                    } else {
                        p.block.assign(
                            block_data,
                            block_data + block_size
                        );

                        p.ok = true;
                        confirmed_not_found[i] = false;
                        ++total_ok;
                    }
                } else if (
                    p.resp_type == msg_type::BLOCK_NOT_FOUND)
                {
                    p.not_found = true;
                    confirmed_not_found[i] = true;
                    p.fail_reason = "block not found";
                } else {
                    p.fail_reason =
                        "unexpected response type " +
                        std::to_string(
                            static_cast<int>(p.resp_type)
                        );
                }

                p.finished = true;

                try {
                    conn.closeStream(p.stream);
                } catch (...) {
                }

                if (remaining > 0) {
                    --remaining;
                }
            } catch (const netplus::NetException& e) {
                p.fail_reason =
                    std::string("recv net: ") + e.what();
                p.finished = true;

                if (remaining > 0) {
                    --remaining;
                }

                mark_connection_failed(i);
            } catch (const std::exception& e) {
                p.fail_reason =
                    std::string("recv exception: ") + e.what();
                p.finished = true;

                if (remaining > 0) {
                    --remaining;
                }

                mark_connection_failed(i);
            } catch (...) {
                p.fail_reason =
                    "recv unknown exception";
                p.finished = true;

                if (remaining > 0) {
                    --remaining;
                }

                mark_connection_failed(i);
            }
        }

        if (total_ok >= k_ || remaining == 0) {
            break;
        }

        /*
         * Wait for activity on all outstanding sockets.
         *
         * Even if poll() reports no readable socket, every QUIC connection is
         * pumped afterwards to ensure PTO timers, retransmissions, and ACK
         * handling continue to make forward progress.
         */
        std::vector<netplus::socket*> wait_sockets;
        std::vector<size_t> wait_indices;

        for (size_t i = 0; i < pending.size(); ++i) {
            auto& p = pending[i];

            if (!p.sent ||
                p.local ||
                p.finished ||
                !connections_[p.node_index].socket)
            {
                continue;
            }

            wait_sockets.push_back(
                connections_[p.node_index].socket.get()
            );

            wait_indices.push_back(i);
        }

        if (wait_sockets.empty()) {
            break;
        }

        std::vector<bool> ready;

        wait_.waitReadMulti(
            wait_sockets,
            ready,
            std::min(poll_ms, 50)
        );

        for (const size_t pending_index : wait_indices) {
            auto& p = pending[pending_index];

            if (!connections_[p.node_index].socket) {
                continue;
            }

            try {
                connections_[p.node_index].socket
                    ->pumpNetwork(MSG_DONTWAIT);
            } catch (...) {
            }
        }

        if (!progress &&
            std::chrono::steady_clock::now() >= retr_deadline)
        {
            break;
        }
    }

    /*
     * Close any streams that did not complete before the retrieval deadline.
     *
     * Connections without a valid response are discarded rather than reused
     * because their protocol state may no longer be reliable.
     */
    for (auto& p : pending) {
        if (!p.sent ||
            p.local ||
            p.finished ||
            !connections_[p.node_index].socket)
        {
            continue;
        }

        p.fail_reason = "response timeout";

        try {
            connections_[p.node_index].socket
                ->closeStream(p.stream);
        } catch (...) {
        }

        connections_[p.node_index].socket.reset();
        connections_[p.node_index].connected = false;

        p.finished = true;
    }

    std::vector<std::vector<uint8_t>> fetched;
    std::vector<size_t> fetch_indices;

    fetched.reserve(total_ok);
    fetch_indices.reserve(total_ok);

    for (auto& p : pending) {
        if (!p.ok || p.block.empty()) {
            continue;
        }

        fetch_indices.push_back(p.node_index);
        fetched.push_back(std::move(p.block));
    }

    if (fetched.size() < k_) {
        size_t not_found_count = 0;

        for (size_t i = 0; i < pending.size(); ++i) {
            if (pending[i].not_found) {
                confirmed_not_found[i] = true;
            }

            if (pending[i].ok) {
                confirmed_not_found[i] = false;
            }

            if (confirmed_not_found[i]) {
                ++not_found_count;
            }
        }

        /*
         * Retry only if reconstruction is still theoretically possible.
         *
         * If fewer than k blocks have been permanently confirmed as missing,
         * another attempt may succeed after a temporary transport failure or
         * after a node reconnects.
         */
        should_retry =
            (n - not_found_count >= k_);

        if (parity_diag_enabled()) {
            std::string diag;

            for (const auto& p : pending) {
                if (!diag.empty()) {
                    diag += ", ";
                }

                diag +=
                    "node" +
                    std::to_string(p.node_index) +
                    "=";

                if (p.ok) {
                    diag += "ok";
                } else if (p.not_found) {
                    diag += "not_found";
                } else if (!p.fail_reason.empty()) {
                    diag += p.fail_reason;
                } else {
                    diag += "no response";
                }
            }

            std::cerr
                << "[PARITY-DIAG] retrieve_stripe fail"
                << " gid=" << group_id
                << " stripe=" << stripe_index
                << " fetched=" << fetched.size()
                << "/" << n
                << " k=" << k_
                << " should_retry=" << should_retry
                << " diag=" << diag
                << "\n";
        }

        return false;
    }

    if (fetched.size() > k_) {
        fetched.resize(k_);
        fetch_indices.resize(k_);
    }

    /*
     * Fast path.
     *
     * If all original data shards 0..k-1 are available, reconstruction is
     * unnecessary and the original payload can be returned immediately.
     */
    bool all_data_blocks_present = true;

    for (size_t required = 0;
         required < k_;
         ++required)
    {
        if (std::find(
                fetch_indices.begin(),
                fetch_indices.end(),
                required
            ) == fetch_indices.end())
        {
            all_data_blocks_present = false;
            break;
        }
    }

    std::vector<std::vector<uint8_t>> data_blocks;

    if (all_data_blocks_present) {
        data_blocks.resize(k_);

        for (size_t i = 0;
             i < fetch_indices.size();
             ++i)
        {
            const size_t block_index =
                fetch_indices[i];

            if (block_index < k_) {
                data_blocks[block_index] =
                    std::move(fetched[i]);
            }
        }
    } else {
        /*
         * Reconstruction path.
         *
         * At least one data shard is missing, therefore Reed-Solomon decoding
         * is used to reconstruct the original data from the surviving shards.
         * All surviving shards must have identical sizes.
         */
        const size_t block_size =
            fetched.front().size();

        for (const auto& block : fetched) {
            if (block.size() != block_size) {
                throw std::runtime_error(
                    "retrieve_stripe_once: "
                    "surviving block size mismatch"
                );
            }
        }

        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::unique_lock<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 {
            // Write stripe 0 LAST, not first. Stripe 0 carries the 8-byte
            // length header that retrieve() trusts to know how many stripes
            // to read back — and store() never truncates stale trailing
            // stripes left over from a previous, longer store() at this
            // same group_id. Writing forward, a partial failure (throw
            // partway through) leaves an already-written new header
            // pointing past the failure point into those leftover old
            // stripes: retrieve() silently splices new+stale bytes into one
            // "valid" (length checks out) but corrupted result, no error
            // raised anywhere. Writing backwards means a partial failure
            // always leaves the last-committed header pointing at a fully
            // intact version — either the untouched old one (if we never
            // got to rewrite stripe 0) or the complete new one (if every
            // other stripe already succeeded) — never a torn mix.
            uint32_t stripes_done = 0;
            for (uint32_t stripe_idx = total_stripes; stripe_idx-- > 0; ) {
                size_t offset = static_cast<size_t>(stripe_idx) * STRIPE_SIZE;
                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);
                ++stripes_done;
                if (stripes_done % 100 == 0 || stripes_done == total_stripes) {
                    DBG_LOG("[PARITY]   stripe " << stripes_done << "/" << total_stripes << "\n");
                }
                // Release mutex between stripes so concurrent reads/retrieves
                // can proceed instead of waiting for the full store to finish
                if (stripe_idx > 0) { guard.unlock(); guard.lock(); }
            }
            DBG_LOG("[PARITY] store gid=" << group_id << " complete\n");
            return; // success
        } catch (const std::exception& e) {
            const std::string err = e.what();
            // Fail fast on response-timeout dominated failures: a full second
            // store pass doubles import stall time without improving outcome.
            if (err.find("response timeout") != std::string::npos) {
                throw;
            }
            if (attempt == 0) {
                DBG_LOG("[PARITY] store attempt 1 failed: " << err
                          << ", 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::chrono::steady_clock::time_point deadline) {
    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, deadline)) break;
        result.insert(result.end(), stripe.begin(), stripe.end());
        if (original_size > 0 && result.size() >= original_size) break;
        // Release mutex between stripes so other client operations
        // (concurrent retrieves, stores, status queries) can proceed.
        if (guard.owns_lock()) { guard.unlock(); guard.lock(); }
    }

    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::chrono::steady_clock::time_point deadline) {
    std::unique_lock<std::mutex> guard(mutex_);
    std::vector<uint8_t> result;
    result.reserve(STRIPE_SIZE * 4);

    bool definitively_not_found = false;
    uint64_t framed_size = 0;
    for (uint32_t s = 0; ; ++s) {
        std::vector<uint8_t> stripe;
        if (!retrieve_stripe(guard, group_id, s, stripe, deadline, &definitively_not_found)) break;
        result.insert(result.end(), stripe.begin(), stripe.end());
        if (framed_size == 0 && result.size() >= 8) {
            uint64_t payload_size = protocol::decode_u64(result.data());
            if (payload_size > std::numeric_limits<uint64_t>::max() - 8)
                throw std::runtime_error("retrieve: length header overflow");
            framed_size = payload_size + 8;
        }
        if (framed_size > 0 && result.size() >= framed_size) break;
        // Release mutex between stripes so other client operations
        // (concurrent retrieves, stores, status queries) can proceed.
        if (guard.owns_lock()) { guard.unlock(); guard.lock(); }
    }

    // No stripes found at all: only treat this as "key does not exist" when
    // explicit misses leave fewer than k possible shards for stripe 0. If
    // we gave up while k shards may still exist on unanswered nodes, that's a
    // connectivity failure masquerading as a missing key — surface it as an
    // error instead of silently handing the caller an empty result
    // (this was the root cause of "index fetch failed" giving zero
    // diagnostic detail: retrieve() looked identical whether "index" had
    // genuinely never been written, which never happens, or the cluster
    // just couldn't reach quorum on stripe 0 this attempt).
    if (result.empty()) {
        if (definitively_not_found) return {};
        throw std::runtime_error("retrieve: gid=" + std::to_string(group_id)
            + " stripe 0 unreachable — retries exhausted while k shards may "
              "still exist on unanswered nodes (connectivity issue, not a missing key)");
    }

    // 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::chrono::steady_clock::time_point deadline) {
    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, deadline)) {
            throw std::runtime_error("retrieve_range: stripe " + std::to_string(s) + " not found");
        }
        buf.insert(buf.end(), stripe.begin(), stripe.end());
        // Release mutex between stripes for concurrency
        if (s < last_stripe && guard.owns_lock()) { guard.unlock(); guard.lock(); }
    }

    // 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);

    std::vector<std::future<void>> pending;
    pending.reserve(nodes_.size());
    for (size_t i = 0; i < nodes_.size(); i++) {
        if (static_cast<int>(i) == local_node_index_ && local_store_) {
            local_store_->remove_group(group_id);
            continue;
        }
        pending.emplace_back(std::async(std::launch::async, [this, i, payload]() {
            try {
                send_recv(i, msg_type::DELETE_GROUP, payload);
            } catch (...) {
                // Best-effort deletion
            }
        }));
    }
    for (auto& future : pending) future.get();
}

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);

    struct node_list_result {
        size_t node_index;
        bool answered = false;
        std::vector<uint32_t> blocks;
    };
    std::vector<std::future<node_list_result>> pending;
    pending.reserve(nodes_.size());

    for (size_t i = 0; i < nodes_.size(); i++) {
        if (static_cast<int>(i) == local_node_index_ && local_store_) {
            result.emplace_back(i, local_store_->list_blocks(group_id));
            continue;
        }

        // Every task owns a different node connection. Run the probes in
        // parallel so one dead peer contributes one timeout to the scrub,
        // rather than serially delaying every other node.
        pending.emplace_back(std::async(std::launch::async,
            [this, i, payload]() -> node_list_result {
                node_list_result item{i};
                try {
                    auto response = send_recv(i, msg_type::LIST_BLOCKS, payload);
                    if (response.size() < protocol::HEADER_SIZE) return item;

                    msg_type type;
                    uint32_t resp_len;
                    if (!protocol::decode_header(response.data(), response.size(),
                                                 type, resp_len)
                        || type != msg_type::BLOCK_LIST || resp_len < 12) {
                        return item;
                    }

                    const uint8_t* rp = response.data() + protocol::HEADER_SIZE;
                    uint32_t count = protocol::decode_u32(rp + 8);
                    if (resp_len < 12 + count * 4) return item;
                    item.blocks.reserve(count);
                    for (uint32_t j = 0; j < count; ++j) {
                        item.blocks.push_back(protocol::decode_u32(rp + 12 + j * 4));
                    }
                    item.answered = true;
                } catch (...) {
                    // Omitted means unreachable. An explicit empty list is
                    // represented and can be diagnosed as missing blocks.
                }
                return item;
            }));
    }

    for (auto& future : pending) {
        auto item = future.get();
        if (item.answered) {
            result.emplace_back(item.node_index, std::move(item.blocks));
        }
    }
    std::sort(result.begin(), result.end(),
              [](const auto& a, const auto& b) { return a.first < b.first; });
    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_);
    return list_groups_on_node_impl(node_index, groups);
}

bool client::list_groups_on_node_impl(size_t node_index, std::vector<uint64_t>& groups) {
    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;
    }
}

std::vector<bool> client::list_groups_on_nodes(const std::vector<size_t>& node_indices) {
    std::lock_guard<std::mutex> guard(mutex_);

    std::vector<bool> reachable(node_indices.size(), false);
    std::vector<std::vector<uint64_t>> discard(node_indices.size());
    std::vector<std::future<void>> futures;
    futures.reserve(node_indices.size());

    for (size_t i = 0; i < node_indices.size(); ++i) {
        futures.push_back(std::async(std::launch::async, [this, i, &node_indices, &reachable, &discard]() {
            reachable[i] = list_groups_on_node_impl(node_indices[i], discard[i]);
        }));
    }
    for (auto& f : futures) f.get();

    return reachable;
}

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);
        }
    }

    // A node can be fully reachable and still be silently under-replicated
    // (e.g. missed some stripe writes) — group-ID presence checks elsewhere
    // don't catch this either, since a node only needs to hold one block of
    // a group to count as "present" for it. Surface it here by comparing
    // each node's self-reported block_count against its healthiest peer:
    // a node reporting fewer blocks than the max is behind and should not
    // read as OK.
    uint32_t max_blocks = 0;
    for (auto& [idx, info] : health.node_statuses) {
        if (info.status == node_status::OK)
            max_blocks = std::max(max_blocks, info.block_count);
    }
    for (auto& [idx, info] : health.node_statuses) {
        if (info.status == node_status::OK && info.block_count < max_blocks) {
            info.status = node_status::DEGRADED;
        }
    }

    for (auto& [idx, info] : health.node_statuses) {
        if (info.status == node_status::DEGRADED ||
            info.status == node_status::REBUILD) {
            health.overall = 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;
    std::vector<std::future<bool>> pending;
    pending.reserve(nodes_.size());

    for (size_t i = 0; i < nodes_.size(); i++) {
        if (static_cast<int>(i) == local_node_index_ && local_store_) {
            if (local_store_->vacuum()) ++ok_count;
            continue;
        }
        pending.emplace_back(std::async(std::launch::async, [this, i]() {
            try {
                auto resp = send_recv(i, msg_type::VACUUM, {});
                if (resp.size() >= protocol::HEADER_SIZE) {
                    msg_type rtype;
                    uint32_t rlen;
                    return protocol::decode_header(resp.data(), resp.size(), rtype, rlen)
                           && rtype == msg_type::OK;
                }
            } catch (...) {}
            return false;
        }));
    }
    for (auto& future : pending) {
        if (future.get()) ++ok_count;
    }
    return ok_count;
}

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

    // Discover the full universe of known group IDs. A group missing
    // entirely from a node won't appear in that node's list, so union
    // across every node's group list rather than trusting any single one.
    std::set<uint64_t> all_groups;
    {
        std::lock_guard<std::mutex> guard(mutex_);
        std::vector<std::vector<uint64_t>> groups_by_node(nodes_.size());
        std::vector<bool> answered(nodes_.size(), false);
        std::vector<std::future<void>> pending;
        pending.reserve(nodes_.size());
        for (size_t i = 0; i < nodes_.size(); ++i) {
            pending.emplace_back(std::async(std::launch::async,
                [this, i, &groups_by_node, &answered]() {
                    answered[i] = list_groups_on_node_impl(i, groups_by_node[i]);
                }));
        }
        for (auto& future : pending) future.get();

        for (size_t i = 0; i < nodes_.size(); ++i) {
            if (!answered[i]) continue;
            for (uint64_t gid : groups_by_node[i]) {
                if (gid != lock_gid) all_groups.insert(gid);
            }
        }
    }

    for (uint64_t gid : all_groups) {
        auto per_node = list(gid); // [(node_index, block_indices)], unreachable/errored nodes omitted

        // A transport timeout is not proof that a block is absent. Only an
        // explicit empty BLOCK_LIST from a reachable node is actionable.
        // Repairing an inconclusive group can otherwise churn healthy data
        // and keep a temporarily unavailable node permanently degraded.
        if (per_node.size() < n) {
            std::cerr << "[PARITY] check_groups: group " << gid
                      << " inconclusive — " << per_node.size() << "/" << n
                      << " nodes answered LIST_BLOCKS; skipping repair\n";
            continue;
        }

        // Each stripe contributes exactly one block per node (storage is
        // erasure-coded — node i only ever holds its own shard, never a
        // full copy). So a group is complete when every node is reachable,
        // no block landed on the wrong node (bidx % n == node_index), and
        // the total block count is a clean multiple of n (every stripe
        // fully represented). Comparing per-node block counts to the union
        // across all nodes is wrong here: each node legitimately holds only
        // 1/n of the union, so that comparison can never succeed for n>1.
        bool misplaced = false;
        size_t total_blocks = 0;
        for (auto& [idx, blocks] : per_node) {
            for (uint32_t bidx : blocks) {
                ++total_blocks;
                if (bidx % n != idx) misplaced = true;
            }
        }

        bool complete = !misplaced && total_blocks > 0 && total_blocks % n == 0;

        if (complete) ++result.already_ok;
        else result.incomplete.push_back(gid);
    }

    return result;
}

client::rebalance_result client::rebalance() {
    rebalance_result result;

    auto check = check_groups();
    result.already_ok = check.already_ok;

    if (check.incomplete.empty())
        return result;

    std::cerr << "[PARITY] rebalance: " << check.incomplete.size() << " group(s) to repair ("
              << result.already_ok << " already ok)" << std::endl;

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

            // 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: 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 the
            // 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.
            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()) {
        mark_connection_failed(node_index);
    }
}

// --- 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 {
        // Lock acquisition is an administrative probe, not a data request.
        // If the cluster cannot give a definitive answer quickly, skipping
        // this scrub is safer than waiting through the full reconstruction
        // retry budget and then risking two concurrent scrubbers.
        auto existing = retrieve(
            gid, std::chrono::steady_clock::now() + std::chrono::seconds(3));
        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 (...) {
        return false;
    }
    // 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
