#include "cluster.h"

#include <netplus/exception.h>

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

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

void Cluster::start() {
    if (running_) return;
    running_ = true;
    server_thread_ = std::thread(&Cluster::server_loop, this);
}

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

std::vector<Cluster::NodeGroupInfo> Cluster::list_peer_groups() {
    std::vector<NodeGroupInfo> result;
    if (!pclient_) return result;
    std::unique_lock<std::timed_mutex> lock(client_mutex_, std::chrono::seconds(30));
    if (!lock.owns_lock()) 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;
}

std::unique_lock<std::timed_mutex> Cluster::acquire_client_lock(int timeout_secs) {
    std::unique_lock<std::timed_mutex> lock(client_mutex_, std::chrono::seconds(timeout_secs));
    return lock;
}

bool Cluster::replicate_unlocked(const std::string& key, const uint8_t* data, size_t len) {
    if (!pclient_) return false;
    uint64_t gid = cluster_group_id(key);
    try {
        pclient_->store(gid, data, len);
    } catch (const netplus::NetException& e) {
        return false;
    } catch (const std::exception& e) {
        return false;
    }
    return true;
}

std::unique_ptr<paritypp::store_session> Cluster::begin_replicate_unlocked(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;
    }
}

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 {
        std::unique_lock<std::timed_mutex> lock(client_mutex_, std::chrono::seconds(120));
        if (!lock.owns_lock()) throw std::runtime_error("cluster client lock timeout");
        pclient_->store(gid, data, len);
    } catch (const netplus::NetException& e) {
        return false;
    } catch (const std::exception& e) {
        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 {
        // Acquire the lock and transfer ownership into a wrapper so the
        // lock is held for the entire lifetime of the store_session.
        // This prevents concurrent replicate() / fetch() / remove() calls
        // from using the shared QUIC connections while we're streaming.
        auto lock = std::make_unique<std::unique_lock<std::timed_mutex>>(
            client_mutex_, std::chrono::seconds(120));
        if (!lock->owns_lock()) throw std::runtime_error("cluster client lock timeout");
        auto inner = pclient_->begin_store(gid, total_size);
        if (!inner) return nullptr;

        struct LockedSession : paritypp::store_session {
            std::unique_ptr<std::unique_lock<std::timed_mutex>> lock;
            std::unique_ptr<paritypp::store_session> inner;
            LockedSession(std::unique_ptr<std::unique_lock<std::timed_mutex>> l,
                          std::unique_ptr<paritypp::store_session> s)
                : lock(std::move(l)), inner(std::move(s)) {}
            void feed(const uint8_t* data, size_t len) override { inner->feed(data, len); }
            void finalize() override {
                inner->finalize();
                // Release the lock immediately after finalize so that
                // subsequent metadata replicate() calls (store buffer,
                // index) can acquire it without deadlocking.
                if (lock && lock->owns_lock()) lock->unlock();
            }
            bool ok() const override { return inner->ok(); }
            std::string error_message() const override { return inner->error_message(); }
        };
        return std::make_unique<LockedSession>(std::move(lock), std::move(inner));
    } catch (const std::exception& e) {
        return nullptr;
    }
}

bool Cluster::fetch(const std::string& key, std::vector<uint8_t>& out) {
    if (!pclient_) return false;
    uint64_t gid = cluster_group_id(key);
    try {
        std::unique_lock<std::timed_mutex> lock(client_mutex_, std::chrono::seconds(60));
        if (!lock.owns_lock()) throw std::runtime_error("cluster client lock timeout");
        out = pclient_->retrieve(gid);
        return !out.empty();
    } catch (const netplus::NetException& e) {
    } catch (const std::exception& e) {
    }
    return false;
}

bool Cluster::fetch_from_peers(const std::string& key, std::vector<uint8_t>& out) {
    if (!pclient_) return false;
    uint64_t gid = cluster_group_id(key);
    try {
        std::unique_lock<std::timed_mutex> lock(client_mutex_, std::chrono::seconds(60));
        if (!lock.owns_lock()) throw std::runtime_error("cluster client lock timeout");
        out = pclient_->retrieve(gid);
        return !out.empty();
    } catch (const netplus::NetException& e) {
    } catch (const std::exception& e) {
    }
    return false;
}

void Cluster::remove(const std::string& key) {
    if (!pclient_) return;
    uint64_t gid = cluster_group_id(key);
    try {
        std::unique_lock<std::timed_mutex> lock(client_mutex_, std::chrono::seconds(30));
        if (!lock.owns_lock()) throw std::runtime_error("cluster client lock timeout");
        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 (!pclient_) return 0;
    std::unique_lock<std::timed_mutex> lock(client_mutex_, std::chrono::seconds(120));
    if (!lock.owns_lock()) return 0;
    return pclient_->vacuum_all_nodes();
}

void Cluster::server_loop() {
    try {
        server_->run();
    } catch (const netplus::NetException& e) {
    } catch (const std::exception& e) {
    }
    running_ = false;
}

} // namespace mediadb
