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

#include <cstring>
#include <stdexcept>
#include <iostream>
#include <thread>
#include <chrono>

namespace paritypp {

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

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

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

client::~client() = default;

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

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

    // Wait for handshake using event-driven wait (epoll/kqueue/IOCP/poll)
    for (int i = 0; i < 50; i++) {
        if (nc.socket->hasPendingWrite()) {
            if (!wait_.waitWrite(*nc.socket, 100)) continue;
        } else {
            if (!wait_.waitRead(*nc.socket, 100)) continue;
        }
        nc.socket->pumpNetwork(MSG_DONTWAIT);
        if (nc.socket->getHandshakeDone()) {
            nc.connected = true;
            // Authenticate if credentials provided
            if (!creds_.client_name.empty()) {
                if (!authenticate_node(node_index)) {
                    nc.socket.reset();
                    nc.connected = false;
                    throw std::runtime_error("Authentication failed on node " +
                                              nodes_[node_index].address);
                }
            }
            return;
        }
    }
    nc.socket.reset();
    nc.connected = false;
    throw std::runtime_error("Failed to connect to node " + nodes_[node_index].address);
}

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

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

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

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

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

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

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

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

    return resp_type == msg_type::AUTH_OK;
}

void client::ensure_connected(size_t node_index) {
    if (connections_[node_index].connected && connections_[node_index].socket) {
        return;
    }
    // Reset broken state before reconnecting
    connections_[node_index].socket.reset();
    connections_[node_index].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) {
    ensure_connected(node_index);
    auto& conn = *connections_[node_index].socket;

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

    // Non-blocking send: wait for writable, then send
    for (int sw = 0; sw < 150; sw++) {
        if (wait_.waitWrite(conn, 100)) {
            conn.sendStreamData(stream, msg, true);
            break;
        }
        conn.pumpNetwork(MSG_DONTWAIT);
    }

    // Wait for response using event-driven waitRead (epoll/kqueue/IOCP/poll)
    std::vector<uint8_t> response;
    for (int attempt = 0; attempt < 150; attempt++) {
        if (!wait_.waitRead(conn, 100)) continue;
        conn.pumpNetwork(MSG_DONTWAIT);
        if (conn.hasStreamData(stream)) {
            // Read header
            uint8_t hdr[protocol::HEADER_SIZE];
            size_t r = conn.recvStreamData(stream, hdr, protocol::HEADER_SIZE);
            if (r < protocol::HEADER_SIZE) continue;

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

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

            size_t total_read = 0;
            int body_attempts = 0;
            while (total_read < resp_len && body_attempts++ < 150) {
                if (!wait_.waitRead(conn, 100)) continue;
                conn.pumpNetwork(MSG_DONTWAIT);
                size_t rd = conn.recvStreamData(stream,
                                                 response.data() + protocol::HEADER_SIZE + total_read,
                                                 resp_len - total_read);
                total_read += rd;
            }
            break;
        }
    }
    conn.closeStream(stream);
    // No response after all attempts — mark connection as broken so
    // subsequent calls (e.g. get_cluster_status) fail fast instead of
    // retrying the expensive connect + send loop on a dead node.
    if (response.empty()) {
        connections_[node_index].connected = false;
        connections_[node_index].socket.reset();
    }
    return response;
}

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) {
    // Pad to STRIPE_SIZE
    std::vector<uint8_t> padded(STRIPE_SIZE, 0);
    std::memcpy(padded.data(), stripe_data, std::min(stripe_len, STRIPE_SIZE));

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

    size_t n = k_ + m_;

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

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

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

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

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

            // Non-blocking send
            for (int sw = 0; sw < 150; sw++) {
                if (wait_.waitWrite(conn, 100)) {
                    conn.sendStreamData(stream, msg, true);
                    // Pump immediately to flush queued UDP datagrams
                    conn.pumpNetwork(MSG_DONTWAIT);
                    pending[i].sent = true;
                    break;
                }
                conn.pumpNetwork(MSG_DONTWAIT);
            }
            // Close leaked stream if send failed
            if (!pending[i].sent) {
                pending[i].fail_reason = "send timeout (waitWrite)";
                try { conn.closeStream(stream); } catch (...) {}
            }
        } catch (const netplus::NetException& e) {
            pending[i].fail_reason = std::string("net: ") + e.what();
            connections_[i].socket.reset();
            connections_[i].connected = false;
        } catch (const std::exception& e) {
            pending[i].fail_reason = std::string("exception: ") + e.what();
            connections_[i].socket.reset();
            connections_[i].connected = false;
        } catch (...) {
            pending[i].fail_reason = "unknown exception";
            connections_[i].socket.reset();
            connections_[i].connected = false;
        }
    };

    // Phase 1: store local blocks + open streams and send to remote nodes
    for (size_t i = 0; i < n; i++) {
        send_to_node(i);
    }

    // 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].socket.reset();
            connections_[i].connected = false;
            send_to_node(i);
        }
    }

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

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

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

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

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

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

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

    if (failed > m_) {
        throw std::runtime_error("store_stripe: too many node failures ("
                                 + std::to_string(failed) + " > " + std::to_string(m_)
                                 + "): " + diag);
    }
}

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

    for (size_t i = 0; i < n && fetched.size() < k_; i++) {
        try {
            std::vector<uint8_t> block;
            uint32_t bidx = stripe_index * static_cast<uint32_t>(k_ + m_) + static_cast<uint32_t>(i);
            if (fetch_block_from_node(i, group_id, bidx, block) && !block.empty()) {
                fetched.push_back(std::move(block));
                indices.push_back(i);
                any_found = true;
            }
        } catch (const std::exception& e) {
            connections_[i].socket.reset();
            connections_[i].connected = false;
        }
    }

    if (!any_found) return false;
    if (fetched.size() < k_) {
        throw std::runtime_error("retrieve_stripe: not enough blocks ("
                                 + std::to_string(fetched.size()) + " < " + std::to_string(k_) + ")");
    }

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

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

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

