use crate::PublicKey;

use thiserror::Error;

pub(crate) type PrivateKeyBytes = [u8; 32];

/// Functions required to use an `ExtendedPrivateKey`.
pub trait PrivateKey: Sized {
    /// A type implementing [`PublicKey`] associated with Self.
    type PublicKey: PublicKey;

    /// The error returned by [`PrivateKey::derive_child()`].
    type Err: std::error::Error;

    /// Create a Self from bytes.
    ///
    /// # Examples
    /// ```rust
    /// # use keyfork_derive_util::{
    /// #   *,
    /// #   private_key::TestPrivateKey as OurPrivateKey,
    /// # };
    /// let key_data: &[u8; 32] = //
    /// #   b"AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA";
    /// let private_key = OurPrivateKey::from_bytes(key_data);
    /// ```
    fn from_bytes(b: &PrivateKeyBytes) -> Result<Self, Self::Err>;

    /// Convert a &Self to bytes.
    ///
    /// # Examples
    /// ```rust
    /// # use keyfork_derive_util::{
    /// #   *,
    /// #   private_key::TestPrivateKey as OurPrivateKey,
    /// # };
    /// let key_data: &[u8; 32] = //
    /// #   b"AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA";
    /// let private_key = OurPrivateKey::from_bytes(key_data).unwrap();
    /// assert_eq!(key_data, &private_key.to_bytes());
    /// ```
    fn to_bytes(&self) -> PrivateKeyBytes;

    /*
     * Freak of nature, unsupported?
    /// Whether or not zero is a valid public key (such as with ed25519 keys).
    fn is_zero_valid_public_key() -> bool {
        false
    }
    */

    /// The initial key for BIP-0032 and SLIP-0010 derivation, such as secp256k1's "Bitcoin seed".
    ///
    /// # Examples
    /// ```rust
    /// # use keyfork_derive_util::{
    /// #   *,
    /// #   private_key::TestPrivateKey as OurPrivateKey,
    /// # };
    /// assert_eq!(OurPrivateKey::key(), "testing seed");
    /// ```
    fn key() -> &'static str;

    /// Generate a [`Self::PublicKey`].
    ///
    /// # Examples
    /// ```rust
    /// # use keyfork_derive_util::{
    /// #   *,
    /// #   private_key::TestPrivateKey as OurPrivateKey,
    /// # };
    /// let key_data: &[u8; 32] = //
    /// #   b"AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA";
    /// let private_key = OurPrivateKey::from_bytes(key_data).unwrap();
    /// let public_key = private_key.public_key();
    /// ```
    fn public_key(&self) -> Self::PublicKey;

    /// Derive a child [`PrivateKey`] with given `PrivateKeyBytes`. The implementation of
    /// derivation is algorithm-specific and a specification should be consulted when implementing
    /// this method.
    ///
    /// # Errors
    ///
    /// An error may be returned if:
    /// * An all-zero `other` is provided.
    /// * An error specific to the given algorithm was encountered.
    fn derive_child(&self, other: &PrivateKeyBytes) -> Result<Self, Self::Err>;

    /// Whether the algorithm requires hardened derivation, such as for ed25519.
    fn requires_hardened_derivation() -> bool {
        false
    }
}

/// Errors associated with creating and arithmetic on private keys. This specific error is only
/// intended to be used by the implementations in this crate.
#[derive(Clone, Debug, Error)]
pub enum PrivateKeyError {
    /// For the given algorithm, the private key must be nonzero.
    #[error("The provided private key must be nonzero, but is not")]
    NonZero,

    /// A scalar could not be constructed for the given algorithm.
    #[error("A scalar could not be constructed for the given algorithm")]
    InvalidScalar,
}

#[cfg(feature = "secp256k1")]
use k256::NonZeroScalar;

#[cfg(feature = "secp256k1")]
impl PrivateKey for k256::SecretKey {
    type Err = PrivateKeyError;
    type PublicKey = k256::PublicKey;

    fn key() -> &'static str {
        "Bitcoin seed"
    }

    fn from_bytes(b: &PrivateKeyBytes) -> Result<Self, Self::Err> {
        Self::from_slice(b).map_err(|_| PrivateKeyError::InvalidScalar)
    }

    fn to_bytes(&self) -> PrivateKeyBytes {
        // Note: Safety assured by type returned from EncodedPoint
        self.to_bytes().into()
    }

    fn public_key(&self) -> Self::PublicKey {
        self.public_key()
    }

    fn derive_child(&self, other: &PrivateKeyBytes) -> Result<Self, Self::Err> {
        use k256::elliptic_curve::ScalarPrimitive;
        use k256::{Scalar, Secp256k1};

        // Construct a scalar from bytes
        let scalar = ScalarPrimitive::<Secp256k1>::from_bytes(other.into());
        let scalar = Option::<ScalarPrimitive<Secp256k1>>::from(scalar);
        let scalar = scalar.ok_or(PrivateKeyError::InvalidScalar)?;
        let scalar = Scalar::from(scalar);

        let derived_scalar = self.to_nonzero_scalar().as_ref() + scalar.as_ref();
        let nonzero_scalar = Option::<NonZeroScalar>::from(NonZeroScalar::new(derived_scalar))
            .ok_or(PrivateKeyError::NonZero)?;
        Ok(Self::from(nonzero_scalar))
    }
}

#[cfg(feature = "ed25519")]
impl PrivateKey for ed25519_dalek::SigningKey {
    type Err = PrivateKeyError;
    type PublicKey = ed25519_dalek::VerifyingKey;

    fn key() -> &'static str {
        "ed25519 seed"
    }

    fn from_bytes(b: &PrivateKeyBytes) -> Result<Self, Self::Err> {
        Ok(Self::from_bytes(b))
    }

    fn to_bytes(&self) -> PrivateKeyBytes {
        self.to_bytes()
    }

    fn public_key(&self) -> Self::PublicKey {
        self.verifying_key()
    }

    fn derive_child(&self, other: &PrivateKeyBytes) -> Result<Self, Self::Err> {
        // SLIP-0010: No arithmetic required for ed25519 keys.
        Ok(Self::from_bytes(other))
    }

    fn requires_hardened_derivation() -> bool {
        true
    }
}

use crate::public_key::TestPublicKey;

/// A private key that can be used for testing purposes. Does not utilize any significant
/// cryptographic operations.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct TestPrivateKey {
    key: [u8; 32],
}

impl TestPrivateKey {
    pub(crate) fn public_key(&self) -> TestPublicKey {
        let mut bytes = [0u8; 33];
        for (i, byte) in self.key.iter().enumerate() {
            bytes[i + 1] = byte ^ 0xFF;
        }
        TestPublicKey { key: bytes }
    }
}

impl PrivateKey for TestPrivateKey {
    type PublicKey = TestPublicKey;
    type Err = PrivateKeyError;

    fn from_bytes(b: &PrivateKeyBytes) -> Result<Self, Self::Err> {
        Ok(Self { key: *b })
    }

    fn to_bytes(&self) -> PrivateKeyBytes {
        self.key
    }

    fn key() -> &'static str {
        "testing seed"
    }

    fn public_key(&self) -> Self::PublicKey {
        self.public_key()
    }

    fn derive_child(&self, other: &PrivateKeyBytes) -> Result<Self, Self::Err> {
        Ok(Self { key: *other })
    }
}