use crate::reed_solomon::{
    DecoderResult, EncoderResult, Error, Plan, RecoveryDecoderResult,
    engine::DefaultEngine,
    rate::{DefaultRate, DefaultRateDecoder, DefaultRateEncoder, Rate, RateDecoder, RateEncoder},
};

/// Reed-Solomon encoder using [`DefaultEngine`] and [`DefaultRate`].
///
/// [`DefaultEngine`]: crate::reed_solomon::engine::DefaultEngine
pub struct Encoder(DefaultRateEncoder<DefaultEngine>);

impl Encoder {
    /// Adds one original shard to the encoder.
    ///
    /// Original shards have indexes `0..original_count` corresponding to the order in which they
    /// are added, and these same indexes must be used when decoding.
    ///
    /// See [basic usage](crate::reed_solomon#basic-usage) for an example.
    pub fn add_original_shard<T: AsRef<[u8]>>(&mut self, original_shard: T) -> Result<(), Error> {
        self.0.add_original_shard(original_shard)
    }

    /// Encodes the added original shards, returning [`EncoderResult`] which contains the
    /// generated recovery shards.
    ///
    /// When the returned [`EncoderResult`] is dropped the encoder is automatically [`reset`] and
    /// ready for a new round of encoding. A result that is leaked instead of dropped (for example
    /// with [`core::mem::forget`]) leaves the encoder reporting a complete round, and the next
    /// [`encode`](Encoder::encode) would transform the previous round's buffers again. Call
    /// [`reset`] before reusing an encoder whose result was leaked.
    ///
    /// See [basic usage](crate::reed_solomon#basic-usage) for an example.
    ///
    /// [`reset`]: Encoder::reset
    pub fn encode(&mut self) -> Result<EncoderResult<'_>, Error> {
        self.0.encode()
    }

    /// Creates a new encoder with the given configuration and allocates the required working
    /// space.
    ///
    /// See [working space](crate::reed_solomon#working-space) for its size and the configurations
    /// a coder accepts, and [basic usage](crate::reed_solomon#basic-usage) for an example.
    pub fn new(
        original_count: usize,
        recovery_count: usize,
        shard_bytes: usize,
    ) -> Result<Self, Error> {
        Ok(Self(DefaultRateEncoder::new(
            original_count,
            recovery_count,
            shard_bytes,
            DefaultEngine::new(),
            None,
        )?))
    }

    /// Resets the encoder to the given configuration.
    ///
    /// - Added original shards are forgotten.
    /// - Existing working space is reused if it is large enough, or reallocated otherwise.
    pub fn reset(
        &mut self,
        original_count: usize,
        recovery_count: usize,
        shard_bytes: usize,
    ) -> Result<(), Error> {
        self.0.reset(original_count, recovery_count, shard_bytes)
    }

    /// Returns `true` if the given `original_count` / `recovery_count` combination is supported.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use commonware_cryptography::reed_solomon::Encoder;
    ///
    /// assert_eq!(Encoder::supports(60_000, 4_000), true);
    /// assert_eq!(Encoder::supports(60_000, 5_000), false);
    /// ```
    pub fn supports(original_count: usize, recovery_count: usize) -> bool {
        DefaultRate::<DefaultEngine>::supports(original_count, recovery_count)
    }
}

/// Reed-Solomon decoder using [`DefaultEngine`] and [`DefaultRate`].
///
/// [`DefaultEngine`]: crate::reed_solomon::engine::DefaultEngine
pub struct Decoder(DefaultRateDecoder<DefaultEngine>);

impl Decoder {
    /// Adds one original shard to the decoder.
    ///
    /// - Shards can be added in any order.
    /// - The index must be the same one used in encoding.
    ///
    /// See [basic usage](crate::reed_solomon#basic-usage) for an example.
    pub fn add_original_shard<T: AsRef<[u8]>>(
        &mut self,
        index: usize,
        original_shard: T,
    ) -> Result<(), Error> {
        self.0.add_original_shard(index, original_shard)
    }

    /// Adds one recovery shard to the decoder.
    ///
    /// - Shards can be added in any order.
    /// - The index must be the same one used in encoding.
    ///
    /// See [basic usage](crate::reed_solomon#basic-usage) for an example.
    pub fn add_recovery_shard<T: AsRef<[u8]>>(
        &mut self,
        index: usize,
        recovery_shard: T,
    ) -> Result<(), Error> {
        self.0.add_recovery_shard(index, recovery_shard)
    }

