#pragma once

#include <atomic>
#include <chrono>
#include <cstdint>
#include <iostream>
#include <map>
#include <memory>
#include <mutex>
#include <set>
#include <string>
#include <thread>
#include <vector>

#ifndef _WIN32
#include <unistd.h>
#include <fcntl.h>
#include <ifaddrs.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#endif

#include <crypto/sha.h>
#include <paritypp/server.h>
#include <paritypp/client.h>
#include <paritypp/auth.h>
#include <paritypp/protocol.h>

namespace mediadb {

struct ClusterNode {
    std::string address;
    int         port;
};

struct ClusterConfig {
    bool                     enabled = false;
    std::vector<ClusterNode> peers;
    std::string              bind_address = "0.0.0.0";
    int                      port = 4433;
    std::string              cert_file;
    std::string              key_file;
    std::string              client_name;
    std::string              client_key;
    std::string              store_path;
    size_t                   data_blocks  = 2;
    size_t                   parity_blocks = 1;
};

// Deterministic group_id from a string key (e.g. "index" or store_id)
static inline uint64_t cluster_group_id(const std::string& key) {
    std::vector<uint8_t> input(key.begin(), key.end());
    auto hash = netplus::sha256_hash(input);
    uint64_t id = 0;
    for (int i = 0; i < 8; i++)
        id = (id << 8) | hash[i];
    return id;
}

class Cluster {
public:
    Cluster() = default;
    ~Cluster() { stop(); }

    void init(const ClusterConfig& cfg);
    void start();
    void stop();

    // Wait until enough peers are online to reconstruct data.
    // Returns the number of online nodes. Blocks up to timeout_s seconds.
    size_t waitForPeers(int timeout_s = 120);

    bool isRunning() const { return running_; }

    paritypp::block_store& getStore() { return *store_; }
    std::unique_ptr<paritypp::client>& getClient() { return pclient_; }
    const ClusterConfig& getConfig() const { return cfg_; }
    int getSelfNodeIndex() const { return self_node_index_; }
    size_t getTotalPeers() const { return cfg_.peers.size(); }
    const std::unordered_map<size_t, size_t>& getPeerToClientIndex() const { return peer_to_client_index_; }

    // Register an additional block store on the parity server (multi-store)
    void addStore(uint8_t store_id, std::shared_ptr<paritypp::block_store> store) {
        if (server_) server_->add_store(store_id, std::move(store));
    }

    struct NodeGroupInfo {
        size_t index;
        std::string address;
        int port;
        bool online;
        std::vector<uint64_t> groups;
    };
    std::vector<NodeGroupInfo> list_peer_groups();

    bool replicate(const std::string& key, const uint8_t* data, size_t len);
    // One-shot store on the dedicated import client.
    bool replicate_import(const std::string& key, const uint8_t* data, size_t len);
    std::unique_ptr<paritypp::store_session> begin_replicate(const std::string& key, size_t total_size);
    // Streaming store session on the dedicated import client — does not block
    // pclient_ so normal writes (add_media, delete, metadata) remain responsive.
    std::unique_ptr<paritypp::store_session> begin_replicate_import(const std::string& key, size_t total_size);

    // Reserve a block for streaming writes on the import client.
    // Warms up the connection, then returns a store_session.
    // On failure retries up to max_retries times.
    std::unique_ptr<paritypp::store_session> reserve_block(const std::string& key,
                                                            size_t total_size,
                                                            int max_retries = 3);

    // Re-establish QUIC connections on the import client.
    void warmup_import_client();
    // Re-establish QUIC connections on read pool and pclient_ (used before sync).
    void warmup_read_clients();
    bool fetch(const std::string& key, std::vector<uint8_t>& out);
    bool fetch_range(const std::string& key, uint64_t offset, uint64_t length,
                      std::vector<uint8_t>& out);
    bool fetch_from_peers(const std::string& key, std::vector<uint8_t>& out);
    void remove(const std::string& key);

    // Compact the local block store file, reclaiming space from tombstoned records.
    bool vacuum_local_store();

    // Compact block stores on all cluster nodes.
    size_t vacuum_all_nodes();

