#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>

namespace paritypp {

class block_store; // forward declaration

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

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

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

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

    uint64_t generate_group_id();
};

} // namespace paritypp
