//! Basic types used by the v1 client puzzle

use crate::pk::HsBlindId;
use crate::pow::v1::err::SolutionErrorV1;
use rand::prelude::*;
use serde::{Deserialize, Serialize};
use tor_bytes::{EncodeResult, Readable, Reader, Writeable, Writer};

/// Effort setting, a u32 value with linear scale
///
/// The numerical value is roughly the expected number of times we will
/// need to invoke the underlying solver (Equi-X) for the v1 proof-of-work
/// protocol to find a solution.
#[derive(derive_more::AsRef, derive_more::From, derive_more::Into)] //
#[derive(Serialize, Deserialize)] //
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq, Ord, PartialOrd)]
pub struct Effort(u32);

impl Effort {
    /// Const constructor for `Effort` from any u32
    ///
    /// `Effort` can also be constructed using into() conversions,
    /// but those are not available in const expressions.
    ///
    pub const fn new(value: u32) -> Self {
        Self(value)
    }

    /// Const constructor for zero effort
    ///
    /// Zero effort is used in suggested effort to mean the puzzle is
    /// available but shouldn't be used on a first connection attempt.
    ///
    /// If zero is actually used for solve and verify, it will have equivalent
    /// performance to an effort of 1.
    ///
    pub const fn zero() -> Self {
        Self(0)
    }

    /// Multiply this effort by an integer of the same size, saturating on overflow
    pub const fn saturating_mul_u32(&self, rhs: u32) -> Self {
        Self(self.0.saturating_mul(rhs))
    }

    /// Multiply this effort by a float, with saturating cast back to Effort
    pub fn saturating_mul_f32(&self, rhs: f32) -> Self {
        // 'as' cast from float to int saturates, and NaN becomes 0.
        Self((self.0 as f32 * rhs) as u32)
    }
}

/// Length of the random seed generated by servers and included in HsDir
pub const SEED_LEN: usize = 32;

/// The random portion of a challenge, distributed through HsDir
#[derive(
    derive_more::AsRef, derive_more::From, Deserialize, Serialize, Debug, Clone, Eq, PartialEq,
)]
pub struct Seed([u8; SEED_LEN]);

impl Seed {
    /// Generate a new, random seed.
    ///
    /// If old_seed is given, avoid generating a seed that shares a seed head.
    pub fn new<R: Rng + CryptoRng>(rng: &mut R, old_seed: Option<&Seed>) -> Self {
        let mut new = Seed(rng.random());
        while Some(new.head()) == old_seed.as_ref().map(|seed| seed.head()) {
            new = Seed(rng.random());
        }
        new
    }

    /// Make a new [`SeedHead`] from a prefix of this seed
    pub fn head(&self) -> SeedHead {
        SeedHead(
            self.0[..SEED_HEAD_LEN]
                .try_into()
                .expect("slice length correct"),
        )
    }
}

impl std::fmt::Display for Seed {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "{}", hex::encode(self.0))
    }
}

/// Error when converting a string to a [`Seed`].
#[non_exhaustive]
pub enum ParseSeedError {
    /// The string was not a valid hex string.
    InvalidHex(hex::FromHexError),
    /// The string was not the correct length for a [`Seed`].
    InvalidLength,
}

impl std::str::FromStr for Seed {
    type Err = ParseSeedError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let bytes = hex::decode(s).map_err(ParseSeedError::InvalidHex)?;
        Ok(Seed(
            bytes
                .try_into()
                .map_err(|_| ParseSeedError::InvalidLength)?,
        ))
    }
}

/// Length of a seed prefix used to identify the entire seed
pub const SEED_HEAD_LEN: usize = 4;

/// A short seed prefix used in solutions to reference the complete seed
#[derive(derive_more::AsRef, derive_more::From, Debug, Clone, Copy, Eq, PartialEq, Hash)]
pub struct SeedHead([u8; SEED_HEAD_LEN]);

/// Length of the nonce value generated by clients and included in the solution
pub const NONCE_LEN: usize = 16;

/// Generated randomly by solvers and included in the solution
#[derive(derive_more::AsRef, derive_more::From, Debug, Clone, Eq, PartialEq)]
pub struct Nonce([u8; NONCE_LEN]);

/// Generate [`Readable`] and [`Writeable`] implementations for trivial newtype wrappers
macro_rules! impl_readable_writeable_newtype {
    ($t:ty) => {
        impl Readable for $t {
            fn take_from(b: &mut Reader<'_>) -> tor_bytes::Result<Self> {
                Ok(Self(Readable::take_from(b)?))
            }
        }
        impl Writeable for $t {
            fn write_onto<B: Writer + ?Sized>(&self, b: &mut B) -> EncodeResult<()> {
                self.0.write_onto(b)
            }
        }
    };
}

impl_readable_writeable_newtype!(Effort);
impl_readable_writeable_newtype!(Seed);
impl_readable_writeable_newtype!(SeedHead);
impl_readable_writeable_newtype!(Nonce);

