//! Provides a private identity of a Bitmessage message sender
//! and a public identity of a Bitmessage message receiver.

use std::{
    convert::TryInto,
    fmt,
    sync::{
        atomic::{AtomicBool, Ordering},
        Arc,
    },
};

use crate::{
    address::{
        count_zeros, Address, Error as AddressError, Version as AddressVersion, DEFAULT_ZEROS,
        VERSION as CURRENT_ADDRESS_VERSION,
    },
    config::{default_nonce_trials_per_byte, default_payload_length_extra_bytes, Config},
    crypto::{KeyPair, PrivateKey, PrivateKeyError, PublicKey},
    hash::{ripemd160_sha512, sha512},
    io::WriteTo,
    pow::{NonceTrialsPerByte, PayloadLengthExtraBytes},
    stream::{StreamNumber, ROOT as ROOT_STREAM},
    var_type::VarInt,
};

pub use crate::feature::Features;

/// A public identity of a Bitmessage message receiver.
#[derive(Clone, PartialEq, Eq, Hash, Debug)]
pub struct Public {
    version: AddressVersion,
    stream: StreamNumber,

    features: Features,
    public_signing_key: PublicKey,
    public_encryption_key: PublicKey,
    nonce_trials_per_byte: NonceTrialsPerByte,
    payload_length_extra_bytes: PayloadLengthExtraBytes,
}

impl fmt::Display for Public {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.address().fmt(f)
    }
}

impl Public {
    /// Creates a public identity.
    pub fn new(
        version: AddressVersion,
        stream: StreamNumber,
        features: Features,
        public_signing_key: PublicKey,
        public_encryption_key: PublicKey,
        nonce_trials_per_byte: NonceTrialsPerByte,
        payload_length_extra_bytes: PayloadLengthExtraBytes,
    ) -> Result<Self, AddressError> {
        let _address = Address::from_public_keys(
            version,
            stream,
            &public_signing_key,
            &public_encryption_key,
        )?;
        Ok(Self {
            version,
            stream,
            features,
            public_signing_key,
            public_encryption_key,
            nonce_trials_per_byte,
            payload_length_extra_bytes,
        })
    }

    /// Returns the address
    pub fn address(&self) -> Address {
        Address::from_public_keys(
            self.version,
            self.stream,
            &self.public_signing_key,
            &self.public_encryption_key,
        )
        .unwrap()
    }

    /// Returns the address version.
    pub fn version(&self) -> AddressVersion {
        self.version
    }

    /// Returns the stream number.
    pub fn stream(&self) -> StreamNumber {
        self.stream
    }

    /// Returns the features.
    pub fn features(&self) -> Features {
        self.features
    }

    /// Returns the public signing key.
    pub fn public_signing_key(&self) -> &PublicKey {
        &self.public_signing_key
    }

    /// Returns the public encryption key.
    pub fn public_encryption_key(&self) -> &PublicKey {
        &self.public_encryption_key
    }

    /// Returns the nonce trials per byte.
    pub fn nonce_trials_per_byte(&self) -> NonceTrialsPerByte {
        self.nonce_trials_per_byte
    }

    /// Returns the payload length extra bytes.
    pub fn payload_length_extra_bytes(&self) -> PayloadLengthExtraBytes {
        self.payload_length_extra_bytes
    }
}

/// A private identity of a Bitmessage message sender.
#[derive(Clone, PartialEq, Eq, Hash, Debug)]
pub struct Private {
    version: AddressVersion,
    stream: StreamNumber,

    features: Features,
    nonce_trials_per_byte: NonceTrialsPerByte,
    payload_length_extra_bytes: PayloadLengthExtraBytes,

    private_signing_key: PrivateKey,
    private_encryption_key: PrivateKey,
    chan: bool,
}

impl fmt::Display for Private {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.address().fmt(f)
    }
}

impl Private {
    /// Creates a private identity.
    #[allow(clippy::too_many_arguments)]
    pub fn new(
        version: AddressVersion,
        stream: StreamNumber,
        features: Features,
        nonce_trials_per_byte: NonceTrialsPerByte,
        payload_length_extra_bytes: PayloadLengthExtraBytes,
        private_signing_key: PrivateKey,
        private_encryption_key: PrivateKey,
        chan: bool,
    ) -> Result<Self, AddressError> {
        let _address = Address::from_public_keys(
            version,
            stream,
            &private_signing_key.public_key(),
            &private_encryption_key.public_key(),
        )?;
        Ok(Self {
            version,
            stream,
            features,
            nonce_trials_per_byte,
            payload_length_extra_bytes,
            private_signing_key,
            private_encryption_key,
            chan,
        })
    }

