//! ProverResult
use std::io;

use crate::chain::{Base16DecodedBytes, Base16EncodedBytes};
use crate::serialization::{
    sigma_byte_reader::SigmaByteRead, sigma_byte_writer::SigmaByteWrite, SerializationError,
    SigmaSerializable,
};
#[cfg(feature = "json")]
use serde::{Deserialize, Serialize};

use super::ContextExtension;

/// Serialized proof generated by ['Prover']
#[cfg_attr(feature = "json", derive(Serialize, Deserialize))]
#[cfg_attr(
    feature = "json",
    serde(into = "Base16EncodedBytes", try_from = "Base16DecodedBytes")
)]
#[derive(PartialEq, Eq, Hash, Debug, Clone)]
pub enum ProofBytes {
    /// Empty proof
    Empty,
    /// Non-empty proof
    Some(Vec<u8>),
}

impl Into<Base16EncodedBytes> for ProofBytes {
    fn into(self) -> Base16EncodedBytes {
        match self {
            ProofBytes::Empty => Base16EncodedBytes::new(vec![].as_slice()),
            ProofBytes::Some(bytes) => Base16EncodedBytes::new(&bytes),
        }
    }
}

impl From<Base16DecodedBytes> for ProofBytes {
    fn from(v: Base16DecodedBytes) -> Self {
        if v.0.is_empty() {
            ProofBytes::Empty
        } else {
            ProofBytes::Some(v.0)
        }
    }
}

impl SigmaSerializable for ProofBytes {
    fn sigma_serialize<W: SigmaByteWrite>(&self, w: &mut W) -> Result<(), io::Error> {
        match self {
            ProofBytes::Empty => w.put_u16(0)?,
            ProofBytes::Some(bytes) => {
                w.put_u16(bytes.len() as u16)?;
                w.write_all(&bytes)?;
            }
        }
        Ok(())
    }

    fn sigma_parse<R: SigmaByteRead>(r: &mut R) -> Result<Self, SerializationError> {
        let proof_len = r.get_u16()?;
        Ok(if proof_len == 0 {
            ProofBytes::Empty
        } else {
            let mut bytes = vec![0; proof_len as usize];
            r.read_exact(&mut bytes)?;
            ProofBytes::Some(bytes)
        })
    }
}

/// Proof of correctness of tx spending
#[derive(Debug, PartialEq, Eq, Clone)]
#[cfg_attr(feature = "json", derive(Serialize, Deserialize))]
pub struct ProverResult {
    /// proof that satisfies final sigma proposition
    #[cfg_attr(feature = "json", serde(rename = "proofBytes"))]
    pub proof: ProofBytes,
    /// user-defined variables to be put into context
    #[cfg_attr(feature = "json", serde(rename = "extension"))]
    pub extension: ContextExtension,
}

impl SigmaSerializable for ProverResult {
    fn sigma_serialize<W: SigmaByteWrite>(&self, w: &mut W) -> Result<(), io::Error> {
        self.proof.sigma_serialize(w)?;
        self.extension.sigma_serialize(w)?;
        Ok(())
    }
    fn sigma_parse<R: SigmaByteRead>(r: &mut R) -> Result<Self, SerializationError> {
        let proof = ProofBytes::sigma_parse(r)?;
        let extension = ContextExtension::sigma_parse(r)?;
        Ok(ProverResult { proof, extension })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::serialization::sigma_serialize_roundtrip;
    use proptest::{collection::vec, prelude::*};

    impl Arbitrary for ProofBytes {
        type Parameters = ();
        type Strategy = BoxedStrategy<Self>;

        fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
            prop_oneof![
                Just(ProofBytes::Empty),
                (vec(any::<u8>(), 32..100))
                    .prop_map(ProofBytes::Some)
                    .boxed()
            ]
            .boxed()
        }
    }

    impl Arbitrary for ProverResult {
        type Parameters = ();
        type Strategy = BoxedStrategy<Self>;

        fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
            (any::<ProofBytes>(), any::<ContextExtension>())
                .prop_map(|(proof, extension)| Self { proof, extension })
                .boxed()
        }
    }
    proptest! {

        #[test]
        fn ser_roundtrip(v in any::<ProverResult>()) {
            prop_assert_eq![sigma_serialize_roundtrip(&v), v];
        }
    }
}