    /// Decodes the added shards, returning `Some(`[`DecoderResult`]`)` with the restored original
    /// shards, or `None` if every original shard was already provided (nothing to reconstruct).
    ///
    /// When the returned [`DecoderResult`] is dropped the decoder is automatically [`reset`] and
    /// ready for a new round of decoding. When every original shard was supplied, this returns
    /// `None` and clears the received shards immediately. A result that is leaked instead of
    /// dropped (for example with [`core::mem::forget`]) leaves the decoder reporting a complete
    /// round, and the next decode would transform the previous round's buffers again. Call
    /// [`reset`] before reusing a decoder whose result was leaked.
    ///
    /// See [basic usage](crate::reed_solomon#basic-usage) for an example.
    ///
    /// [`reset`]: Decoder::reset
    pub fn decode(&mut self) -> Result<Option<DecoderResult<'_>>, Error> {
        self.0.decode(false)
    }

    /// Like [`decode`](Decoder::decode), but reads erasure coefficients from `plan` instead of
    /// computing them.
    ///
    /// `plan` must match this decoder's shard counts and received shard indices. Shard byte
    /// length is not part of the plan. A mismatched plan returns
    /// [`Error::PlanMismatch`] and leaves the added shards in place, so the decoder can
    /// retry with the matching plan or [`decode`](Decoder::decode).
    pub fn decode_with_plan(&mut self, plan: &Plan) -> Result<Option<DecoderResult<'_>>, Error> {
        self.0.decode_with_plan(false, plan)
    }

    /// Like [`decode`](Decoder::decode), but also reconstructs the missing recovery shards,
    /// returning `Some(`[`RecoveryDecoderResult`]`)` that additionally exposes them via
    /// [`RecoveryDecoderResult::recovery`] / [`recovery_iter`]. This adds a multiplication of each
    /// missing recovery shard by a field element, so prefer [`decode`](Decoder::decode) when only
    /// the original data is needed.
    ///
    /// [`recovery_iter`]: RecoveryDecoderResult::recovery_iter
    pub fn decode_with_recovery(&mut self) -> Result<Option<RecoveryDecoderResult<'_>>, Error> {
        Ok(self.0.decode(true)?.map(RecoveryDecoderResult::new))
    }

    /// Like [`decode_with_recovery`](Decoder::decode_with_recovery), but reads erasure
    /// coefficients from `plan` as [`decode_with_plan`](Decoder::decode_with_plan) does.
    pub fn decode_with_recovery_plan(
        &mut self,
        plan: &Plan,
    ) -> Result<Option<RecoveryDecoderResult<'_>>, Error> {
        Ok(self
            .0
            .decode_with_plan(true, plan)?
            .map(RecoveryDecoderResult::new))
    }

    /// Creates a new decoder with the given configuration and allocates the required working
    /// space.
    ///
    /// See [working space](crate::reed_solomon#working-space) for its size and the configurations
    /// a coder accepts, and [basic usage](crate::reed_solomon#basic-usage) for an example.
    pub fn new(
        original_count: usize,
        recovery_count: usize,
        shard_bytes: usize,
    ) -> Result<Self, Error> {
        Ok(Self(DefaultRateDecoder::new(
            original_count,
            recovery_count,
            shard_bytes,
            DefaultEngine::new(),
            None,
        )?))
    }

    /// Resets the decoder to the given configuration.
    ///
    /// - Added shards are forgotten.
    /// - Existing working space is reused if it is large enough, or reallocated otherwise.
    pub fn reset(
        &mut self,
        original_count: usize,
        recovery_count: usize,
        shard_bytes: usize,
    ) -> Result<(), Error> {
        self.0.reset(original_count, recovery_count, shard_bytes)
    }

    /// Returns `true` if the given `original_count` / `recovery_count` combination is supported.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use commonware_cryptography::reed_solomon::Decoder;
    ///
    /// assert_eq!(Decoder::supports(60_000, 4_000), true);
    /// assert_eq!(Decoder::supports(60_000, 5_000), false);
    /// ```
    pub fn supports(original_count: usize, recovery_count: usize) -> bool {
        DefaultRate::<DefaultEngine>::supports(original_count, recovery_count)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::reed_solomon::test_util;
    use fixedbitset::FixedBitSet;
    use std::collections::BTreeMap;

    /// Encodes `original_count` 1024-byte originals, checks the recovery hash, decodes from the
    /// given original and recovery indices, and checks every original not given to the decoder.
    fn roundtrip(
        encoder: &mut Encoder,
        decoder: &mut Decoder,
        original_count: usize,
        recovery_hash: &str,
        decoder_original: &[usize],
        decoder_recovery: &[usize],
        seed: u8,
    ) {
        let original = test_util::generate_original(original_count, 1024, seed);

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

        let result = encoder.encode().unwrap();
        let recovery: Vec<_> = result.recovery_iter().collect();

        test_util::assert_hash(&recovery, recovery_hash);

        let mut original_received = FixedBitSet::with_capacity(original_count);

        for i in decoder_original {
            decoder.add_original_shard(*i, &original[*i]).unwrap();
            original_received.set(*i, true);
        }

        for i in decoder_recovery {
            decoder.add_recovery_shard(*i, recovery[*i]).unwrap();
        }

        let decoded = decoder.decode();
        let restored: BTreeMap<_, _> = match &decoded {
            Ok(Some(result)) => result.original_iter().collect(),
            Ok(None) => BTreeMap::new(),
            Err(e) => panic!("decode failed: {e:?}"),
        };

        for i in 0..original_count {
            if !original_received[i] {
                assert_eq!(restored[&i], original[i]);
            }
        }
    }

    #[test]
    fn maximum_decoder_indices() {
        for (original_count, recovery_count) in [(3, 2), (2, 3)] {
            let original = test_util::generate_original(original_count, 2, 123);
            let mut encoder = Encoder::new(original_count, recovery_count, 2).unwrap();
            for shard in &original {
                encoder.add_original_shard(shard).unwrap();
            }
            let encoded = encoder.encode().unwrap();
            let mut decoder = Decoder::new(original_count, recovery_count, 2).unwrap();
            assert_eq!(
                decoder.add_original_shard(usize::MAX, [0; 2]),
                Err(Error::InvalidOriginalShardIndex {
                    original_count,
                    index: usize::MAX
                })
            );
            assert_eq!(
                decoder.add_recovery_shard(usize::MAX, [0; 2]),
                Err(Error::InvalidRecoveryShardIndex {
                    recovery_count,
                    index: usize::MAX
                })
            );
            for (index, shard) in original.iter().enumerate().skip(1) {
                decoder.add_original_shard(index, shard).unwrap();
            }
            decoder
                .add_recovery_shard(0, encoded.recovery(0).unwrap())
                .unwrap();
            let decoded = decoder.decode_with_recovery().unwrap().unwrap();
            assert_eq!(decoded.original(0), Some(original[0].as_slice()));
            assert_eq!(decoded.original(usize::MAX), None);
            assert_eq!(decoded.recovery(usize::MAX), None);
        }
    }

    #[test]
    fn unrepresentable_working_space() {
        for (original_count, recovery_count) in [(3, 2), (2, 3)] {
            for shard_bytes in [usize::MAX - 1, 1 << (usize::BITS - 2)] {
                assert!(matches!(
                    Encoder::new(original_count, recovery_count, shard_bytes),
                    Err(Error::InvalidShardSize { .. })
                ));
                assert!(matches!(
                    Decoder::new(original_count, recovery_count, shard_bytes),
                    Err(Error::InvalidShardSize { .. })
                ));
            }
        }
    }

    #[test]
    fn roundtrip_two_rounds_reset_low_to_high() {
        let mut encoder = Encoder::new(2, 3, 1024).unwrap();
        let mut decoder = Decoder::new(2, 3, 1024).unwrap();

        roundtrip(
            &mut encoder,
            &mut decoder,
            2,
            test_util::LOW_2_3,
            &[],
            &[0, 1],
            123,
        );

        encoder.reset(3, 2, 1024).unwrap();
        decoder.reset(3, 2, 1024).unwrap();

        roundtrip(
            &mut encoder,
            &mut decoder,
            3,
            test_util::HIGH_3_2,
            &[1],
            &[0, 1],
            132,
        );
    }

    #[test]
    fn failed_encoder_reset_preserves_state() {
        let original = test_util::generate_original(2, 1024, 123);
        let mut encoder = Encoder::new(2, 3, 1024).unwrap();
        encoder.add_original_shard(&original[0]).unwrap();

        for (original_count, recovery_count) in [(2, 3), (3, 2)] {
            for shard_bytes in [0, 3, usize::MAX - 1] {
                assert_eq!(
                    encoder.reset(original_count, recovery_count, shard_bytes),
                    Err(Error::InvalidShardSize { shard_bytes })
                );
            }
        }

        for shard in original.iter().skip(1) {
            encoder.add_original_shard(shard).unwrap();
        }
        let result = encoder.encode().unwrap();
        let recovery: Vec<_> = result.recovery_iter().collect();

        test_util::assert_hash(&recovery, test_util::LOW_2_3);
    }

    #[test]
    fn failed_decoder_reset_preserves_state() {
        let original = test_util::generate_original(2, 1024, 123);
        let mut encoder = Encoder::new(2, 3, 1024).unwrap();
        for shard in &original {
            encoder.add_original_shard(shard).unwrap();
        }
        let result = encoder.encode().unwrap();
        let recovery: Vec<_> = result.recovery_iter().map(<[u8]>::to_vec).collect();

        let mut decoder = Decoder::new(2, 3, 1024).unwrap();
        decoder.add_recovery_shard(0, &recovery[0]).unwrap();

        for (original_count, recovery_count) in [(2, 3), (3, 2)] {
            for shard_bytes in [0, 3, usize::MAX - 1] {
                assert_eq!(
                    decoder.reset(original_count, recovery_count, shard_bytes),
                    Err(Error::InvalidShardSize { shard_bytes })
                );
            }
        }

        decoder.add_recovery_shard(1, &recovery[1]).unwrap();
        let decoded = decoder.decode();
        let restored: BTreeMap<_, _> = match &decoded {
            Ok(Some(result)) => result.original_iter().collect(),
            Ok(None) => BTreeMap::new(),
            Err(e) => panic!("decode failed: {e:?}"),
        };

        assert_eq!(restored[&0], original[0]);
        assert_eq!(restored[&1], original[1]);
    }

    #[test]
    fn supports() {
        assert!(Encoder::supports(4096, 61440));
        assert!(Encoder::supports(61440, 4096));

        assert!(Decoder::supports(4096, 61440));
        assert!(Decoder::supports(61440, 4096));
    }
}