#include "cluster.h"

#include <iostream>
#include <netplus/exception.h>

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

namespace mediadb {

Cluster* g_Cluster = nullptr;

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

    if (!cfg_.store_path.empty()) {
        store_ = std::make_shared<paritypp::file_block_store>(cfg_.store_path);
    } else {
        store_ = std::make_shared<paritypp::memory_block_store>();
    }

    auth_ = std::make_shared<paritypp::auth_store>();
    if (!cfg_.client_name.empty()) {
        auth_->add_client(cfg_.client_name, cfg_.client_key);
    }

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

    server_ = std::make_unique<paritypp::server>(scfg, store_, auth_);

    size_t peer_count = 0;
    if (!cfg_.peers.empty()) {
        std::vector<std::string> local_addrs;
        {
            char hostname[256] = {};
            gethostname(hostname, sizeof(hostname));
            local_addrs.emplace_back(hostname);
        }
#ifndef _WIN32
        {
            struct ifaddrs* ifaddr = nullptr;
            if (getifaddrs(&ifaddr) == 0) {
                for (struct ifaddrs* ifa = ifaddr; ifa; ifa = ifa->ifa_next) {
                    if (!ifa->ifa_addr) continue;
                    char buf[INET6_ADDRSTRLEN] = {};
                    if (ifa->ifa_addr->sa_family == AF_INET) {
                        inet_ntop(AF_INET,
                                  &reinterpret_cast<struct sockaddr_in*>(ifa->ifa_addr)->sin_addr,
                                  buf, sizeof(buf));
                        local_addrs.emplace_back(buf);
                    } else if (ifa->ifa_addr->sa_family == AF_INET6) {
                        inet_ntop(AF_INET6,
                                  &reinterpret_cast<struct sockaddr_in6*>(ifa->ifa_addr)->sin6_addr,
                                  buf, sizeof(buf));
                        local_addrs.emplace_back(buf);
                    }
                }
                freeifaddrs(ifaddr);
            }
        }
#endif
        std::vector<paritypp::client::node_info> nodes;
        for (size_t pi = 0; pi < cfg_.peers.size(); ++pi) {
            const auto& peer = cfg_.peers[pi];
            bool is_self = false;
            for (const auto& la : local_addrs) {
                if (peer.address == la && peer.port == cfg_.port) {
                    is_self = true;
                    break;
                }
            }
            if (is_self) {
                self_node_index_ = static_cast<int>(pi);
            }
            peer_to_client_index_[pi] = nodes.size();
            nodes.push_back({peer.address, peer.port});
        }
        peer_count = cfg_.peers.size();

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

        if (!nodes.empty()) {
            pclient_ = std::make_unique<paritypp::client>(
                cfg_.data_blocks, cfg_.parity_blocks, nodes, pcreds);
            if (self_node_index_ >= 0) {
                pclient_->set_local_node(
                    peer_to_client_index_[static_cast<size_t>(self_node_index_)], store_);
            }
            // Dedicated read client — separate QUIC connections, never blocked
            // by replicate/store holding pclient_->mutex_
            read_client_ = std::make_unique<paritypp::client>(
                cfg_.data_blocks, cfg_.parity_blocks, nodes, pcreds);
            if (self_node_index_ >= 0) {
                read_client_->set_local_node(
                    peer_to_client_index_[static_cast<size_t>(self_node_index_)], store_);
            }
            // Dedicated scrub/rebalance client — avoids blocking pclient_
            // during long scrub operations
            scrub_client_ = std::make_unique<paritypp::client>(
                cfg_.data_blocks, cfg_.parity_blocks, nodes, pcreds);
            if (self_node_index_ >= 0) {
                scrub_client_->set_local_node(
                    peer_to_client_index_[static_cast<size_t>(self_node_index_)], store_);
            }
            // Dedicated import client — avoids blocking pclient_ during
            // long streaming imports so normal writes remain responsive
            import_client_ = std::make_unique<paritypp::client>(
                cfg_.data_blocks, cfg_.parity_blocks, nodes, pcreds);
            if (self_node_index_ >= 0) {
                import_client_->set_local_node(
                    peer_to_client_index_[static_cast<size_t>(self_node_index_)], store_);
            }
            // NOTE: warmup() is deferred to start() so QUIC connections
            // are established in the child process (after fork).
        }
    }

