use serde::{Serialize, Deserialize};
use sha2::{Sha256, Digest};
use data_encoding::BASE64URL;

/// Methods to use for PKCE challenges.
#[derive(PartialEq, Eq, Copy, Clone, Debug, Serialize, Deserialize, Default)]
pub enum PKCEMethod {
    /// Base64-encoded SHA256 hash of an ASCII string.
    #[default]
    S256,
    /// Plain string by itself. Please don't use this.
    #[serde(rename = "snake_case")]
    Plain
}

impl PKCEMethod {
    /// Return a string representing a PKCE method as it would be serialized.
    pub fn as_str(&self) -> &'static str {
        match self {
            PKCEMethod::S256 => "S256",
            PKCEMethod::Plain => "plain"
        }
    }
}
/// A PKCE verifier string that should be kept in secret until the end
/// of the authentication ceremony, where it is revealed to prove that
/// the one who uses the grant is the same entity who it was given to.
#[derive(Debug, PartialEq, Eq, Clone, Serialize, Deserialize)]
pub struct PKCEVerifier(pub(super) String);

impl<T: Into<String>> From<T> for PKCEVerifier {
    fn from(t: T) -> Self {
        Self(t.into())
    }
}
impl AsRef<str> for PKCEVerifier {
    fn as_ref(&self) -> &str {
        self.0.as_str()
    }
}

// We don't need Display here, this is not a user-facing value.
#[allow(clippy::to_string_trait_impl)]
impl ToString for PKCEVerifier {
    fn to_string(&self) -> String {
        self.0.clone()
    }
}

impl PKCEVerifier {
    /// Generate a new PKCE verifier string of 128 bytes in length.
    #[allow(clippy::new_without_default)]
    pub fn new() -> Self {
        Self::from_rng(&mut rand::thread_rng())
    }

    /// Generate a new PKCE string of 128 bytes in length, using
    /// the provided random number generator.
    pub fn from_rng(rng: &mut (impl rand::CryptoRng + rand::Rng)) -> Self {
        use rand::{Rng, distributions::Alphanumeric};

        let bytes = rng
            .sample_iter(&Alphanumeric)
            .take(128)
            .collect::<Vec<u8>>();
        Self(String::from_utf8(bytes).unwrap())
    }

}

/// A PKCE challenge as described in [RFC7636].
///
/// [RFC7636]: https://tools.ietf.org/html/rfc7636
#[derive(Eq, PartialEq, Debug, Clone, Serialize, Deserialize)]
pub struct PKCEChallenge {
    code_challenge: String,
    #[serde(rename = "code_challenge_method")]
    method: PKCEMethod
}

impl PKCEChallenge {
    /// Create a new challenge from a [PKCEVerifier] using a certain
    /// [PKCEMethod].
    pub fn new(code_verifier: &PKCEVerifier, method: PKCEMethod) -> Self {
        Self {
            code_challenge: match method {
                PKCEMethod::S256 => {
                    let mut hasher = Sha256::new();
                    hasher.update(code_verifier.as_ref());
                    let mut challenge = BASE64URL.encode(&hasher.finalize());
                    // Trim padding
                    challenge.retain(|c| c != '=');

                    challenge
                },
                PKCEMethod::Plain => code_verifier.to_string(),
            },
            method
        }
    }

    /// Verify that the [PKCEVerifier] corresponds to this challenge,
    /// by creating a second challenge string and comparing it against
    /// this challenge data.
    ///
    /// ```rust
    /// use kittybox_indieauth::{PKCEVerifier, PKCEMethod, PKCEChallenge};
    ///
    /// let verifier = PKCEVerifier::new();
    /// let challenge = PKCEChallenge::new(&verifier, PKCEMethod::default());
    /// // Meanwhile, at the token endpoint, in the end of the ceremony...
    /// // ...the challenge gets retrieved from the stored data and verified
    /// assert!(challenge.verify(verifier))
    /// ```
    #[must_use]
    pub fn verify(&self, code_verifier: PKCEVerifier) -> bool {
        Self::new(&code_verifier, self.method) == *self
    }

    /// Return a reference to the code challenge string.
    pub fn as_str(&self) -> &str {
        self.code_challenge.as_str()
    }

    /// Return the method used to create this challenge.
    pub fn method(&self) -> PKCEMethod {
        self.method
    }
}

#[cfg(test)]
mod tests {
    use super::{PKCEMethod, PKCEVerifier, PKCEChallenge};

    #[test]
    /// A snapshot test generated using [Aaron Parecki's PKCE
    /// tools](https://example-app.com/pkce) that checks for a
    /// conforming challenge.
    fn test_pkce_challenge_verification() {
        let verifier = PKCEVerifier("ec03310e4e90f7bc988af05384060c3c1afeae4bb4d0f648c5c06b63".to_owned());

        let challenge = PKCEChallenge::new(&verifier, PKCEMethod::S256);

        assert_eq!(challenge.as_str(), "aB8OG20Rh8UoQ9gFhI0YvPkx4dDW2MBspBKGXL6j6Wg");
    }
}