#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_);
            }
            // Pool of dedicated read clients — separate QUIC connections each,
            // concurrent preview/download requests don't serialize.
            read_pool_.resize(READ_POOL_SIZE);
            for (size_t ri = 0; ri < READ_POOL_SIZE; ++ri) {
                read_pool_[ri] = std::make_unique<paritypp::client>(
                    cfg_.data_blocks, cfg_.parity_blocks, nodes, pcreds);
                if (self_node_index_ >= 0) {
                    read_pool_[ri]->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_);
            }
            // Cache for reset_import_client()
            client_nodes_ = nodes;
            client_creds_ = pcreds;
            // 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)};
    });
}

paritypp::client* Cluster::next_read_client() {
    if (read_pool_.empty()) return nullptr;
    size_t idx = read_pool_rr_.fetch_add(1, std::memory_order_relaxed) % read_pool_.size();
    return read_pool_[idx].get();
}

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();
    for (auto& rc : read_pool_) if (rc) rc->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::reset_import_client() {
    if (client_nodes_.empty()) return;
    std::lock_guard<std::mutex> lk(import_client_mutex_);
    std::cerr << "[CLUSTER] reset_import_client: recreating import_client_ to clear error state\n";
    import_client_ = std::make_unique<paritypp::client>(
        cfg_.data_blocks, cfg_.parity_blocks, client_nodes_, client_creds_);
    if (self_node_index_ >= 0) {
        import_client_->set_local_node(
            peer_to_client_index_[static_cast<size_t>(self_node_index_)], store_);
    }
    import_client_->warmup();
}

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

void Cluster::warmup_read_clients() {
    for (auto& rc : read_pool_) if (rc) rc->warmup();
    if (pclient_) pclient_->warmup();
}

bool Cluster::fetch(const std::string& key, std::vector<uint8_t>& out) {
    uint64_t gid = cluster_group_id(key);
    // Overall deadline so HTTP requests don't hang.
    auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(30);
    auto timed_out = [&]() {
        return std::chrono::steady_clock::now() >= deadline;
    };

    // Try a read pool client (round-robin across independent connections)
    auto* rcli = next_read_client();
    if (rcli) {
        try {
            out = rcli->retrieve(gid);
            if (!out.empty()) return true;
        } catch (const netplus::NetException& e) {
            std::cerr << "[CLUSTER] fetch read_pool NetException key=" << key
                      << " gid=" << gid << ": " << e.what() << "\n";
            if (!timed_out()) {
                try {
                    rcli->warmup();
                    out = rcli->retrieve(gid);
                    if (!out.empty()) return true;
                } catch (const std::exception& e2) {
                    std::cerr << "[CLUSTER] fetch read_pool retry failed key=" << key
                              << " gid=" << gid << ": " << e2.what() << "\n";
                } catch (...) {
                    std::cerr << "[CLUSTER] fetch read_pool 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_pool exception key=" << key
                      << " gid=" << gid << ": " << e.what() << "\n";
        }
    }

    // Fallback to pclient_ (write client) — may succeed if read pool
    // connections are stale (e.g. after fork or network glitch)
    if (pclient_ && !timed_out()) {
        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";
            if (!timed_out()) {
                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) {
    // Overall deadline: fail fast so HTTP requests don't hang.
    auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(30);
    uint64_t gid = cluster_group_id(key);

    auto timed_out = [&]() {
        return std::chrono::steady_clock::now() >= deadline;
    };

    auto* rcli = next_read_client();
    if (rcli) {
        try {
            out = rcli->retrieve_range(gid, offset, length);
            if (!out.empty()) return true;
        } catch (const netplus::NetException& e) {
            std::cerr << "[CLUSTER] fetch_range read_pool NetException key=" << key
                      << " gid=" << gid << ": " << e.what() << "\n";
            if (!timed_out()) {
                try {
                    rcli->warmup();
                    out = rcli->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_pool exception key=" << key
                      << " gid=" << gid << ": " << e.what() << "\n";
        }
    }

    if (pclient_ && !timed_out()) {
        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";
            if (!timed_out()) {
                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* rcli = next_read_client();
    paritypp::client* cli = rcli ? rcli : pclient_.get();
    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 without polluting blocks.bin.
            // Previous approach wrote a sentinel record every health cycle
            // which accumulated thousands of dead records in the data file.
            bool store_ok = true;
            if (!cfg_.store_path.empty()) {
                try {
                    std::string probe_path = cfg_.store_path + "/.health_probe";
#ifdef _WIN32
                    int probe_fd = ::_open(probe_path.c_str(),
                                           _O_WRONLY | _O_CREAT | _O_TRUNC,
                                           _S_IREAD | _S_IWRITE);
                    if (probe_fd >= 0) {
                        store_ok = (::_write(probe_fd, "OK", 2) == 2);
                        ::_close(probe_fd);
                        ::_unlink(probe_path.c_str());
                    } else {
                        store_ok = false;
                    }
#else
                    int probe_fd = ::open(probe_path.c_str(),
                                          O_WRONLY | O_CREAT | O_TRUNC, 0644);
                    if (probe_fd >= 0) {
                        store_ok = (::write(probe_fd, "OK", 2) == 2);
                        ::close(probe_fd);
                        ::unlink(probe_path.c_str());
                    } else {
                        store_ok = false;
                    }
#endif
                } catch (...) {
                    store_ok = false;
                }
                if (!store_ok && !critical_) {
                    std::cerr << "[CLUSTER] CRITICAL — local store not writable\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;
            }

            // SAFETY: store FIRST so correct blocks exist on correct nodes.
            // If this throws, original data remains on old nodes.
            scrub_client_->store(gid, data.data(), data.size());

            // Data is now safely placed. Remove all copies and re-store
            // for clean distribution without leftover misplaced blocks.
            try {
                scrub_client_->remove(gid);
                scrub_client_->store(gid, data.data(), data.size());
            } catch (const std::exception& ce) {
                // Cleanup pass failed — data is still safe from the
                // first store(). Misplaced blocks may remain.
                std::cerr << "[SCRUB] group " << gid
                          << " — cleanup pass failed: " << ce.what()
                          << " (data is safe)\n";
            }
            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
