#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>
#include <mutex>
#include <unordered_map>

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

    // Rebalance block placement across all nodes.
    // For each group: retrieve data, delete the group, re-store so blocks land
    // on their correct nodes (block_index % n == node_index).
    // Returns {groups_rebalanced, groups_failed}.
    struct rebalance_result {
        size_t rebalanced = 0;
        size_t failed = 0;
        size_t already_ok = 0;
    };
    rebalance_result rebalance();

    // Distributed scrub/rebalance lock using a well-known parity group.
    // Prevents concurrent scrub/rebalance across cluster nodes.
    // lock_ttl_seconds: how long the lock is considered valid (default 5 min).
    // Returns true if lock was acquired.
    bool try_scrub_lock(int lock_ttl_seconds = 300);

    // Release the distributed scrub lock.
    void release_scrub_lock();

    // Well-known group_id used for the scrub lock.
    static uint64_t scrub_lock_group_id();

    // Mark a node as dead (skip it for the cooldown period)
    void mark_node_dead(size_t node_index);

    // Pre-establish QUIC connections + authentication to all nodes.
    // Call once after construction to avoid first-request latency.
    void warmup();

    // Set the store_id used for all subsequent operations on this client.
    // Default is 0 (primary store).
    void set_store_id(uint8_t id) { store_id_ = id; }
    uint8_t get_store_id() const { return store_id_; }

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

    // --- Batch API ---
    // Store multiple items packed into a single parity group.
    // Each item has a uint64_t key and a data payload.
    // Format: [4B 'BTCH'][4B count][for each: 8B key, 4B len, data...]
    struct batch_item {
        uint64_t key;
        std::vector<uint8_t> data;
    };

    // Store a batch of items as one parity group under batch_gid.
    void store_batch(uint64_t batch_gid, const std::vector<batch_item>& items);

    // Retrieve all items from a batch group.
    std::vector<batch_item> retrieve_batch(uint64_t batch_gid);

    // Retrieve a single item by key from a batch group.
    // Returns true if found, false otherwise.
    bool retrieve_batch_item(uint64_t batch_gid, uint64_t item_key,
                             std::vector<uint8_t>& out);

    // Encode items into batch payload (static utility).
    static std::vector<uint8_t> encode_batch(const std::vector<batch_item>& items);

    // Decode batch payload into items (static utility).
    static bool decode_batch(const std::vector<uint8_t>& payload,
                             std::vector<batch_item>& items);

    friend class client_store_session;

private:
    size_t k_;
    size_t m_;
    uint8_t store_id_ = 0;
    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;
        // Track dead nodes to skip them quickly
        std::chrono::steady_clock::time_point dead_until{};
    };

    std::vector<quic_connection> connections_;
    netplus::socketwait          wait_;
    mutable std::mutex           mutex_;  // thread-safety for all public methods

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

    // Unlocked node status probe (for internal parallel use)
    bool get_node_status_impl(size_t node_index, status_info& info);

    // 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.
    // Retries with backoff if the cluster is rebuilding.
    // The caller's unique_lock is released during retry sleeps so that other
    // client operations are not blocked by long backoff waits.
    bool retrieve_stripe(std::unique_lock<std::mutex>& lk,
                         uint64_t group_id, uint32_t stripe_index,
                         std::vector<uint8_t>& out);

    // Single-attempt stripe retrieval (internal).
    // Sets should_retry=true when some blocks were found but < k.
    bool retrieve_stripe_once(uint64_t group_id, uint32_t stripe_index,
                               std::vector<uint8_t>& out, bool& should_retry);

    uint64_t generate_group_id();
};

} // namespace paritypp