    bool isDegraded() const { return degraded_; }
    bool isCritical() const { return critical_; }

    // Scrub: verify and repair all data groups across the cluster.
    // For each group, checks that every node has its blocks. Missing blocks
    // are fetched from peers and re-replicated.
    struct ScrubResult {
        size_t groups_checked  = 0;
        size_t groups_repaired = 0;
        size_t groups_failed   = 0;
    };
    ScrubResult scrub();

    struct RebalanceResult {
        size_t rebalanced = 0;
        size_t already_ok = 0;
        size_t failed     = 0;
    };

    // Rebalance: re-distribute under-replicated groups across all nodes.
    RebalanceResult rebalance();

    // Last scrub/rebalance results (accessible for status display)
    ScrubResult lastScrubResult() const {
        std::lock_guard<std::mutex> lk(result_mutex_);
        return last_scrub_;
    }
    std::string lastScrubTime() const {
        std::lock_guard<std::mutex> lk(result_mutex_);
        return last_scrub_time_;
    }
    void setLastScrubResult(const ScrubResult& r, const std::string& time) {
        std::lock_guard<std::mutex> lk(result_mutex_);
        last_scrub_ = r;
        last_scrub_time_ = time;
    }

    RebalanceResult lastRebalanceResult() const {
        std::lock_guard<std::mutex> lk(result_mutex_);
        return last_rebalance_;
    }
    std::string lastRebalanceTime() const {
        std::lock_guard<std::mutex> lk(result_mutex_);
        return last_rebalance_time_;
    }
    void setLastRebalanceResult(const RebalanceResult& r, const std::string& time) {
        std::lock_guard<std::mutex> lk(result_mutex_);
        last_rebalance_ = r;
        last_rebalance_time_ = time;
    }

    // Try to acquire the scrub/rebalance lock (non-blocking).
    std::unique_lock<std::mutex> tryScrubLock() {
        return std::unique_lock<std::mutex>(scrub_mutex_, std::try_to_lock);
    }

    bool isScrubRunning() const { return scrub_running_.load(); }

    // External flag: set by App when an import is in progress, so the
    // health loop can tolerate transient probe failures (peers busy
    // handling import replication traffic).
    void setImportRunning(bool v) { import_active_.store(v); }
    bool isImportRunning() const { return import_active_.load(); }

private:
    ClusterConfig cfg_;
    std::atomic<bool> running_{false};
    std::atomic<bool> degraded_{false};
    std::atomic<bool> critical_{false};
    std::atomic<bool> import_active_{false};
    std::thread server_thread_;
    std::thread health_thread_;
    std::mutex scrub_mutex_;
    std::atomic<bool> scrub_running_{false};
    mutable std::mutex result_mutex_;
    ScrubResult last_scrub_;
    std::string last_scrub_time_;
    RebalanceResult last_rebalance_;
    std::string last_rebalance_time_;

    std::shared_ptr<paritypp::block_store> store_;
    std::shared_ptr<paritypp::auth_store> auth_;
    std::unique_ptr<paritypp::server> server_;

    std::unique_ptr<paritypp::client> pclient_;

    // Pool of dedicated read clients — each has independent QUIC connections
    // and its own mutex, so concurrent preview/download requests don't
    // serialize on a single client.
    static constexpr size_t READ_POOL_SIZE = 4;
    std::vector<std::unique_ptr<paritypp::client>> read_pool_;
    std::atomic<size_t> read_pool_rr_{0};  // round-robin index

    paritypp::client* next_read_client();

    // Dedicated client for scrub/rebalance — avoids blocking pclient_
    // (used by replicate and health probing) during long scrub operations
    std::unique_ptr<paritypp::client> scrub_client_;

    // Dedicated client for import streaming — avoids blocking pclient_
    // during long import operations so normal writes remain responsive
    std::unique_ptr<paritypp::client> import_client_;

    int self_node_index_ = -1;
    std::unordered_map<size_t, size_t> peer_to_client_index_;  // cfg peer idx -> pclient node idx

    void server_loop();
    void health_loop();
};

// Global cluster instance (nullptr when cluster is disabled)
extern Cluster* g_Cluster;

} // namespace mediadb