    server_->set_peer_status_callback([peer_count]() -> std::pair<uint32_t, uint32_t> {
        return {0, static_cast<uint32_t>(peer_count)};
    });
}

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

    // Compact the local store on startup to reclaim tombstoned space
    auto* fbs = dynamic_cast<paritypp::file_block_store*>(store_.get());
    if (fbs) {
        DBG_LOG("[CLUSTER] startup vacuum ... ");
        if (fbs->vacuum())
            DBG_LOG("done\n");
        else
            DBG_LOG("skipped (below threshold)\n");
    }

    // Establish QUIC connections now (post-fork) so they survive in the
    // child process.  Doing this in init() (pre-fork) leaves the
    // connections dead after fork, causing empty retrieve results.
    if (pclient_) pclient_->warmup();
    if (read_client_) read_client_->warmup();
    if (scrub_client_) scrub_client_->warmup();
    if (import_client_) import_client_->warmup();

    running_ = true;
    server_thread_ = std::thread(&Cluster::server_loop, this);
    health_thread_ = std::thread(&Cluster::health_loop, this);
}

void Cluster::stop() {
    if (!running_) return;
    running_ = false;
    if (server_) server_->stop();
    if (health_thread_.joinable()) health_thread_.join();
    if (server_thread_.joinable()) server_thread_.join();
}

size_t Cluster::waitForPeers(int timeout_s) {
    if (!pclient_) return 0;
    size_t k = cfg_.data_blocks;
    size_t n = cfg_.data_blocks + cfg_.parity_blocks;
    size_t best_online = 0;
    for (int elapsed = 0; elapsed < timeout_s; ++elapsed) {
        try {
            auto health = pclient_->get_cluster_status();
            if (health.nodes_online > best_online)
                best_online = health.nodes_online;
            std::cerr << "[CLUSTER] " << health.nodes_online << "/"
                      << health.nodes_total << " nodes online ("
                      << elapsed << "s/" << timeout_s << "s)\n";
            if (health.nodes_online >= n) {
                std::cerr << "[CLUSTER] all nodes online\n";
                degraded_ = false;
                critical_ = false;
                return health.nodes_online;
            }
            if (health.nodes_online >= k) {
                std::cerr << "[CLUSTER] DEGRADED — " << health.nodes_online
                          << "/" << n << " nodes online, continuing\n";
                degraded_ = true;
                critical_ = false;
                return health.nodes_online;
            }
        } catch (const std::exception& e) {
            std::cerr << "[CLUSTER] peer probe failed: " << e.what() << "\n";
        }
        std::this_thread::sleep_for(std::chrono::seconds(1));
    }
    std::cerr << "[CLUSTER] FAILED — only " << best_online << "/" << n
              << " nodes reachable after " << timeout_s << "s (need " << k
              << ")\n";
    return best_online;
}

std::vector<Cluster::NodeGroupInfo> Cluster::list_peer_groups() {
    std::vector<NodeGroupInfo> result;
    if (!pclient_) return result;
    size_t idx = 0;
    for (const auto& peer : cfg_.peers) {
        NodeGroupInfo ngi;
        ngi.index = idx;
        ngi.address = peer.address;
        ngi.port = peer.port;
        ngi.online = false;
        auto it = peer_to_client_index_.find(idx);
        if (it != peer_to_client_index_.end() &&
            pclient_->list_groups_on_node(it->second, ngi.groups)) {
            ngi.online = true;
        }
        result.push_back(std::move(ngi));
        ++idx;
    }
    return result;
}

bool Cluster::replicate(const std::string& key, const uint8_t* data, size_t len) {
    if (!pclient_) return false;
    uint64_t gid = cluster_group_id(key);
    try {
        DBG_LOG("[CLUSTER] replicate key=" << key << " gid=" << gid
                  << " size=" << len << "\n");
        pclient_->store(gid, data, len);
        DBG_LOG("[CLUSTER] replicate key=" << key << " store done\n");
    } catch (const netplus::NetException& e) {
        DBG_LOG("[CLUSTER] replicate NetException key=" << key << ": " << e.what() << "\n");
        return false;
    } catch (const std::exception& e) {
        DBG_LOG("[CLUSTER] replicate exception key=" << key << ": " << e.what() << "\n");
        return false;
    }
    return true;
}