/// One instance of this proof-of-work puzzle
///
/// Identified uniquely by the combination of onion service blinded Id key
/// plus a rotating seed chosen by the service.
#[derive(Debug, Clone)]
pub struct Instance {
    /// Blinded public Id key, binding this puzzle to a specific onion service
    service: HsBlindId,
    /// Seed value distributed in the HsDir by that service
    seed: Seed,
}

impl Instance {
    /// A new puzzle instance, wrapping a service Id and service-chosen seed
    pub fn new(service: HsBlindId, seed: Seed) -> Self {
        Self { service, seed }
    }

    /// Get the [`HsBlindId`] identifying the service this puzzle is for.
    pub fn service(&self) -> &HsBlindId {
        &self.service
    }

    /// Get the rotating random [`Seed`] used in this puzzle instance.
    pub fn seed(&self) -> &Seed {
        &self.seed
    }
}

/// One potential solution to some puzzle [`Instance`]
///
/// The existence of a [`Solution`] guarantees that the solution is well formed
/// (for example, the correct length, the correct order  in the Equi-X solution)
/// but it makes no guarantee to actually solve any specific puzzle instance.
#[derive(Debug, Clone, Eq, PartialEq)]
pub struct Solution {
    /// Arbitrary value chosen by the solver to reach a valid solution
    ///
    /// Services are responsible for remembering used values to prevent replay.
    nonce: Nonce,

    /// The effort chosen by the client
    ///
    /// This is validated against the actual effort spent by the client using
    /// a combination of two checks:
    ///
    /// - We can ensure the effort value here was chosen prior to successfully
    ///   solving the Equi-X puzzle just by verifying the Equi-X proof.
    ///   Effort values are part of the [`crate::pow::v1::challenge::Challenge`]
    ///   string the puzzle is constructed around.
    ///
    /// - We can ensure, on average, that the proper proportion of Equi-X
    ///   solutions have been discarded. The proof and challenge are hashed,
    ///   and the resulting digest is effectively a random variable that must
    ///   fit within a range inversely proportional to the effort. This test
    ///   happens in [`crate::pow::v1::challenge::Challenge::check_effort`].
    effort: Effort,

    /// Prefix of the [`Seed`] used in this puzzle Instance
    ///
    /// A service will normally have two active [`Seed`] values at once.
    /// This prefix is sufficient to distinguish between them. (Services
    /// skip seeds which would have the same prefix as the last seed.)
    seed_head: SeedHead,

    /// Equi-X solution which claims to prove the above effort choice
    proof: equix::Solution,
}

impl Solution {
    /// Construct a new Solution around a well-formed [`equix::Solution`] proof.
    pub fn new(nonce: Nonce, effort: Effort, seed_head: SeedHead, proof: equix::Solution) -> Self {
        Solution {
            nonce,
            effort,
            seed_head,
            proof,
        }
    }

    /// Try to build a [`Solution`] from an unvalidated [`equix::SolutionByteArray`].
    ///
    /// This will either return a [`Solution`] or a [`SolutionErrorV1::Order`].
    pub fn try_from_bytes(
        nonce: Nonce,
        effort: Effort,
        seed_head: SeedHead,
        bytes: &equix::SolutionByteArray,
    ) -> Result<Self, SolutionErrorV1> {
        Ok(Self::new(
            nonce,
            effort,
            seed_head,
            equix::Solution::try_from_bytes(bytes).map_err(|_| SolutionErrorV1::Order)?,
        ))
    }

    /// Get the winning [`Nonce`] value used in this solution
    pub fn nonce(&self) -> &Nonce {
        &self.nonce
    }

    /// Get the client-chosen and provable [`Effort`] value used in this solution
    pub fn effort(&self) -> Effort {
        self.effort
    }

    /// Get the [`SeedHead`] value identifying the puzzle this solution is for
    pub fn seed_head(&self) -> SeedHead {
        self.seed_head
    }

    /// Reference the [`equix::Solution`] backing the proof portion
    pub fn proof(&self) -> &equix::Solution {
        &self.proof
    }

    /// Clone the proof portion of the solution in its canonical byte string format
    pub fn proof_to_bytes(&self) -> equix::SolutionByteArray {
        self.proof.to_bytes()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn new_seed_head_collision() {
        use tor_basic_utils::test_rng::Config;

        let mut rng = Config::Seeded([
            191, 169, 93, 126, 223, 128, 184, 217, 105, 3, 80, 87, 181, 242, 206, 38, 152, 149,
            116, 71, 249, 142, 6, 196, 188, 141, 20, 139, 97, 45, 230, 55,
        ])
        .into_rng();

        let seed_1 = Seed::new(&mut rng, None);
        let seed_2_shared_head = Seed::new(&mut rng, None);

        // Verify that this RNG seed does actually result in a collision.
        // If the code to generate the seed is changed, we may need to find a new RNG seed to
        // generate this collision.
        assert_eq!(seed_1.head(), seed_2_shared_head.head());

        let mut rng = Config::Seeded([0x00; 32]).into_rng();

        let seed_1 = Seed::new(&mut rng, None);
        let seed_2 = Seed::new(&mut rng, Some(&seed_1));

        assert_ne!(seed_1.head(), seed_2.head());
    }
}
