#pragma once

#include <socket.h>

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

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

#ifdef _WIN32
#include <sys/types.h>
typedef intptr_t ssize_t;
#endif

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.
// When max_bytes > 0, acts as a byte-budget LRU cache: stores are admitted
// up to the budget, and least-recently-used groups are evicted to make room.
class memory_block_store : public block_store {
public:
    // Unbounded (original behaviour)
    memory_block_store() = default;

    // Byte-budget LRU mode: evicts oldest groups when used_bytes > max_bytes
    explicit memory_block_store(size_t max_bytes);

    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;

    size_t used_bytes() const;

private:
    size_t max_bytes_ = 0;   // 0 = unbounded
    size_t used_bytes_ = 0;
    std::mutex mutex_;
    // group_id -> (block_index -> data)
    std::map<uint64_t, std::map<uint32_t, std::vector<uint8_t>>> blocks_;
    // LRU order: front = most recently used, back = least recently used
    std::list<uint64_t> lru_;
    std::unordered_map<uint64_t, std::list<uint64_t>::iterator> lru_map_;

    // Remove the least-recently-used group (caller must hold mutex_)
    void evict_lru();
    // Promote group to MRU position (caller must hold mutex_)
    void touch(uint64_t group_id);
    // Byte size of a group (caller must hold mutex_)
    size_t group_bytes(uint64_t group_id) const;
};

// 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.
//
// Write-hole protection:
//   - Per-record CRC32 checksum to detect corruption and truncated writes
//   - Write-ahead journal (blocks.wal) logs intent before data modification;
//     on startup, incomplete journal entries are replayed or rolled back
//   - fdatasync after every flush to ensure durability
//   - Crash-safe vacuum: fsync tmp file + directory after rename
class file_block_store : public block_store {
public:
    explicit file_block_store(const std::string& base_dir, size_t write_buf_threshold = DEFAULT_WRITE_BUF_THRESHOLD);
    ~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;

    // Flush buffered writes to disk immediately.
    void flush();

    // CRC32 computation (IEEE polynomial)
    static uint32_t crc32(const uint8_t* data, size_t len);

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][4B crc32]
    static constexpr size_t RECORD_HEADER_SIZE = 21;
    // Write buffer threshold: flush when buffer exceeds this size (0 = immediate flush)
    static constexpr size_t DEFAULT_WRITE_BUF_THRESHOLD = 256 * 1024; // 256 KB

    struct block_entry {
        off_t  offset;    // file offset of the record header
        uint32_t data_len;
        bool in_write_buf; // true if data is still in write_buf_ (not yet flushed)
    };

    std::string base_dir_;
    std::string data_path_;
    std::string wal_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;
    off_t dead_bytes_ = 0;     // accumulated tombstoned bytes (header + data)
    off_t file_size_  = 0;    // current file size (including unflushed)

    // Write coalescing buffer
    std::vector<uint8_t> write_buf_;
    off_t write_buf_base_ = 0; // file offset where write_buf_ starts
    size_t write_buf_threshold_ = DEFAULT_WRITE_BUF_THRESHOLD;

    // Deferred tombstones: file offsets of old records to mark DELETED
    // after the replacement record is durably flushed. This prevents data
    // loss when a crash occurs between tombstoning the old record and
    // flushing the new one.
    std::vector<off_t> pending_tombstones_;

    // WAL journal operations
    // WAL record: [1B op][8B group_id][4B block_index][4B data_len][4B crc32]
    //   op=0x01: STORE (intent to append)
    //   op=0x02: TOMBSTONE (intent to mark deleted)
    //   op=0x03: COMMIT (previous ops are durable)
    static constexpr uint8_t WAL_OP_STORE     = 0x01;
    static constexpr uint8_t WAL_OP_TOMBSTONE = 0x02;
    static constexpr uint8_t WAL_OP_COMMIT    = 0x03;
    static constexpr size_t  WAL_RECORD_SIZE  = 21; // 1+8+4+4+4

    int wal_fd_ = -1;
    void wal_open();
    void wal_close();
    void wal_log(uint8_t op, uint64_t group_id, uint32_t block_index,
                 uint32_t data_len, uint32_t crc);
    void wal_commit();  // write COMMIT marker + fdatasync
    void wal_clear();   // truncate WAL after successful data flush
    void wal_recover(); // replay/rollback incomplete WAL entries on startup

    bool scan_and_build_index();
    void maybe_auto_vacuum();
    bool vacuum_impl(); // internal vacuum, caller must hold mutex_
    void flush_impl();  // flush write buffer, caller must hold mutex_
    void fsync_dir();   // fsync the base directory (for rename durability)
};

// 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
        // On startup, wait up to this many seconds for at least one peer
        // to become reachable before accepting client requests.
        // Prevents data-unavailable errors during full cluster restarts.
        // 0 = don't wait (legacy behaviour).
        int startup_peer_wait_secs = 0;
    };

    // 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();

    // Register an additional block store under the given store_id.
    // store_id 0 is the primary store passed to the constructor.
    void add_store(uint8_t store_id, std::shared_ptr<block_store> store);

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

    // Set rebuilding flag externally (e.g. during cluster scrub)
    void set_rebuilding(bool v) { rebuilding_ = v; }

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

    // Stop the server
    void stop();

private:
    config cfg_;
    std::shared_ptr<block_store> store_;     // primary store (store_id=0)
    std::map<uint8_t, std::shared_ptr<block_store>> stores_; // all stores by id
    std::shared_ptr<auth_store> auth_store_;
    std::unique_ptr<netplus::quic> socket_;
    std::atomic<bool> running_{false};
    std::atomic<bool> rebuilding_{false};
    std::atomic<bool> degraded_{false};  // true when startup peer wait timed out
    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,
                          uint8_t store_id);

    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