bool Cluster::replicate_import(const std::string& key, const uint8_t* data, size_t len) {
    if (!import_client_) return false;
    uint64_t gid = cluster_group_id(key);
    try {
        DBG_LOG("[CLUSTER] replicate_import key=" << key << " gid=" << gid
                  << " size=" << len << "\n");
        import_client_->store(gid, data, len);
        DBG_LOG("[CLUSTER] replicate_import key=" << key << " store done\n");
    } catch (const netplus::NetException& e) {
        DBG_LOG("[CLUSTER] replicate_import NetException key=" << key << ": " << e.what() << "\n");
        return false;
    } catch (const std::exception& e) {
        DBG_LOG("[CLUSTER] replicate_import exception key=" << key << ": " << e.what() << "\n");
        return false;
    }
    return true;
}

std::unique_ptr<paritypp::store_session> Cluster::begin_replicate(const std::string& key, size_t total_size) {
    if (!pclient_) return nullptr;
    uint64_t gid = cluster_group_id(key);
    try {
        return pclient_->begin_store(gid, total_size);
    } catch (const std::exception& e) {
        return nullptr;
    }
}

std::unique_ptr<paritypp::store_session> Cluster::begin_replicate_import(const std::string& key, size_t total_size) {
    if (!import_client_) return nullptr;
    uint64_t gid = cluster_group_id(key);
    try {
        return import_client_->begin_store(gid, total_size);
    } catch (const std::exception& e) {
        return nullptr;
    }
}

std::unique_ptr<paritypp::store_session> Cluster::reserve_block(const std::string& key,
                                                                 size_t total_size,
                                                                 int max_retries) {
    if (!import_client_) return nullptr;
    uint64_t gid = cluster_group_id(key);
    for (int attempt = 0; attempt < max_retries; ++attempt) {
        if (attempt > 0) {
            import_client_->warmup();
            std::this_thread::sleep_for(
                std::chrono::milliseconds(200 * (1 << (attempt - 1))));
        }
        try {
            auto session = import_client_->begin_store(gid, total_size);
            if (session) return session;
        } catch (const std::exception& e) {
            DBG_LOG("[CLUSTER] reserve_block attempt " << attempt
                      << " failed key=" << key << ": " << e.what() << "\n");
        }
    }
    return nullptr;
}

void Cluster::warmup_import_client() {
    if (import_client_) import_client_->warmup();
}

void Cluster::warmup_read_clients() {
    if (read_client_) read_client_->warmup();
    if (pclient_) pclient_->warmup();
}

bool Cluster::fetch(const std::string& key, std::vector<uint8_t>& out) {
    uint64_t gid = cluster_group_id(key);

    // Try read_client_ first (dedicated read connections)
    if (read_client_) {
        try {
            out = read_client_->retrieve(gid);
            if (!out.empty()) return true;
        } catch (const netplus::NetException& e) {
            std::cerr << "[CLUSTER] fetch read_client NetException key=" << key
                      << " gid=" << gid << ": " << e.what() << "\n";
            try {
                read_client_->warmup();
                out = read_client_->retrieve(gid);
                if (!out.empty()) return true;
            } catch (const std::exception& e2) {
                std::cerr << "[CLUSTER] fetch read_client retry failed key=" << key
                          << " gid=" << gid << ": " << e2.what() << "\n";
            } catch (...) {
                std::cerr << "[CLUSTER] fetch read_client retry failed key=" << key
                          << " gid=" << gid << ": unknown exception\n";
            }
        } catch (const std::exception& e) {
            // Data-level error (e.g. "stripe not found") — warmup won't help.
            std::cerr << "[CLUSTER] fetch read_client exception key=" << key
                      << " gid=" << gid << ": " << e.what() << "\n";
        }
    }

    // Fallback to pclient_ (write client) — may succeed if read_client_
    // connections are stale (e.g. after fork or network glitch)
    if (pclient_) {
        try {
            out = pclient_->retrieve(gid);
            if (!out.empty()) {
                std::cerr << "[CLUSTER] fetch key=" << key
                          << " succeeded via pclient_ fallback\n";
                return true;
            }
        } catch (const netplus::NetException& e) {
            std::cerr << "[CLUSTER] fetch pclient NetException key=" << key
                      << " gid=" << gid << ": " << e.what() << "\n";
            try {
                pclient_->warmup();
                out = pclient_->retrieve(gid);
                if (!out.empty()) return true;
            } catch (const std::exception& e2) {
                std::cerr << "[CLUSTER] fetch pclient retry failed key=" << key
                          << " gid=" << gid << ": " << e2.what() << "\n";
            } catch (...) {
                std::cerr << "[CLUSTER] fetch pclient retry failed key=" << key
                          << " gid=" << gid << ": unknown exception\n";
            }
        } catch (const std::exception& e) {
            // Data-level error — skip warmup retry.
            std::cerr << "[CLUSTER] fetch pclient exception key=" << key
                      << " gid=" << gid << ": " << e.what() << "\n";
        }
    }

    return false;
}

