#pragma once

#include <socket.h>
#include <eventapi.h>

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

#include <string>
#include <vector>
#include <memory>
#include <functional>
#include <fstream>

namespace paritypp {

class block_store; // forward declaration

// Incremental store session: feed chunks, finalize computes parity + distributes
class store_session {
public:
    virtual ~store_session() = default;
    virtual void feed(const uint8_t* data, size_t len) = 0;
    virtual void finalize() = 0;
    virtual bool ok() const = 0;
    virtual std::string error_message() const { return {}; }
};

// Fixed stripe size for chunk-level parity encoding (4 MB)
static constexpr size_t STRIPE_SIZE = 4096 * 1024;

// QUIC client for distributing parity-encoded blocks across storage nodes
class client {
public:
    struct node_info {
        std::string address;
        int port;
    };

    // Authentication credentials (optional)
    struct credentials {
        std::string client_name;
        std::vector<uint8_t> key;
    };

    // Create client with given parity configuration
    // data_blocks: number of data blocks (k)
    // parity_blocks: number of parity blocks (m)
    // nodes: list of storage node addresses (should have at least k+m entries)
    client(size_t data_blocks, size_t parity_blocks,
           const std::vector<node_info>& nodes,
           const credentials& creds = {});
    ~client();

    // Store data: splits into blocks, computes parity, distributes across nodes
    // Returns a group_id that can be used to retrieve/delete the data
    uint64_t store(const uint8_t* data, size_t data_len);
    uint64_t store(const std::vector<uint8_t>& data);

    // Store with a caller-supplied group_id (deterministic key)
    // Prepends an 8-byte length header so retrieve can strip padding
    void store(uint64_t group_id, const uint8_t* data, size_t data_len);
    void store(uint64_t group_id, const std::vector<uint8_t>& data);

    // Retrieve data: fetches blocks from nodes, reconstructs if needed
    // original_size: original data size before padding
    std::vector<uint8_t> retrieve(uint64_t group_id, size_t original_size);

    // Retrieve with automatic size detection (reads 8-byte length header)
    std::vector<uint8_t> retrieve(uint64_t group_id);

    // Retrieve a byte range from stored data (fetches only needed stripes)
    // offset/length are in terms of the original data (before framing header)
    std::vector<uint8_t> retrieve_range(uint64_t group_id,
                                         uint64_t offset, uint64_t length);

    // Delete all blocks for a group from all nodes
    void remove(uint64_t group_id);

    // Check which blocks are available for a group
    std::vector<std::pair<size_t, std::vector<uint32_t>>> list(uint64_t group_id);

    // Replicate a raw block to all peer nodes (for domain-level replication)
    // Returns number of nodes that accepted the block
    size_t replicate_to_peers(uint64_t group_id, uint32_t block_index,
                            const uint8_t* data, size_t len);

    // Fetch a raw block from any peer node that has it
    bool fetch_from_peers(uint64_t group_id, uint32_t block_index,
                          std::vector<uint8_t>& out);

    // Query status of a specific node (returns false if unreachable)
    bool get_node_status(size_t node_index, status_info& info);

    // Query all nodes and compute overall cluster status
    struct cluster_health {
        node_status overall;              // OK, DEGRADED, or REBUILD
        uint32_t nodes_total;
        uint32_t nodes_online;
        std::vector<std::pair<size_t, status_info>> node_statuses; // (index, info)
    };
    cluster_health get_cluster_status();

    // List all group IDs stored on a specific node
    bool list_groups_on_node(size_t node_index, std::vector<uint64_t>& groups);

    // Compact block stores on all nodes, reclaiming space from deleted records.
    // Returns number of nodes that completed vacuum successfully.
    size_t vacuum_all_nodes();

    // Designate a node index as "local" — bypass QUIC, use block_store directly
    void set_local_node(size_t node_index, std::shared_ptr<block_store> store);

    // Begin an incremental store session: feed data in chunks, each STRIPE_SIZE
    // chunk is independently parity-encoded and distributed.
    // Peak RAM = STRIPE_SIZE * (k+m) ≈ tiny.
    // total_size: total payload size (needed for length framing header).
    std::unique_ptr<store_session> begin_store(uint64_t group_id, size_t total_size);

    friend class client_store_session;

private:
    size_t k_;
    size_t m_;
    std::vector<node_info> nodes_;
    credentials creds_;
    parity parity_engine_;

    // Local node bypass (optional)
    int local_node_index_ = -1;
    std::shared_ptr<block_store> local_store_;

    struct quic_connection {
        std::unique_ptr<netplus::quic> socket;
        bool connected = false;
    };

    std::vector<quic_connection> connections_;
    netplus::socketwait          wait_;

    void connect_to_node(size_t node_index);
    void ensure_connected(size_t node_index);
    bool authenticate_node(size_t node_index);

    // Send a message and receive response on a node connection
    std::vector<uint8_t> send_recv(size_t node_index,
                                    msg_type type,
                                    const std::vector<uint8_t>& payload);

    bool store_block_on_node(size_t node_index, uint64_t group_id,
                              uint32_t block_index,
                              const uint8_t* data, size_t len);

    bool fetch_block_from_node(size_t node_index, uint64_t group_id,
                                uint32_t block_index,
                                std::vector<uint8_t>& out);

    // Store one stripe (STRIPE_SIZE bytes, zero-padded) across all nodes
    void store_stripe(uint64_t group_id, uint32_t stripe_index,
                      const uint8_t* stripe_data, size_t stripe_len);

    // Retrieve and reconstruct one stripe from nodes
    bool retrieve_stripe(uint64_t group_id, uint32_t stripe_index,
                         std::vector<uint8_t>& out);

    uint64_t generate_group_id();
};

} // namespace paritypp
