use crate::private_key::PrivateKeyBytes;

use digest::Digest;
use ripemd::Ripemd160;
use sha2::Sha256;
use thiserror::Error;

pub type PublicKeyBytes = [u8; 33];

pub trait PublicKey: Sized {
    type Err: std::error::Error;

    fn from_bytes(b: &PublicKeyBytes) -> Self;
    fn to_bytes(&self) -> PublicKeyBytes;

    /// Derive a child [`PublicKey`] with given `PrivateKeyBytes`.
    ///
    /// # Errors
    ///
    /// An error may be returned if:
    /// * A nonzero `other` is provided.
    /// * An error specific to the given algorithm was encountered.
    fn derive_child(&self, other: PrivateKeyBytes) -> Result<Self, Self::Err>;

    fn fingerprint(&self) -> [u8; 4] {
        let hash = Sha256::new().chain_update(self.to_bytes()).finalize();
        let hash = Ripemd160::new().chain_update(hash).finalize();
        // Note: Safety assured by type returned from Ripemd160
        hash[..4].try_into().unwrap()
    }
}

#[derive(Clone, Debug, Error)]
pub enum PublicKeyError {
    #[error("The provided public key must be nonzero, but is not")]
    NonZero,

    #[error("Unable to convert point to key")]
    PointToKey(#[from] k256::elliptic_curve::Error),
}

#[cfg(feature = "secp256k1")]
use k256::{
    elliptic_curve::{group::prime::PrimeCurveAffine, sec1::ToEncodedPoint},
    AffinePoint, NonZeroScalar,
};

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

    fn from_bytes(b: &PublicKeyBytes) -> Self {
        Self::from_sec1_bytes(b).expect("Invalid public key bytes")
    }

    fn to_bytes(&self) -> PublicKeyBytes {
        // Note: Safety assured by type returned from EncodedPoint
        self.to_encoded_point(true).as_bytes().try_into().unwrap()
    }

    fn derive_child(&self, other: PrivateKeyBytes) -> Result<Self, Self::Err> {
        if other.iter().all(|n| n == &0) {
            return Err(PublicKeyError::NonZero);
        }
        // Checked: See above
        let scalar = Option::<NonZeroScalar>::from(NonZeroScalar::from_repr(other.into()))
            .expect("Should have been able to get a NonZeroScalar");

        let point = self.to_projective() + (AffinePoint::generator() * *scalar);
        Self::from_affine(point.into()).map_err(From::from)
    }
}