bool Cluster::fetch_range(const std::string& key, uint64_t offset, uint64_t length,
                           std::vector<uint8_t>& out) {
    uint64_t gid = cluster_group_id(key);

    if (read_client_) {
        try {
            out = read_client_->retrieve_range(gid, offset, length);
            if (!out.empty()) return true;
        } catch (const netplus::NetException& e) {
            std::cerr << "[CLUSTER] fetch_range read_client NetException key=" << key
                      << " gid=" << gid << ": " << e.what() << "\n";
            try {
                read_client_->warmup();
                out = read_client_->retrieve_range(gid, offset, length);
                if (!out.empty()) return true;
            } catch (...) {}
        } catch (const std::exception& e) {
            // Data-level error (e.g. "stripe not found") — warmup won't help,
            // skip retry to avoid blocking 30+ s per dead node.
            std::cerr << "[CLUSTER] fetch_range read_client exception key=" << key
                      << " gid=" << gid << ": " << e.what() << "\n";
        }
    }

    if (pclient_) {
        try {
            out = pclient_->retrieve_range(gid, offset, length);
            if (!out.empty()) return true;
        } catch (const netplus::NetException& e) {
            std::cerr << "[CLUSTER] fetch_range pclient NetException key=" << key
                      << " gid=" << gid << ": " << e.what() << "\n";
            try {
                pclient_->warmup();
                out = pclient_->retrieve_range(gid, offset, length);
                if (!out.empty()) return true;
            } catch (...) {}
        } catch (const std::exception& e) {
            // Data-level error — skip warmup retry (same reason as above).
            std::cerr << "[CLUSTER] fetch_range pclient exception key=" << key
                      << " gid=" << gid << ": " << e.what() << "\n";
        }
    }

    return false;
}

bool Cluster::fetch_from_peers(const std::string& key, std::vector<uint8_t>& out) {
    auto& cli = read_client_ ? read_client_ : pclient_;
    if (!cli) return false;
    uint64_t gid = cluster_group_id(key);
    try {
        out = cli->retrieve(gid);
        return !out.empty();
    } catch (const netplus::NetException& e) {
        DBG_LOG("[CLUSTER] fetch_from_peers NetException key=" << key << ": " << e.what() << "\n");
    } catch (const std::exception& e) {
        DBG_LOG("[CLUSTER] fetch_from_peers exception key=" << key << ": " << e.what() << "\n");
    }
    return false;
}

void Cluster::remove(const std::string& key) {
    if (!pclient_) return;
    uint64_t gid = cluster_group_id(key);
    try {
        pclient_->remove(gid);
    } catch (...) {}
}

