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

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
    // deadline: overall wall-clock cutoff for the whole call (every stripe,
    // every retry). Defaults to none (existing behavior: each stripe gets
    // its own independent ~30s retry budget) — pass one when a caller
    // needs a genuinely bounded worst-case time, e.g. Cluster::fetch()
    // stacking a read-pool attempt and a pclient_ fallback and wanting
    // both together to cost ~30s, not ~30s each.
    std::vector<uint8_t> retrieve(uint64_t group_id, size_t original_size,
                                   std::chrono::steady_clock::time_point deadline =
                                       (std::chrono::steady_clock::time_point::max)());

    // Retrieve with automatic size detection (reads 8-byte length header)
    std::vector<uint8_t> retrieve(uint64_t group_id,
                                   std::chrono::steady_clock::time_point deadline =
                                       (std::chrono::steady_clock::time_point::max)());

    // 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,
                                         std::chrono::steady_clock::time_point deadline =
                                             (std::chrono::steady_clock::time_point::max)());

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

    // Liveness-check multiple nodes concurrently (LIST_GROUPS, response
    // discarded). Probes run in parallel so total latency is bounded by the
    // slowest single node rather than the sum of all node timeouts — use
    // this instead of looping list_groups_on_node() when checking several
    // nodes, e.g. for cluster health monitoring.
    // Returns one bool per entry in node_indices, same order.
    std::vector<bool> list_groups_on_nodes(const std::vector<size_t>& node_indices);

    // Check every known group's actual block placement via list(), rather
    // than group-ID presence alone. A node can hold *some* of a group's
    // blocks and still show up as "present" for that group in
    // list_groups_on_node, so presence-only checks cannot detect a node
    // that is quietly missing a subset of its blocks. incomplete lists
    // every group where a block landed on the wrong node, the total block
    // count isn't a clean multiple of n, or fewer than n nodes report any
    // blocks at all (including a group missing from a node entirely).
    struct group_check_result {
        std::vector<uint64_t> incomplete;
        size_t already_ok = 0;
    };
    group_check_result check_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{};
        // Last time this connection was known-alive (successful connect or
        // successful send_recv response). Nothing pumps idle connections in
        // the background here — pumpNetwork() only ever runs inline inside
        // an active call — so the QUIC layer's own PING keepalive (fires
        // after 15s idle, socket.h) never gets a chance to run for a
        // connection that just sits between infrequent calls (e.g. the
        // ~30s CLUSTER-SYNC interval). Past _max_idle_timeout (30s) with no
        // traffic either side may have silently discarded connection state
        // per RFC 9000 SS10.1, and the only way ensure_connected() would
        // otherwise find out is by eating a full 5s send_recv timeout on
        // the next use. Tracking this lets ensure_connected() reconnect
        // proactively instead.
        std::chrono::steady_clock::time_point last_activity{};
        // Consecutive connection failures since the last successful connect.
        // Drives an adaptive cooldown (see mark_connection_failed) so a
        // single transient blip doesn't cost the same 10s blackout as a
        // node that's actually, persistently down.
        int consecutive_failures = 0;
    };

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

    // Reset a node's connection and apply an adaptive cooldown after a
    // failed connect/send/recv. Cooldown starts at 1s and doubles per
    // consecutive failure, capped at 10s — a single blip recovers quickly,
    // while a node that keeps failing backs off to the old flat behaviour.
    // Caller must already hold mutex_ (this is called from within
    // already-locked retrieve/store paths, including from async worker
    // threads that only ever touch their own node_index's slot).
    void mark_connection_failed(size_t node_index);

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

    // Unlocked LIST_GROUPS probe (for internal parallel use)
    bool list_groups_on_node_impl(size_t node_index, std::vector<uint64_t>& groups);

    // 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, bounded by
    // `deadline` (defaults to none, i.e. the retry budget is only capped
    // by MAX_RETRIES as before).
    // The caller's unique_lock is released during retry sleeps so that other
    // client operations are not blocked by long backoff waits.
    // definitively_not_found, when provided, is set on a false return to
    // distinguish "enough nodes gave an explicit BLOCK_NOT_FOUND/local-miss
    // answer that fewer than k possible shards remain" (true — no readable
    // value can exist) from "k shards may still exist on unanswered nodes"
    // (false — this is a connectivity problem, not proof of non-existence).
    // Left null by callers that don't need the distinction (unchanged
    // behavior).
    bool retrieve_stripe(std::unique_lock<std::mutex>& lk,
                         uint64_t group_id, uint32_t stripe_index,
                         std::vector<uint8_t>& out,
                         std::chrono::steady_clock::time_point deadline =
                             (std::chrono::steady_clock::time_point::max)(),
                         bool* definitively_not_found = nullptr);

    // Single-attempt stripe retrieval (internal).
    // Sets should_retry=true when some blocks were found but < k.
    // `deadline` clamps this attempt's own network wait (normally 5s) to
    // whatever time remains, so a caller-supplied deadline is never
    // overrun by a single attempt.
    bool retrieve_stripe_once(uint64_t group_id, uint32_t stripe_index,
                               std::vector<uint8_t>& out, bool& should_retry,
                               size_t candidate_offset,
                               std::vector<bool>& confirmed_not_found,
                               std::chrono::steady_clock::time_point deadline =
                                   (std::chrono::steady_clock::time_point::max)());

    uint64_t generate_group_id();
};

} // namespace paritypp
