#include "parity.h"
#include "galois.h"
#include "matrix.h"
#include <cstring>
#include <algorithm>
#include <stdexcept>

namespace paritypp {

parity::parity(size_t data_blocks, size_t parity_blocks)
    : k_(data_blocks), m_(parity_blocks) {
    if (k_ == 0) throw std::invalid_argument("Need at least 1 data block");
    if (m_ == 0) throw std::invalid_argument("Need at least 1 parity block");
    if (k_ + m_ > 255) throw std::invalid_argument("Total blocks cannot exceed 255 for GF(2^8)");

    galois::init();
    build_encoding_matrix();
}

void parity::build_encoding_matrix() {
    // Use a Cauchy-based encoding matrix for optimal MDS properties
    // Top k rows = identity (data passes through)
    // Bottom m rows = Cauchy parity matrix
    size_t n = k_ + m_;
    encode_matrix_.resize(n * k_);

    // Identity for data rows
    for (size_t r = 0; r < k_; r++) {
        for (size_t c = 0; c < k_; c++) {
            encode_matrix_[r * k_ + c] = (r == c) ? 1 : 0;
        }
    }

    // Cauchy submatrix for parity rows
    // M[i][j] = 1 / (x_i XOR y_j) where x_i and y_j are distinct elements
    for (size_t r = 0; r < m_; r++) {
        for (size_t c = 0; c < k_; c++) {
            uint8_t x = static_cast<uint8_t>(k_ + r); // parity indices
            uint8_t y = static_cast<uint8_t>(c);       // data indices
            encode_matrix_[(k_ + r) * k_ + c] = galois::inv(galois::add(x, y));
        }
    }
}

std::vector<std::vector<uint8_t>> parity::encode(
    const std::vector<std::vector<uint8_t>>& data_in) const {

    if (data_in.size() != k_) {
        throw std::invalid_argument("Expected " + std::to_string(k_) + " data blocks");
    }

    size_t block_size = data_in[0].size();
    for (const auto& block : data_in) {
        if (block.size() != block_size) {
            throw std::invalid_argument("All data blocks must have the same size");
        }
    }

    // Compute parity blocks using the parity rows of the encoding matrix
    std::vector<std::vector<uint8_t>> parity_out(m_, std::vector<uint8_t>(block_size, 0));

    for (size_t p = 0; p < m_; p++) {
        for (size_t d = 0; d < k_; d++) {
            uint8_t coeff = encode_matrix_[(k_ + p) * k_ + d];
            galois::addmul_region(parity_out[p].data(), data_in[d].data(), coeff, block_size);
        }
    }

    return parity_out;
}

std::vector<std::vector<uint8_t>> parity::decode(
    const std::vector<std::vector<uint8_t>>& surviving_blocks,
    const std::vector<size_t>& surviving_indices,
    size_t block_size) const {

    if (surviving_blocks.size() < k_) {
        throw std::invalid_argument("Need at least " + std::to_string(k_) +
                                    " surviving blocks to decode");
    }
    if (surviving_blocks.size() != surviving_indices.size()) {
        throw std::invalid_argument("Block count and index count mismatch");
    }

    // Take exactly k surviving blocks
    size_t use_count = k_;

    // Build the submatrix from the encoding matrix using surviving row indices
    matrix sub(use_count, k_);
    for (size_t r = 0; r < use_count; r++) {
        size_t idx = surviving_indices[r];
        if (idx >= k_ + m_) {
            throw std::invalid_argument("Invalid block index");
        }
        for (size_t c = 0; c < k_; c++) {
            sub.at(r, c) = encode_matrix_[idx * k_ + c];
        }
    }

    // Invert the submatrix
    matrix inv = sub.invert();

    // Multiply inverse by surviving data to get original data blocks
    std::vector<std::vector<uint8_t>> decoded(k_, std::vector<uint8_t>(block_size, 0));

    for (size_t d = 0; d < k_; d++) {
        for (size_t s = 0; s < use_count; s++) {
            uint8_t coeff = inv.at(d, s);
            galois::addmul_region(decoded[d].data(), surviving_blocks[s].data(),
                                  coeff, block_size);
        }
    }

    return decoded;
}

bool parity::verify(const std::vector<std::vector<uint8_t>>& data_blocks,
                    const std::vector<std::vector<uint8_t>>& parity_blocks) const {
    if (data_blocks.size() != k_ || parity_blocks.size() != m_) {
        return false;
    }

    auto computed = encode(data_blocks);
    for (size_t p = 0; p < m_; p++) {
        if (computed[p] != parity_blocks[p]) {
            return false;
        }
    }
    return true;
}

std::vector<std::vector<uint8_t>> parity::split(const uint8_t* data,
                                                  size_t data_len,
                                                  size_t k) {
    size_t block_size = (data_len + k - 1) / k;
    std::vector<std::vector<uint8_t>> blocks(k, std::vector<uint8_t>(block_size, 0));

    for (size_t i = 0; i < k; i++) {
        size_t offset = i * block_size;
        size_t copy_len = std::min(block_size, data_len > offset ? data_len - offset : 0);
        if (copy_len > 0) {
            std::memcpy(blocks[i].data(), data + offset, copy_len);
        }
    }
    return blocks;
}

std::vector<uint8_t> parity::join(const std::vector<std::vector<uint8_t>>& blocks,
                                   size_t original_size) {
    size_t total = 0;
    for (const auto& b : blocks) total += b.size();
    if (original_size > 0 && original_size < total) total = original_size;

    std::vector<uint8_t> result(total);
    size_t offset = 0;
    for (const auto& b : blocks) {
        size_t copy_len = std::min(b.size(), total - offset);
        if (copy_len > 0) {
            std::memcpy(result.data() + offset, b.data(), copy_len);
        }
        offset += copy_len;
        if (offset >= total) break;
    }
    return result;
}

} // namespace paritypp
