use reed_solomon_simd::{ReedSolomonEncoder, ReedSolomonDecoder};
use std::error::Error;
use rayon::prelude::*;

fn unzip<'a>(data: &'a Vec<&'a [u8]>, k: usize) -> Vec<Vec<&'a [u8]>> {
    (0..k)
        .map(|i| data.iter().skip(i).step_by(k).copied().collect())
        .collect()
}

pub fn encode(segment: &mut Vec<u8>, original_shards_count: usize, recovery_shards_count: usize) -> Result<Vec<Vec<u8>>, Box<dyn Error + Send + Sync>>{
    let length = segment.len();

    let w_e = 2*original_shards_count;

    if length%w_e != 0 {
        let target_size = ((length / w_e) + 1) * w_e;
        let padding_size = target_size - length;
        segment.extend(std::iter::repeat(0).take(padding_size));
    }

    let new_len = segment.len();
    let bytes_per_chunk =  new_len / w_e;
    let octet_pairs: Vec<&[u8]> = segment.chunks_exact(2).collect();

    let p = unzip(&octet_pairs, bytes_per_chunk);

    let mut original_shards: Vec<Vec<u8>> = vec![Vec::with_capacity(2 * bytes_per_chunk); original_shards_count];
    let mut recovery_shards: Vec<Vec<u8>> = vec![Vec::with_capacity(2 * bytes_per_chunk); recovery_shards_count];

    let mut encoder = ReedSolomonEncoder::new(original_shards_count, recovery_shards_count, 2)?;
    for i in 0..bytes_per_chunk{
        encoder.reset(original_shards_count, recovery_shards_count, 2)?;
        for j in 0..original_shards_count {
            let shard = p[i][j];
            encoder.add_original_shard(shard)?;
            original_shards[j].extend_from_slice(&shard);
        }

        let encoded = encoder.encode()?;
        for (recovery_idx, shard) in encoded.recovery_iter().enumerate() {
            recovery_shards[recovery_idx].extend_from_slice(shard);
        }
    }

    return Ok(original_shards.into_iter().chain(recovery_shards.into_iter()).collect());
}

pub fn decode(data: Vec<(Vec<u8>, usize)>, original_shards_count: usize, recovery_shards_count: usize) -> Result<Vec<u8>, Box<dyn Error + Send + Sync>>{

    let bytes_per_chunk =  data[0].0.len() / (2 * original_shards_count);
    let mut shards: Vec<Vec<u8>> = vec![Vec::with_capacity(2 * bytes_per_chunk); original_shards_count];

    let rows = data.len();
    let shard_len = data[0].0.len();
    let mut decoder = ReedSolomonDecoder::new(original_shards_count, recovery_shards_count, 2)?;

    for cols in (0..shard_len).step_by(2){

        decoder.reset(original_shards_count, recovery_shards_count, 2)?;

        for i in 0..rows{
            let pair = [data[i].0[cols], data[i].0[cols+1]];
            let index = data[i].1;

            if index < original_shards_count {
                decoder.add_original_shard(index, pair)?;
                shards[index].extend_from_slice(&pair);
            }else {
                decoder.add_recovery_shard(index - original_shards_count, pair)?;
            }
        }

        let decoded = decoder.decode()?;
        for (idx, shard) in decoded.restored_original_iter() {
            shards[idx].extend_from_slice(shard);
        }
    }

    return Ok(shards.iter().flat_map(|v| v.iter().copied()).collect());

}

pub fn multi_encoder(mut segments: Vec<Vec<u8>>, original_shards_count: usize, recovery_shards_count: usize) -> Result<Vec<Vec<Vec<u8>>>, Box<dyn Error + Send + Sync>> {
    segments
        .par_iter_mut()
        .map(|segment| encode(segment, original_shards_count, recovery_shards_count))
        .collect()
}
