use crate::rate::EncoderWork;

// ======================================================================
// EncoderResult - PUBLIC

/// Result of encoding. Contains the generated recovery shards.
///
/// This struct is created by [`ReedSolomonEncoder::encode`]
/// and [`RateEncoder::encode`].
///
/// [`RateEncoder::encode`]: crate::rate::RateEncoder::encode
/// [`ReedSolomonEncoder::encode`]: crate::ReedSolomonEncoder::encode
pub struct EncoderResult<'a> {
    work: &'a mut EncoderWork,
}

impl<'a> EncoderResult<'a> {
    /// Returns recovery shard with given `index`
    /// or `None` if `index >= recovery_count`.
    ///
    /// Recovery shards have indexes `0..recovery_count`
    /// and these same indexes must be used when decoding.
    pub fn recovery(&self, index: usize) -> Option<&[u8]> {
        self.work.recovery(index)
    }

    /// Returns iterator over all recovery shards ordered by their indexes.
    ///
    /// Recovery shards have indexes `0..recovery_count`
    /// and these same indexes must be used when decoding.
    pub fn recovery_iter(&self) -> Recovery {
        Recovery::new(self.work)
    }
}

// ======================================================================
// EncoderResult - CRATE

impl<'a> EncoderResult<'a> {
    pub(crate) fn new(work: &'a mut EncoderWork) -> Self {
        Self { work }
    }
}

// ======================================================================
// EncoderResult - IMPL DROP

impl<'a> Drop for EncoderResult<'a> {
    fn drop(&mut self) {
        self.work.reset_received();
    }
}

// ======================================================================
// Recovery - PUBLIC

/// Iterator over generated recovery shards.
///
/// This struct is created by [`EncoderResult::recovery_iter`].
pub struct Recovery<'a> {
    ended: bool,
    next_index: usize,
    work: &'a EncoderWork,
}

// ======================================================================
// Recovery - IMPL Iterator

impl<'a> Iterator for Recovery<'a> {
    type Item = &'a [u8];
    fn next(&mut self) -> Option<&'a [u8]> {
        if self.ended {
            None
        } else if let Some(next) = self.work.recovery(self.next_index) {
            self.next_index += 1;
            Some(next)
        } else {
            self.ended = true;
            None
        }
    }
}

// ======================================================================
// Recovery - CRATE

impl<'a> Recovery<'a> {
    pub(crate) fn new(work: &'a EncoderWork) -> Self {
        Self {
            ended: false,
            next_index: 0,
            work,
        }
    }
}

// ======================================================================
// TESTS

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{test_util, ReedSolomonEncoder};

    #[test]
    // EncoderResult::recovery
    // EncoderResult::recovery_iter
    // Recovery
    fn encoder_result() {
        let original = test_util::generate_original(2, 1024, 123);
        let mut encoder = ReedSolomonEncoder::new(2, 3, 1024).unwrap();

        for original in &original {
            encoder.add_original_shard(original).unwrap();
        }

        let result: EncoderResult = encoder.encode().unwrap();

        let mut all = Vec::new();
        all.push(result.recovery(0).unwrap());
        all.push(result.recovery(1).unwrap());
        all.push(result.recovery(2).unwrap());
        assert!(result.recovery(3).is_none());
        test_util::assert_hash(all, test_util::LOW_2_3);

        let mut iter: Recovery = result.recovery_iter();
        let mut all = Vec::new();
        all.push(iter.next().unwrap());
        all.push(iter.next().unwrap());
        all.push(iter.next().unwrap());
        assert!(iter.next().is_none());
        test_util::assert_hash(all, test_util::LOW_2_3);
    }
}
