#pragma once

#include <cstdint>
#include <cstddef>
#include <vector>
#include <string>

namespace paritypp {

// Reed-Solomon erasure coding with flexible parity block count
//
// k = number of data blocks
// m = number of parity blocks (flexible, 1..255-k)
// Can tolerate loss of any m blocks (data or parity)
class parity {
public:
    // Create encoder/decoder for k data blocks and m parity blocks
    parity(size_t data_blocks, size_t parity_blocks);

    size_t data_block_count() const { return k_; }
    size_t parity_block_count() const { return m_; }
    size_t total_block_count() const { return k_ + m_; }

    // Encode: compute m parity blocks from k data blocks
    // data_in: vector of k data blocks (each of block_size bytes)
    // Returns: vector of m parity blocks
    std::vector<std::vector<uint8_t>> encode(
        const std::vector<std::vector<uint8_t>>& data_in) const;

    // Decode: reconstruct missing blocks from surviving blocks
    // surviving_blocks: at least k blocks (data or parity)
    // surviving_indices: which block indices survived (0..k-1 = data, k..k+m-1 = parity)
    // block_size: size of each block in bytes
    // Returns: reconstructed k data blocks
    std::vector<std::vector<uint8_t>> decode(
        const std::vector<std::vector<uint8_t>>& surviving_blocks,
        const std::vector<size_t>& surviving_indices,
        size_t block_size) const;

    // Verify: check if parity blocks are consistent with data blocks
    bool verify(const std::vector<std::vector<uint8_t>>& data_blocks,
                const std::vector<std::vector<uint8_t>>& parity_blocks) const;

    // Split a data buffer into k equal-size blocks (padding with zeroes if needed)
    static std::vector<std::vector<uint8_t>> split(const uint8_t* data,
                                                    size_t data_len,
                                                    size_t k);

    // Join k data blocks back into a single buffer
    // original_size: size before padding (to strip trailing zeroes)
    static std::vector<uint8_t> join(const std::vector<std::vector<uint8_t>>& blocks,
                                     size_t original_size = 0);

private:
    size_t k_; // data blocks
    size_t m_; // parity blocks

    // The encoding matrix: (k+m) x k, top k rows = identity, bottom m rows = parity coefficients
    std::vector<uint8_t> encode_matrix_;

    void build_encoding_matrix();
};

} // namespace paritypp
