#include "server.h"
#include "client.h"

#include <crypto/x509.h>
#include <crypto/rsa.h>
#include <exception.h>

#include <cstring>
#include <fstream>
#include <sstream>
#include <filesystem>
#include <algorithm>
#include <thread>
#include <chrono>
#include <set>

namespace paritypp {

// --- memory_block_store ---

bool memory_block_store::store(uint64_t group_id, uint32_t block_index,
                                const uint8_t* data, size_t len) {
    std::lock_guard<std::mutex> lock(mutex_);
    blocks_[group_id][block_index].assign(data, data + len);
    return true;
}

bool memory_block_store::fetch(uint64_t group_id, uint32_t block_index,
                                std::vector<uint8_t>& out) {
    std::lock_guard<std::mutex> lock(mutex_);
    auto git = blocks_.find(group_id);
    if (git == blocks_.end()) return false;
    auto bit = git->second.find(block_index);
    if (bit == git->second.end()) return false;
    out = bit->second;
    return true;
}

bool memory_block_store::remove_group(uint64_t group_id) {
    std::lock_guard<std::mutex> lock(mutex_);
    return blocks_.erase(group_id) > 0;
}

std::vector<uint32_t> memory_block_store::list_blocks(uint64_t group_id) {
    std::lock_guard<std::mutex> lock(mutex_);
    std::vector<uint32_t> result;
    auto git = blocks_.find(group_id);
    if (git != blocks_.end()) {
        for (const auto& kv : git->second) {
            result.push_back(kv.first);
        }
    }
    return result;
}

uint32_t memory_block_store::total_blocks() {
    std::lock_guard<std::mutex> lock(mutex_);
    uint32_t count = 0;
    for (const auto& g : blocks_) {
        count += static_cast<uint32_t>(g.second.size());
    }
    return count;
}

uint32_t memory_block_store::total_groups() {
    std::lock_guard<std::mutex> lock(mutex_);
    return static_cast<uint32_t>(blocks_.size());
}

std::vector<uint64_t> memory_block_store::list_groups() {
    std::lock_guard<std::mutex> lock(mutex_);
    std::vector<uint64_t> result;
    result.reserve(blocks_.size());
    for (const auto& kv : blocks_) {
        result.push_back(kv.first);
    }
    return result;
}

// --- file_block_store ---

file_block_store::file_block_store(const std::string& base_dir)
    : base_dir_(base_dir) {
    std::filesystem::create_directories(base_dir_);
}

std::string file_block_store::group_dir(uint64_t group_id) const {
    std::ostringstream oss;
    oss << base_dir_ << "/" << std::hex << group_id;
    return oss.str();
}

std::string file_block_store::block_path(uint64_t group_id, uint32_t block_index) const {
    std::ostringstream oss;
    oss << group_dir(group_id) << "/" << block_index << ".blk";
    return oss.str();
}

bool file_block_store::store(uint64_t group_id, uint32_t block_index,
                              const uint8_t* data, size_t len) {
    std::string dir = group_dir(group_id);
    std::error_code ec;
    std::filesystem::create_directories(dir, ec);
    if (ec) {
        std::cerr << "file_block_store::store: create_directories failed for "
                  << dir << ": " << ec.message() << std::endl;
        return false;
    }
    std::string path = block_path(group_id, block_index);
    std::ofstream ofs(path, std::ios::binary);
    if (!ofs) {
        std::cerr << "file_block_store::store: cannot open " << path
                  << ": " << strerror(errno) << std::endl;
        return false;
    }
    ofs.write(reinterpret_cast<const char*>(data), static_cast<std::streamsize>(len));
    if (!ofs.good()) {
        std::cerr << "file_block_store::store: write failed for " << path
                  << ": " << strerror(errno) << std::endl;
        return false;
    }
    return true;
}

bool file_block_store::fetch(uint64_t group_id, uint32_t block_index,
                              std::vector<uint8_t>& out) {
    std::string path = block_path(group_id, block_index);
    std::ifstream ifs(path, std::ios::binary | std::ios::ate);
    if (!ifs) return false;
    auto size = ifs.tellg();
    ifs.seekg(0);
    out.resize(static_cast<size_t>(size));
    ifs.read(reinterpret_cast<char*>(out.data()), size);
    return ifs.good();
}

bool file_block_store::remove_group(uint64_t group_id) {
    std::string dir = group_dir(group_id);
    std::error_code ec;
    return std::filesystem::remove_all(dir, ec) > 0;
}

std::vector<uint32_t> file_block_store::list_blocks(uint64_t group_id) {
    std::vector<uint32_t> result;
    std::string dir = group_dir(group_id);
    std::error_code ec;
    if (!std::filesystem::exists(dir, ec)) return result;

    for (const auto& entry : std::filesystem::directory_iterator(dir, ec)) {
        if (entry.is_regular_file()) {
            std::string stem = entry.path().stem().string();
            try {
                result.push_back(static_cast<uint32_t>(std::stoul(stem)));
            } catch (...) {}
        }
    }
    std::sort(result.begin(), result.end());
    return result;
}

uint32_t file_block_store::total_blocks() {
    uint32_t count = 0;
    std::error_code ec;
    if (!std::filesystem::exists(base_dir_, ec)) return 0;
    for (const auto& group_entry : std::filesystem::directory_iterator(base_dir_, ec)) {
        if (group_entry.is_directory()) {
            for (const auto& blk : std::filesystem::directory_iterator(group_entry.path(), ec)) {
                if (blk.is_regular_file()) ++count;
            }
        }
    }
    return count;
}

uint32_t file_block_store::total_groups() {
    uint32_t count = 0;
    std::error_code ec;
    if (!std::filesystem::exists(base_dir_, ec)) return 0;
    for (const auto& entry : std::filesystem::directory_iterator(base_dir_, ec)) {
        if (entry.is_directory()) ++count;
    }
    return count;
}

std::vector<uint64_t> file_block_store::list_groups() {
    std::vector<uint64_t> result;
    std::error_code ec;
    if (!std::filesystem::exists(base_dir_, ec)) return result;
    for (const auto& entry : std::filesystem::directory_iterator(base_dir_, ec)) {
        if (entry.is_directory()) {
            try {
                result.push_back(std::stoull(entry.path().filename().string(), nullptr, 16));
            } catch (...) {}
        }
    }
    return result;
}

// --- server ---

static std::map<std::string, netplus::ssl::CertificateBundle> load_certs(
    const std::string& cert_file, const std::string& key_file)
{
    std::map<std::string, netplus::ssl::CertificateBundle> certs;
    netplus::ssl::CertificateBundle bundle;

    if (!bundle.cert.loadFromFile(cert_file)) {
        std::cerr << "[paritypp] WARN: failed to load cert from " << cert_file << std::endl;
    }

    std::ifstream kf(key_file, std::ios::binary | std::ios::ate);
    if (kf) {
        auto sz = kf.tellg();
        kf.seekg(0);
        bundle.privateKeyDer.resize(static_cast<size_t>(sz));
        kf.read(reinterpret_cast<char*>(bundle.privateKeyDer.data()), sz);
        bundle.rsa_key = netplus::rsa(bundle.privateKeyDer);
    } else {
        std::cerr << "[paritypp] WARN: failed to open key file " << key_file << std::endl;
    }

    std::string cn = bundle.cert.getSubjectCN();
    std::cerr << "[paritypp] loaded cert CN='" << cn << "' key=" << bundle.privateKeyDer.size() << " bytes" << std::endl;
    certs[cn] = std::move(bundle);
    return certs;
}

server::server(const config& cfg, std::shared_ptr<block_store> store,
               std::shared_ptr<auth_store> auth)
    : cfg_(cfg), store_(std::move(store)), auth_store_(std::move(auth)) {}

server::~server() {
    stop();
    if (rebuild_thread_ && rebuild_thread_->joinable()) {
        rebuild_thread_->join();
    }
}

void server::set_peer_status_callback(peer_status_fn fn) {
    peer_status_fn_ = std::move(fn);
}

void server::set_peers(const std::vector<peer_info>& peers,
                       const peer_credentials& creds) {
    peers_ = peers;
    peer_creds_ = creds;
}

void server::trigger_rebuild() {
    if (rebuilding_) return;
    if (peers_.empty()) return;

    // Join previous rebuild thread if still around
    if (rebuild_thread_ && rebuild_thread_->joinable()) {
        rebuild_thread_->join();
    }

    rebuild_thread_ = std::make_unique<std::thread>([this]() {
        rebuild_from_peers();
    });
}

void server::rebuild_from_peers() {
    rebuilding_ = true;
    std::cerr << "paritypp::server: rebuild started" << std::endl;

    // Build a client to talk to peers
    std::vector<client::node_info> nodes;
    nodes.reserve(peers_.size());
    for (const auto& p : peers_) {
        nodes.push_back({p.address, p.port});
    }

    client::credentials creds;
    creds.client_name = peer_creds_.client_name;
    creds.key = peer_creds_.key;

    try {
        // Use k=1, m=0 — we only need simple fetch, not parity distribution
        client peer_client(1, 0, nodes, creds);

        // Collect all group IDs across all peers
        std::set<uint64_t> remote_groups;
        for (size_t i = 0; i < nodes.size() && running_; i++) {
            std::vector<uint64_t> groups;
            if (peer_client.list_groups_on_node(i, groups)) {
                remote_groups.insert(groups.begin(), groups.end());
            }
        }

        // Get our local groups
        auto local_groups_vec = store_->list_groups();
        std::set<uint64_t> local_groups(local_groups_vec.begin(),
                                         local_groups_vec.end());

        uint32_t synced_groups = 0;
        uint32_t synced_blocks = 0;

        // For each remote group, check which blocks we're missing
        for (uint64_t group_id : remote_groups) {
            if (!running_) break;

            // Get block lists from all peers for this group
            std::map<uint32_t, size_t> block_to_peer; // block_index -> first peer that has it
            for (size_t i = 0; i < nodes.size(); i++) {
                auto listing = peer_client.list(group_id);
                for (const auto& [node_idx, indices] : listing) {
                    for (uint32_t bidx : indices) {
                        if (block_to_peer.find(bidx) == block_to_peer.end()) {
                            block_to_peer[bidx] = node_idx;
                        }
                    }
                }
                break; // list() already queries all nodes
            }

            // Check which blocks we have locally
            auto local_blocks = store_->list_blocks(group_id);
            std::set<uint32_t> local_set(local_blocks.begin(), local_blocks.end());

            bool group_changed = false;
            for (const auto& [bidx, peer_idx] : block_to_peer) {
                if (!running_) break;
                if (local_set.count(bidx)) continue; // Already have it

                // Fetch this block from a peer
                std::vector<uint8_t> block_data;
                if (peer_client.fetch_from_peers(group_id, bidx, block_data)) {
                    if (store_->store(group_id, bidx, block_data.data(), block_data.size())) {
                        synced_blocks++;
                        group_changed = true;
                    }
                }
            }
            if (group_changed) synced_groups++;
        }

        std::cerr << "paritypp::server: rebuild complete — synced "
                  << synced_blocks << " block(s) across "
                  << synced_groups << " group(s)" << std::endl;

    } catch (const std::exception& e) {
        std::cerr << "paritypp::server: rebuild failed: " << e.what() << std::endl;
    }

    rebuilding_ = false;
}

void server::run() {
    auto certs = load_certs(cfg_.cert_file, cfg_.key_file);

    socket_ = std::make_unique<netplus::quic>(
        certs, cfg_.bind_address, cfg_.port, cfg_.max_connections, 0);

    socket_->bind();
    socket_->setNonBlock();

    // Wire up stream callback — fired from inside accept() when
    // child connections receive complete stream data.
    // NOTE: The data parameter contains the full stream payload;
    // recv_buffer has already been moved out by processStreamFrame,
    // so we must use the callback data directly.
    socket_->setStreamCallback(
        [this](netplus::socket* sock, uint64_t stream_id,
               const std::vector<uint8_t>& data, bool /*fin*/) {
            auto* conn = static_cast<netplus::quic*>(sock);
            handle_stream_data(*conn, stream_id, data);
        });

    running_ = true;

    // Auto-rebuild: sync missing blocks from peers on startup
    if (!peers_.empty()) {
        trigger_rebuild();
    }

    // Simple accept loop — QUIC processes datagrams and fires
    // stream callbacks inside accept().
    auto last_resync = std::chrono::steady_clock::now();
    while (running_) {
        try {
            std::unique_ptr<netplus::socket> csock;
            socket_->accept(csock, true);
        } catch (const netplus::NetException& e) {
            if (e.getErrorType() != netplus::NetException::Note) {
                std::cerr << "paritypp::server accept error: " << e.what() << std::endl;
            }
        } catch (...) {}

        // Periodic resync every 60 seconds
        auto now = std::chrono::steady_clock::now();
        if (!peers_.empty() &&
            std::chrono::duration_cast<std::chrono::seconds>(now - last_resync).count() >= 60) {
            last_resync = now;
            trigger_rebuild();
        }

        std::this_thread::sleep_for(std::chrono::milliseconds(1));
    }
}

void server::stop() {
    running_ = false;
}

void server::handle_stream_data(netplus::quic& conn, uint64_t stream_id,
                                 const std::vector<uint8_t>& data) {
    // Parse the message directly from the callback data.
    // The recv_buffer has already been moved out by processStreamFrame,
    // so we cannot call recvStreamData() here.
    if (data.size() < protocol::HEADER_SIZE) return;

    msg_type type;
    uint32_t payload_len;
    if (!protocol::decode_header(data.data(), data.size(), type, payload_len)) return;

    if (protocol::HEADER_SIZE + payload_len > data.size()) return;

    process_message(conn, stream_id, type,
                    data.data() + protocol::HEADER_SIZE, payload_len);
}

bool server::is_authenticated(netplus::quic* conn) {
    if (!cfg_.require_auth || !auth_store_) return true;
    std::lock_guard<std::mutex> lock(auth_mutex_);
    auto it = auth_states_.find(conn);
    return it != auth_states_.end() && it->second.authenticated;
}

void server::handle_auth_request(netplus::quic& conn, uint64_t stream_id,
                                   const uint8_t* payload, size_t len) {
    std::string client_name;
    if (!protocol::parse_auth_request(payload, len, client_name)) {
        send_error(conn, stream_id, 10, "Invalid AUTH_REQUEST");
        return;
    }

    // Check if client exists
    std::vector<uint8_t> key;
    if (!auth_store_->lookup(client_name, key)) {
        send_response(conn, stream_id, msg_type::AUTH_DENIED, {});
        return;
    }

    // Generate nonce and send challenge
    auto nonce = auth::generate_nonce();

    {
        std::lock_guard<std::mutex> lock(auth_mutex_);
        auto& state = auth_states_[&conn];
        state.client_name = client_name;
        state.pending_nonce = nonce;
        state.authenticated = false;
    }

    auto challenge = protocol::make_auth_challenge(nonce);
    send_response(conn, stream_id, msg_type::AUTH_CHALLENGE, challenge);
}

void server::handle_auth_response(netplus::quic& conn, uint64_t stream_id,
                                    const uint8_t* payload, size_t len) {
    std::vector<uint8_t> client_hmac;
    if (!protocol::parse_auth_response(payload, len, client_hmac)) {
        send_error(conn, stream_id, 10, "Invalid AUTH_RESPONSE");
        return;
    }

    std::string client_name;
    std::vector<uint8_t> nonce;
    {
        std::lock_guard<std::mutex> lock(auth_mutex_);
        auto it = auth_states_.find(&conn);
        if (it == auth_states_.end() || it->second.pending_nonce.empty()) {
            send_response(conn, stream_id, msg_type::AUTH_DENIED, {});
            return;
        }
        client_name = it->second.client_name;
        nonce = it->second.pending_nonce;
    }

    // Lookup key and verify
    std::vector<uint8_t> key;
    if (!auth_store_->lookup(client_name, key)) {
        send_response(conn, stream_id, msg_type::AUTH_DENIED, {});
        return;
    }

    if (auth::verify_hmac(key, nonce, client_hmac)) {
        {
            std::lock_guard<std::mutex> lock(auth_mutex_);
            auto& state = auth_states_[&conn];
            state.authenticated = true;
            state.pending_nonce.clear();
        }
        send_response(conn, stream_id, msg_type::AUTH_OK, {});
    } else {
        {
            std::lock_guard<std::mutex> lock(auth_mutex_);
            auth_states_.erase(&conn);
        }
        send_response(conn, stream_id, msg_type::AUTH_DENIED, {});
    }
}

void server::process_message(netplus::quic& conn, uint64_t stream_id,
                               msg_type type, const uint8_t* payload, size_t len) {
    switch (type) {
        case msg_type::AUTH_REQUEST: {
            handle_auth_request(conn, stream_id, payload, len);
            return;
        }
        case msg_type::AUTH_RESPONSE: {
            handle_auth_response(conn, stream_id, payload, len);
            return;
        }
        default:
            break;
    }

    // All other messages require authentication
    if (!is_authenticated(&conn)) {
        send_error(conn, stream_id, 11, "Authentication required");
        return;
    }

    switch (type) {
        case msg_type::STORE_BLOCK: {
            uint64_t group_id;
            uint32_t block_index;
            const uint8_t* data;
            size_t data_len;
            if (!protocol::parse_store_block(payload, len, group_id, block_index, data, data_len)) {
                send_error(conn, stream_id, 1, "Invalid STORE_BLOCK payload");
                return;
            }
            if (store_->store(group_id, block_index, data, data_len)) {
                send_response(conn, stream_id, msg_type::OK, {});
            } else {
                send_error(conn, stream_id, 2, "Store failed");
            }
            break;
        }

        case msg_type::FETCH_BLOCK: {
            uint64_t group_id;
            uint32_t block_index;
            if (!protocol::parse_fetch_block(payload, len, group_id, block_index)) {
                send_error(conn, stream_id, 1, "Invalid FETCH_BLOCK payload");
                return;
            }
            std::vector<uint8_t> block_data;
            if (store_->fetch(group_id, block_index, block_data)) {
                auto resp = protocol::make_block_data(group_id, block_index,
                                                       block_data.data(), block_data.size());
                send_response(conn, stream_id, msg_type::BLOCK_DATA, resp);
            } else {
                std::vector<uint8_t> not_found(12);
                protocol::encode_u64(not_found.data(), group_id);
                protocol::encode_u32(not_found.data() + 8, block_index);
                send_response(conn, stream_id, msg_type::BLOCK_NOT_FOUND, not_found);
            }
            break;
        }

        case msg_type::DELETE_GROUP: {
            if (len < 8) {
                send_error(conn, stream_id, 1, "Invalid DELETE_GROUP payload");
                return;
            }
            uint64_t group_id = protocol::decode_u64(payload);
            store_->remove_group(group_id);
            send_response(conn, stream_id, msg_type::OK, {});
            break;
        }

        case msg_type::LIST_BLOCKS: {
            if (len < 8) {
                send_error(conn, stream_id, 1, "Invalid LIST_BLOCKS payload");
                return;
            }
            uint64_t group_id = protocol::decode_u64(payload);
            auto indices = store_->list_blocks(group_id);

            std::vector<uint8_t> resp(8 + 4 + indices.size() * 4);
            protocol::encode_u64(resp.data(), group_id);
            protocol::encode_u32(resp.data() + 8, static_cast<uint32_t>(indices.size()));
            for (size_t i = 0; i < indices.size(); i++) {
                protocol::encode_u32(resp.data() + 12 + i * 4, indices[i]);
            }
            send_response(conn, stream_id, msg_type::BLOCK_LIST, resp);
            break;
        }

        case msg_type::GET_STATUS: {
            status_info info;
            // Gather block/group counts from the store
            info.block_count = store_->total_blocks();
            info.group_count = store_->total_groups();
            // Peer connectivity from callback
            if (peer_status_fn_) {
                auto [total, online] = peer_status_fn_();
                info.peers_total = total;
                info.peers_online = online;
            }
            // Determine status from peer connectivity
            if (info.peers_total > 0 && info.peers_online < info.peers_total) {
                info.status = node_status::DEGRADED;
            } else {
                info.status = node_status::OK;
            }
            // Override with REBUILD if rebuild is in progress
            if (rebuilding_) {
                info.status = node_status::REBUILD;
            }
            auto resp = protocol::make_status_resp(info);
            send_response(conn, stream_id, msg_type::STATUS_RESP, resp);
            break;
        }

        case msg_type::LIST_GROUPS: {
            auto groups = store_->list_groups();
            auto resp = protocol::make_group_list(groups);
            send_response(conn, stream_id, msg_type::GROUP_LIST, resp);
            break;
        }

        default:
            send_error(conn, stream_id, 3, "Unknown message type");
            break;
    }
}

void server::send_response(netplus::quic& conn, uint64_t stream_id,
                             msg_type type, const std::vector<uint8_t>& payload) {
    auto msg = protocol::encode_message(type, payload);
    conn.sendStreamData(stream_id, msg, true);
}

void server::send_error(netplus::quic& conn, uint64_t stream_id,
                          uint16_t code, const std::string& message) {
    std::vector<uint8_t> payload(2 + message.size());
    payload[0] = static_cast<uint8_t>(code >> 8);
    payload[1] = static_cast<uint8_t>(code);
    std::memcpy(payload.data() + 2, message.data(), message.size());
    auto msg = protocol::encode_message(msg_type::PROTOCOL_ERROR, payload);
    conn.sendStreamData(stream_id, msg, true);
}

} // namespace paritypp