    /// Returns the address
    pub fn address(&self) -> Address {
        Address::from_public_keys(
            self.version,
            self.stream,
            &self.public_signing_key(),
            &self.public_encryption_key(),
        )
        .unwrap()
    }

    /// Returns the address version.
    pub fn version(&self) -> AddressVersion {
        self.version
    }

    /// Returns the stream number.
    pub fn stream(&self) -> StreamNumber {
        self.stream
    }

    /// Returns the features.
    pub fn features(&self) -> Features {
        self.features
    }

    /// Returns the nonce trials per byte.
    pub fn nonce_trials_per_byte(&self) -> NonceTrialsPerByte {
        self.nonce_trials_per_byte
    }

    /// Returns the payload length extra bytes.
    pub fn payload_length_extra_bytes(&self) -> PayloadLengthExtraBytes {
        self.payload_length_extra_bytes
    }

    /// Returns the private signing key.
    pub fn private_signing_key(&self) -> &PrivateKey {
        &self.private_signing_key
    }

    /// Returns the private encryption key.
    pub fn private_encryption_key(&self) -> &PrivateKey {
        &self.private_encryption_key
    }

    /// Returns true if this is chan identity.
    pub fn chan(&self) -> bool {
        self.chan
    }

    /// Returns the public signing key.
    pub fn public_signing_key(&self) -> PublicKey {
        self.private_signing_key.public_key()
    }

    /// Returns the public encryption key.
    pub fn public_encryption_key(&self) -> PublicKey {
        self.private_encryption_key.public_key()
    }

    /// Returns a random identity builder.
    pub fn random_builder() -> RandomBuilder {
        RandomBuilder::new()
    }

    /// Returns a deterministic identity builder.
    pub fn deterministic_builder(password: Vec<u8>) -> DeterministicBuilder {
        DeterministicBuilder::new(password)
    }

    /// Returns a chan identity builder.
    pub fn chan_builder(password: Vec<u8>) -> ChanBuilder {
        ChanBuilder::new(password)
    }
}

impl From<&Private> for Public {
    fn from(v: &Private) -> Public {
        Self {
            version: v.version,
            stream: v.stream,
            features: v.features,
            public_signing_key: v.public_signing_key(),
            public_encryption_key: v.public_encryption_key(),
            nonce_trials_per_byte: v.nonce_trials_per_byte,
            payload_length_extra_bytes: v.payload_length_extra_bytes,
        }
    }
}

/// An identity of a Bitmessage user.
#[derive(Clone, PartialEq, Eq, Hash, Debug)]
pub enum Identity {
    /// A Bitmessage address.
    Address(Address),
    /// A public identity.
    Public(Public),
    /// A private identity.
    Private(Private),
}

impl fmt::Display for Identity {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.address().fmt(f)
    }
}

impl Identity {
    /// Returns the address.
    pub fn address(&self) -> Address {
        match self {
            Identity::Address(address) => address.clone(),
            Identity::Public(public) => public.address(),
            Identity::Private(private) => private.address(),
        }
    }
}

/// An error which can be returned when processing an identity.
#[derive(Clone, Debug)]
pub enum Error {
    /// Indicates that construction of an address from public keys failed.
    AddressError(AddressError),
    /// Indicates that the operation was canceled.
    Canceled,
    /// Indicates that the operation on a private key failed.
    PrivateKeyError(PrivateKeyError),
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::AddressError(err) => err.fmt(f),
            Self::Canceled => "canceled".fmt(f),
            Self::PrivateKeyError(err) => err.fmt(f),
        }
    }
}

impl std::error::Error for Error {}

impl From<AddressError> for Error {
    fn from(err: AddressError) -> Self {
        Self::AddressError(err)
    }
}

impl From<PrivateKeyError> for Error {
    fn from(err: PrivateKeyError) -> Self {
        Self::PrivateKeyError(err)
    }
}

/// A builder for building a random identity.
#[derive(Clone, PartialEq, Eq, Hash, Debug)]
pub struct RandomBuilder {
    version: AddressVersion,
    stream: StreamNumber,
    zeros: usize,
    features: Features,
    nonce_trials_per_byte: NonceTrialsPerByte,
    payload_length_extra_bytes: PayloadLengthExtraBytes,
}