bool Cluster::vacuum_local_store() {
    auto* fbs = dynamic_cast<paritypp::file_block_store*>(store_.get());
    if (!fbs) return false;
    return fbs->vacuum();
}

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

    // Local mutex: prevent concurrent scrub/rebalance/vacuum on this node
    std::unique_lock<std::mutex> scrub_lk(scrub_mutex_, std::try_to_lock);
    if (!scrub_lk.owns_lock()) {
        std::cerr << "[VACUUM] skipped — scrub/rebalance is running\n";
        return 0;
    }

    // Cluster-wide lock: prevent concurrent vacuum/scrub across nodes
    if (!scrub_client_->try_scrub_lock()) {
        std::cerr << "[VACUUM] skipped — another node holds the scrub lock\n";
        return 0;
    }

    size_t result = scrub_client_->vacuum_all_nodes();

    scrub_client_->release_scrub_lock();
    return result;
}

void Cluster::server_loop() {
    try {
        server_->run();
    } catch (const netplus::NetException& e) {
        DBG_LOG("[CLUSTER] server_loop NetException: " << e.what() << "\n");
    } catch (const std::exception& e) {
        DBG_LOG("[CLUSTER] server_loop exception: " << e.what() << "\n");
    }
    running_ = false;
}

void Cluster::health_loop() {
    size_t k = cfg_.data_blocks;
    size_t n = cfg_.data_blocks + cfg_.parity_blocks;
    int healthy_cycles = 0;
    static constexpr int REBALANCE_INTERVAL = 10; // rebalance every ~5min (10 × 30s)
    while (running_) {
        if (pclient_) {
            uint32_t online = 0;
            uint32_t total = static_cast<uint32_t>(cfg_.peers.size());
            for (size_t pi = 0; pi < cfg_.peers.size(); ++pi) {
                auto it = peer_to_client_index_.find(pi);
                if (it == peer_to_client_index_.end()) continue;
                if (static_cast<int>(pi) == self_node_index_) {
                    ++online;
                    continue;
                }
                std::vector<uint64_t> dummy;
                if (pclient_->list_groups_on_node(it->second, dummy)) {
                    ++online;
                } else {
                    // Retry once after short delay — peers may be busy
                    // during import or scrub and miss the first probe
                    std::this_thread::sleep_for(std::chrono::milliseconds(500));
                    dummy.clear();
                    if (pclient_->list_groups_on_node(it->second, dummy))
                        ++online;
                }
            }
            server_->set_peer_status_callback([online, total]() -> std::pair<uint32_t, uint32_t> {
                return {online, total};
            });

            // Verify local store is writable
            bool store_ok = true;
            if (store_) {
                try {
                    const uint64_t sentinel_gid = 0;
                    const uint32_t sentinel_idx = UINT32_MAX;
                    uint8_t probe = 0x01;
                    store_ok = store_->store(sentinel_gid, sentinel_idx, &probe, 1);
                    if (store_ok)
                        store_->remove_group(sentinel_gid);
                } catch (...) {
                    store_ok = false;
                }
                if (!store_ok && !critical_) {
                    std::cerr << "[CLUSTER] CRITICAL — local store not writable\n";
                }
            }

            bool was_degraded = degraded_;
            if (!store_ok) {
                degraded_ = true;
                critical_ = true;
            } else if (online >= n) {
                if (degraded_) {
                    std::cerr << "[CLUSTER] recovered — all " << online << "/"
                              << n << " nodes online, starting scrub\n";
                    std::thread([this](){ try { scrub(); } catch (...) {} }).detach();
                    healthy_cycles = 0;
                }
                degraded_ = false;
                critical_ = false;

                // Periodic rebalance while healthy
                if (++healthy_cycles >= REBALANCE_INTERVAL) {
                    healthy_cycles = 0;
                    std::thread([this](){
                        std::unique_lock<std::mutex> lk(scrub_mutex_, std::try_to_lock);
                        if (!lk.owns_lock()) return;
                        if (!scrub_client_ || !scrub_client_->try_scrub_lock()) return;
                        scrub_running_ = true;
                        if (server_) server_->set_rebuilding(true);
                        try {
                            auto rb = scrub_client_->rebalance();
                            if (rb.rebalanced > 0) {
                                std::cerr << "[CLUSTER] periodic rebalance moved "
                                          << rb.rebalanced << " group(s)\n";
                                scrub_client_->vacuum_all_nodes();
                            }
                        } catch (...) {}
                        scrub_running_ = false;
                        if (server_) server_->set_rebuilding(false);
                        scrub_client_->release_scrub_lock();
                    }).detach();
                }
            } else if (online >= k) {
                // During active imports, peers may be busy handling
                // replication traffic and miss health probes.  Suppress
                // the DEGRADED transition so it doesn't trigger a
                // recovery-scrub when the import finishes normally.
                if (!import_active_.load()) {
                    if (!degraded_) {
                        std::cerr << "[CLUSTER] DEGRADED — " << online << "/"
                                  << n << " nodes online\n";
                    }
                    degraded_ = true;
                } else {
                    if (!degraded_) {
                        std::cerr << "[CLUSTER] probe missed " << (n - online) << " node(s)"
                                  << " during import — suppressing DEGRADED\n";
                    }
                }
                critical_ = false;
                healthy_cycles = 0;
            } else {
                if (!critical_) {
                    std::cerr << "[CLUSTER] CRITICAL — only " << online << "/"
                              << n << " nodes online (need " << k << ")\n";
                }
                degraded_ = true;
                critical_ = true;
            }
        }
        // Probe faster when degraded/critical (5s) vs normal (30s)
        int interval = (degraded_ || critical_) ? 5 : 30;
        for (int i = 0; i < interval && running_; ++i)
            std::this_thread::sleep_for(std::chrono::seconds(1));
    }
}

