#pragma once

#include <socket.h>

#include "protocol.h"
#include "parity.h"
#include "auth.h"

#include <string>
#include <vector>
#include <map>
#include <unordered_map>
#include <mutex>
#include <memory>
#include <functional>
#include <atomic>
#include <thread>

namespace paritypp {

// Block storage backend interface
class block_store {
public:
    virtual ~block_store() = default;

    virtual bool store(uint64_t group_id, uint32_t block_index,
                       const uint8_t* data, size_t len) = 0;
    virtual bool fetch(uint64_t group_id, uint32_t block_index,
                       std::vector<uint8_t>& out) = 0;
    virtual bool remove_group(uint64_t group_id) = 0;
    virtual std::vector<uint32_t> list_blocks(uint64_t group_id) = 0;

    // Status: total number of blocks and groups stored
    virtual uint32_t total_blocks() = 0;
    virtual uint32_t total_groups() = 0;

    // List all group IDs in the store
    virtual std::vector<uint64_t> list_groups() = 0;

    // Compact the store, reclaiming space from deleted records.
    // Returns true on success. Default implementation is a no-op.
    virtual bool vacuum() { return true; }
};

// In-memory block store for testing/lightweight usage
class memory_block_store : public block_store {
public:
    bool store(uint64_t group_id, uint32_t block_index,
               const uint8_t* data, size_t len) override;
    bool fetch(uint64_t group_id, uint32_t block_index,
               std::vector<uint8_t>& out) override;
    bool remove_group(uint64_t group_id) override;
    std::vector<uint32_t> list_blocks(uint64_t group_id) override;
    uint32_t total_blocks() override;
    uint32_t total_groups() override;
    std::vector<uint64_t> list_groups() override;

private:
    std::mutex mutex_;
    // group_id -> (block_index -> data)
    std::map<uint64_t, std::map<uint32_t, std::vector<uint8_t>>> blocks_;
};

// Append-only file block store
// All blocks stored in a single binary file with in-memory index.
// Deleted blocks are tombstoned; vacuum rewrites the file to reclaim space.
class file_block_store : public block_store {
public:
    explicit file_block_store(const std::string& base_dir);
    ~file_block_store() override;

    bool store(uint64_t group_id, uint32_t block_index,
               const uint8_t* data, size_t len) override;
    bool fetch(uint64_t group_id, uint32_t block_index,
               std::vector<uint8_t>& out) override;
    bool remove_group(uint64_t group_id) override;
    std::vector<uint32_t> list_blocks(uint64_t group_id) override;
    uint32_t total_blocks() override;
    uint32_t total_groups() override;
    std::vector<uint64_t> list_groups() override;

    // Rewrite the data file, dropping deleted records to reclaim space.
    // Returns true on success.
    bool vacuum() override;

private:
    static constexpr uint8_t RECORD_LIVE    = 0x00;
    static constexpr uint8_t RECORD_DELETED = 0xFF;
    // Record header: [1B flags][8B group_id][4B block_index][4B data_len]
    static constexpr size_t RECORD_HEADER_SIZE = 17;

    struct block_entry {
        off_t  offset;    // file offset of the record header
        uint32_t data_len;
    };

    std::string data_path_;
    int fd_ = -1;
    std::mutex mutex_;
    // group_id -> (block_index -> entry)
    std::unordered_map<uint64_t,
        std::unordered_map<uint32_t, block_entry>> index_;
    uint32_t total_blocks_ = 0;

    bool scan_and_build_index();
};

// QUIC-based parity storage server
// Listens on a QUIC endpoint and handles block store/fetch/delete requests.
// Runs a simple accept loop — QUIC datagrams are processed inside accept()
// which fires stream callbacks for complete requests.
class server {
public:
    struct config {
        std::string bind_address = "0.0.0.0";
        int port = 4433;
        std::string cert_file;    // X.509 certificate PEM/DER path
        std::string key_file;     // Private key DER path
        int max_connections = 128;
        bool require_auth = false; // Enable client authentication
    };

    // Peer node info for auto-rebuild
    struct peer_info {
        std::string address;
        int port;
    };

    // Credentials for connecting to peer nodes
    struct peer_credentials {
        std::string client_name;
        std::vector<uint8_t> key;
    };

    server(const config& cfg, std::shared_ptr<block_store> store,
           std::shared_ptr<auth_store> auth = nullptr);
    ~server();

    // Set a callback that provides peer connectivity info for status queries.
    // The callback should return (peers_total, peers_online).
    using peer_status_fn = std::function<std::pair<uint32_t, uint32_t>()>;
    void set_peer_status_callback(peer_status_fn fn);

    // Configure peer nodes for auto-rebuild.
    // When set, the server will automatically sync missing blocks from peers
    // on startup and periodically.
    void set_peers(const std::vector<peer_info>& peers,
                   const peer_credentials& creds = {});

    // Trigger a rebuild now (non-blocking, runs in background)
    void trigger_rebuild();

    // Check if a rebuild is currently in progress
    bool is_rebuilding() const { return rebuilding_; }

    // Start the server (blocks until stopped)
    void run();

    // Stop the server
    void stop();

private:
    config cfg_;
    std::shared_ptr<block_store> store_;
    std::shared_ptr<auth_store> auth_store_;
    std::unique_ptr<netplus::quic> socket_;
    std::atomic<bool> running_{false};
    std::atomic<bool> rebuilding_{false};
    peer_status_fn peer_status_fn_;

    // Peer rebuild config
    std::vector<peer_info> peers_;
    peer_credentials peer_creds_;
    std::unique_ptr<std::thread> rebuild_thread_;

    void rebuild_from_peers();

    // Auth state per connection
    struct auth_state {
        bool authenticated = false;
        std::string client_name;
        std::vector<uint8_t> pending_nonce; // Nonce waiting for response
    };
    std::mutex auth_mutex_;
    // keyed by quic* raw pointer for per-connection auth
    std::map<netplus::quic*, auth_state> auth_states_;

    bool is_authenticated(netplus::quic* conn);
    void handle_auth_request(netplus::quic& conn, uint64_t stream_id,
                              const uint8_t* payload, size_t len);
    void handle_auth_response(netplus::quic& conn, uint64_t stream_id,
                               const uint8_t* payload, size_t len);

    void handle_stream_data(netplus::quic& conn, uint64_t stream_id,
                             const std::vector<uint8_t>& data);
    void process_message(netplus::quic& conn, uint64_t stream_id,
                          msg_type type, const uint8_t* payload, size_t len);

    void send_response(netplus::quic& conn, uint64_t stream_id,
                        msg_type type, const std::vector<uint8_t>& payload);
    void send_error(netplus::quic& conn, uint64_t stream_id,
                     uint16_t code, const std::string& message);
};

} // namespace paritypp