impl Default for RandomBuilder {
    fn default() -> Self {
        Self {
            version: CURRENT_ADDRESS_VERSION,
            stream: ROOT_STREAM,
            zeros: DEFAULT_ZEROS,
            features: Features::DOES_ACK,
            nonce_trials_per_byte: default_nonce_trials_per_byte(),
            payload_length_extra_bytes: default_payload_length_extra_bytes(),
        }
    }
}

impl RandomBuilder {
    fn new() -> Self {
        Self::default()
    }

    /// A shortcut to set parameters from a config object.
    pub fn config(&mut self, v: &Config) -> &mut Self {
        self.nonce_trials_per_byte(v.nonce_trials_per_byte())
            .payload_length_extra_bytes(v.payload_length_extra_bytes())
    }

    /// Sets the address version.
    pub fn version(&mut self, v: AddressVersion) -> &mut Self {
        self.version = v;
        self
    }

    /// Sets the stream number.
    pub fn stream(&mut self, v: StreamNumber) -> &mut Self {
        self.stream = v;
        self
    }

    /// Sets the minimum count of leading zeros in the hash of the public keys.
    pub fn zeros(&mut self, v: usize) -> &mut Self {
        self.zeros = v;
        self
    }

    /// Sets the features.
    pub fn features(&mut self, v: Features) -> &mut Self {
        self.features = v;
        self
    }

    /// Sets the nonce trials per byte.
    pub fn nonce_trials_per_byte(&mut self, v: NonceTrialsPerByte) -> &mut Self {
        self.nonce_trials_per_byte = v;
        self
    }

    /// Sets the payload length extra bytes.
    pub fn payload_length_extra_bytes(&mut self, v: PayloadLengthExtraBytes) -> &mut Self {
        self.payload_length_extra_bytes = v;
        self
    }

    /// Returns the private identity this builder represents.
    pub fn build(&self, cancel: Arc<AtomicBool>) -> Result<Private, Error> {
        let signing_key_pair = KeyPair::generate();
        let signing_public_key_bytes = signing_key_pair.public_key().encode();
        while !cancel.load(Ordering::SeqCst) {
            for _ in 0..0x1000 {
                let encryption_key_pair = KeyPair::generate();
                let mut bytes = Vec::new();
                bytes.extend_from_slice(&signing_public_key_bytes);
                bytes.append(&mut encryption_key_pair.public_key().encode());
                let hash = ripemd160_sha512(bytes);
                if count_zeros(hash) < self.zeros {
                    continue;
                }
                return Ok(Private::new(
                    self.version,
                    self.stream,
                    self.features,
                    self.nonce_trials_per_byte,
                    self.payload_length_extra_bytes,
                    signing_key_pair.private_key().clone(),
                    encryption_key_pair.private_key().clone(),
                    false,
                )?);
            }
        }
        Err(Error::Canceled)
    }
}

/// A builder for building a deterministic identity.
#[derive(Clone, PartialEq, Eq, Hash, Debug)]
pub struct DeterministicBuilder {
    version: AddressVersion,
    stream: StreamNumber,
    zeros: usize,
    features: Features,
    nonce_trials_per_byte: NonceTrialsPerByte,
    payload_length_extra_bytes: PayloadLengthExtraBytes,
    password: Vec<u8>,
}

impl DeterministicBuilder {
    fn new(password: Vec<u8>) -> Self {
        Self {
            version: CURRENT_ADDRESS_VERSION,
            stream: ROOT_STREAM,
            zeros: DEFAULT_ZEROS,
            features: Features::DOES_ACK,
            nonce_trials_per_byte: default_nonce_trials_per_byte(),
            payload_length_extra_bytes: default_payload_length_extra_bytes(),
            password,
        }
    }

    /// A shortcut to set parameters from a config object.
    pub fn config(&mut self, v: &Config) -> &mut Self {
        self.nonce_trials_per_byte(v.nonce_trials_per_byte())
            .payload_length_extra_bytes(v.payload_length_extra_bytes())
    }

    /// Sets the address version.
    pub fn version(&mut self, v: AddressVersion) -> &mut Self {
        self.version = v;
        self
    }

    /// Sets the stream number.
    pub fn stream(&mut self, v: StreamNumber) -> &mut Self {
        self.stream = v;
        self
    }

    /// Sets the minimum count of leading zeros in the hash of the public keys.
    pub fn zeros(&mut self, v: usize) -> &mut Self {
        self.zeros = v;
        self
    }