Cluster::RebalanceResult Cluster::rebalance() {
    RebalanceResult result;
    auto& cli = scrub_client_ ? scrub_client_ : pclient_;
    if (!cli) return result;

    std::unique_lock<std::mutex> scrub_lk(scrub_mutex_, std::try_to_lock);
    if (!scrub_lk.owns_lock()) {
        std::cerr << "[CLUSTER] rebalance skipped — another scrub/rebalance is running\n";
        return result;
    }

    scrub_running_ = true;
    if (server_) server_->set_rebuilding(true);
    try {
        cli->warmup();
        auto rb = cli->rebalance();
        result.rebalanced = rb.rebalanced;
        result.already_ok = rb.already_ok;
        result.failed = rb.failed;
        if (rb.rebalanced > 0) {
            std::cerr << "[CLUSTER] rebalance moved " << rb.rebalanced << " group(s)\n";
            cli->vacuum_all_nodes();
        }
    } catch (const std::exception& e) {
        std::cerr << "[CLUSTER] rebalance exception: " << e.what() << "\n";
    }
    scrub_running_ = false;
    if (server_) server_->set_rebuilding(false);
    return result;
}

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

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

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

    size_t n = cfg_.data_blocks + cfg_.parity_blocks;

    // 1. Collect all group IDs from every node
    uint64_t lock_gid = paritypp::client::scrub_lock_group_id();
    std::map<uint64_t, std::set<size_t>> group_nodes;
    for (size_t pi = 0; pi < cfg_.peers.size(); ++pi) {
        auto it = peer_to_client_index_.find(pi);
        if (it == peer_to_client_index_.end()) continue;
        std::vector<uint64_t> node_groups;
        if (scrub_client_->list_groups_on_node(it->second, node_groups)) {
            for (uint64_t gid : node_groups) {
                if (gid == lock_gid) continue;  // skip lock group
                group_nodes[gid].insert(pi);
            }
        }
    }

    std::cerr << "[SCRUB] found " << group_nodes.size()
              << " groups across " << n << " nodes\n";

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

        if (present_nodes.size() >= n) {
            continue;
        }

        std::cerr << "[SCRUB] group " << gid << " present on "
                  << present_nodes.size() << "/" << n
                  << " nodes — repairing\n";

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

            // remove + store ensures correct block placement on all nodes
            scrub_client_->remove(gid);
            scrub_client_->store(gid, data.data(), data.size());
            result.groups_repaired++;
            std::cerr << "[SCRUB] group " << gid << " — repaired\n";
        } catch (const std::exception& e) {
            std::cerr << "[SCRUB] group " << gid
                      << " — repair failed: " << e.what() << "\n";
            result.groups_failed++;
        }
    }

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

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

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

} // namespace mediadb