void client::store(uint64_t group_id, const uint8_t* data, size_t data_len) {
    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 stripe_idx = 0;
    size_t offset = 0;

    while (offset < framed_len) {
        std::vector<uint8_t> stripe(STRIPE_SIZE, 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, STRIPE_SIZE);
            std::memcpy(stripe.data(), header.data() + offset, hdr_copy);
            filled = hdr_copy;
        }
        // Fill from payload
        if (filled < STRIPE_SIZE && offset + filled >= 8) {
            size_t data_off = offset + filled - 8;
            size_t data_copy = std::min<size_t>(STRIPE_SIZE - 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(), STRIPE_SIZE);
        ++stripe_idx;
        offset += STRIPE_SIZE;
    }
}

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 (padded with zeros)
        if (!stripe_buf_.empty()) {
            stripe_buf_.resize(STRIPE_SIZE, 0);
            flush_stripe();
        }
    }

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

private:
    void flush_stripe() {
        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::vector<uint8_t> result;
    for (uint32_t s = 0; ; ++s) {
        std::vector<uint8_t> stripe;
        if (!retrieve_stripe(group_id, s, stripe)) break;
        result.insert(result.end(), stripe.begin(), stripe.end());
    }

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

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

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

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

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

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

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

size_t client::replicate_to_peers(uint64_t group_id, uint32_t block_index,
                                 const uint8_t* data, size_t len) {
    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) {
    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(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::list_groups_on_node(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;
    }
}

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

    for (size_t i = 0; i < nodes_.size(); i++) {
        status_info info;
        if (get_node_status(i, info)) {
            health.nodes_online++;
            health.node_statuses.emplace_back(i, info);
            // If any node reports degraded, note it
            if (info.status == node_status::DEGRADED ||
                info.status == node_status::REBUILD) {
                health.overall = info.status;
            }
        }
    }

    // Determine overall status based on reachability
    if (health.nodes_online == 0) {
        health.overall = node_status::DEGRADED;
    } else if (health.nodes_online < health.nodes_total) {
        // Some nodes unreachable — still degraded even if individual nodes report OK
        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) {
    local_node_index_ = static_cast<int>(node_index);
    local_store_ = std::move(store);
}

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

} // namespace paritypp