    /// Sets the features.
    pub fn features(&mut self, v: Features) -> &mut Self {
        self.features = v;
        self
    }

    /// Sets the nonce trials per byte.
    pub fn nonce_trials_per_byte(&mut self, v: NonceTrialsPerByte) -> &mut Self {
        self.nonce_trials_per_byte = v;
        self
    }

    /// Sets the payload length extra bytes.
    pub fn payload_length_extra_bytes(&mut self, v: PayloadLengthExtraBytes) -> &mut Self {
        self.payload_length_extra_bytes = v;
        self
    }

    /// Returns the private identity this builder represents.
    pub fn build(&self, n: usize, cancel: Arc<AtomicBool>) -> Result<Vec<Private>, Error> {
        let mut identities = Vec::new();
        let mut signing_key_nonce: u64 = 0;
        let mut encryption_key_nonce: u64 = 1;
        for _ in 0..n {
            'outer: loop {
                if cancel.load(Ordering::SeqCst) {
                    return Err(Error::Canceled);
                }
                for _ in 0..0x1000 {
                    let mut bytes = self.password.clone();
                    VarInt::new(signing_key_nonce).write_to(&mut bytes).unwrap();
                    let private_signing_key =
                        PrivateKey::new(sha512(bytes)[..32].try_into().unwrap())?;
                    signing_key_nonce += 2;

                    let mut bytes = self.password.clone();
                    VarInt::new(encryption_key_nonce)
                        .write_to(&mut bytes)
                        .unwrap();
                    let private_encryption_key =
                        PrivateKey::new(sha512(bytes)[..32].try_into().unwrap())?;
                    encryption_key_nonce += 2;

                    let mut bytes = private_signing_key.public_key().encode();
                    bytes.append(&mut private_encryption_key.public_key().encode());
                    let hash = ripemd160_sha512(bytes);
                    if count_zeros(hash) >= self.zeros {
                        let identity = Private::new(
                            self.version,
                            self.stream,
                            self.features,
                            self.nonce_trials_per_byte,
                            self.payload_length_extra_bytes,
                            private_signing_key,
                            private_encryption_key,
                            false,
                        )?;
                        identities.push(identity);
                        break 'outer;
                    }
                }
            }
        }
        Ok(identities)
    }
}

#[test]
fn test_deterministic_builder() {
    let cancel = Arc::new(AtomicBool::new(false));
    let identities = Private::deterministic_builder(b"hello".to_vec())
        .build(2, cancel)
        .unwrap();
    assert_eq!(
        identities[0].address().to_string(),
        "BM-2cWhA72reAp1CBa8JmspqWRCdw93sDLgiS"
    );
    assert_eq!(
        identities[1].address().to_string(),
        "BM-2cWoETgcY3YJHZfKPVLS1avrXtbzJtcvDY"
    );
}

/// A builder for building a chan identity.
#[derive(Clone, PartialEq, Eq, Hash, Debug)]
pub struct ChanBuilder {
    version: AddressVersion,
    stream: StreamNumber,
    zeros: usize,
    features: Features,
    nonce_trials_per_byte: NonceTrialsPerByte,
    payload_length_extra_bytes: PayloadLengthExtraBytes,
    password: Vec<u8>,
}

impl ChanBuilder {
    fn new(password: Vec<u8>) -> Self {
        Self {
            version: CURRENT_ADDRESS_VERSION,
            stream: ROOT_STREAM,
            zeros: 1,
            features: Features::empty(),
            nonce_trials_per_byte: default_nonce_trials_per_byte(),
            payload_length_extra_bytes: default_payload_length_extra_bytes(),
            password,
        }
    }

    /// Sets the address version.
    pub fn version(&mut self, v: AddressVersion) -> &mut Self {
        self.version = v;
        self
    }

    /// Sets the stream number.
    pub fn stream(&mut self, v: StreamNumber) -> &mut Self {
        self.stream = v;
        self
    }

    /// Returns the private identity this builder represents.
    pub fn build(&self, cancel: Arc<AtomicBool>) -> Result<Private, Error> {
        let builder = DeterministicBuilder {
            version: self.version,
            stream: self.stream,
            zeros: self.zeros,
            features: self.features,
            nonce_trials_per_byte: self.nonce_trials_per_byte,
            payload_length_extra_bytes: self.payload_length_extra_bytes,
            password: self.password.clone(),
        };
        let identities = builder.build(1, cancel)?;
        let mut identity = identities[0].clone();
        identity.chan = true;
        Ok(identity)
    }
}