//! Contains a Solana Pubkey implementation with some modifications to allow for efficient
//! serialization.

#[cfg(feature = "pubkey-str")]
use core::fmt::Debug;
use core::hash::{Hash, Hasher};
use core::str::FromStr;

#[cfg(feature = "solana-pubkey-convert-v2")]
use solana_pubkey_v2::Pubkey as SolanaPubkeyV2;

/// Number of bytes in a pubkey.
pub const PUBKEY_BYTES: usize = 32;

/// The address of a [Solana account].
///
/// This is a thin copy of the [`solana_pubkey::Pubkey`] type, used to force more efficient
/// encoding when using `serde`.
///
/// This can be easily converted to and from the [`solana_pubkey::Pubkey`] type if the
/// `solana-pubkey-convert` feature is enabled.
///
/// [Solana account]: https://solana.com/docs/core/accounts
#[repr(transparent)]
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, PartialOrd, Ord)]
pub struct Pubkey(pub(crate) [u8; PUBKEY_BYTES]);

#[cfg(feature = "serde")]
impl serde::Serialize for Pubkey {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        #[cfg(feature = "serde-pubkey-human-base58")]
        if serializer.is_human_readable() {
            return self.to_string().serialize(serializer);
        }

        serde_bytes::serialize(&self.0, serializer)
    }
}

#[cfg(feature = "serde")]
struct PubkeyVisitor;

#[cfg(feature = "serde")]
impl<'de> serde::de::Visitor<'de> for PubkeyVisitor {
    type Value = Pubkey;

    fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result {
        formatter.write_str("32-byte Solana public key as binary")
    }

    fn visit_seq<A>(self, mut seq: A) -> Result<Self::Value, A::Error>
    where
        A: serde::de::SeqAccess<'de>,
    {
        let mut data: [u8; PUBKEY_BYTES] = Default::default();
        for (idx, byte) in data.iter_mut().enumerate() {
            *byte = seq
                .next_element()?
                .ok_or_else(|| serde::de::Error::invalid_length(idx, &self))?;
        }
        Ok(Pubkey(data))
    }

    fn visit_bytes<E>(self, v: &[u8]) -> Result<Self::Value, E>
    where
        E: serde::de::Error,
    {
        if v.len() != PUBKEY_BYTES {
            return Err(E::invalid_length(v.len(), &self));
        }
        let mut data: [u8; PUBKEY_BYTES] = Default::default();
        data.copy_from_slice(v);
        Ok(Pubkey(data))
    }

    #[cfg(feature = "serde-pubkey-human-base58")]
    fn visit_str<E>(self, v: &str) -> Result<Self::Value, E>
    where
        E: serde::de::Error,
    {
        Pubkey::from_str(v).map_err(E::custom)
    }
}

#[cfg(feature = "serde")]
impl<'de> serde::Deserialize<'de> for Pubkey {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        deserializer.deserialize_any(PubkeyVisitor)
    }
}

/// Custom implementation of Hash for Pubkey.
/// This allows us to skip hashing the length of the pubkey, which is always the same.
///
/// Implementation taken from [`solana_pubkey::Pubkey`]'s `Hash` implementation.
impl Hash for Pubkey {
    fn hash<H: Hasher>(&self, state: &mut H) {
        state.write(self.as_array())
    }
}

impl Pubkey {
    /// Constructs a new `Pubkey` from the given byte array.
    pub const fn new_from_array(pubkey_array: [u8; PUBKEY_BYTES]) -> Self {
        Self(pubkey_array)
    }

    /// Decode a string into a Pubkey in a const context.
    ///
    /// Not as efficient as the standard `FromStr` implementation, so should generally only be used
    /// for compile-time constants.
    ///
    /// # Panics
    ///
    /// This function will panic if the input value is not a valid base58 string encoding exactly
    /// 32 bytes.
    ///
    /// # Examples
    ///
    /// ```
    /// use titan_api_types::common::Pubkey;
    ///
    /// let sol = Pubkey::from_str_const("So11111111111111111111111111111111111111111");
    /// ```
    #[cfg(feature = "pubkey-fromstr-const")]
    pub const fn from_str_const(s: &str) -> Self {
        let id_bytes = five8_const::decode_32_const(s);
        Self::new_from_array(id_bytes)
    }

    /// Converts the `Pubkey` to an array of bytes.
    pub const fn to_bytes(self) -> [u8; PUBKEY_BYTES] {
        self.0
    }

    /// Return a reference to the `Pubkey`'s byte array.
    pub const fn as_array(&self) -> &[u8; PUBKEY_BYTES] {
        &self.0
    }
}

/// Convenient macro to define a static public key.
///
/// Based on the same macro from the [`solana_pubkey::pubkey`] crate.
///
/// Simply calls [`Pubkey::from_str_const`] on the specified input.
#[cfg(feature = "pubkey-fromstr-const")]
#[macro_export]
macro_rules! pubkey {
    ($input:literal) => {
        $crate::common::Pubkey::from_str_const($input)
    };
}

impl From<[u8; PUBKEY_BYTES]> for Pubkey {
    #[inline]
    fn from(value: [u8; PUBKEY_BYTES]) -> Self {
        Self(value)
    }
}

impl TryFrom<&[u8]> for Pubkey {
    type Error = std::array::TryFromSliceError;

    #[inline]
    fn try_from(value: &[u8]) -> Result<Self, Self::Error> {
        <[u8; PUBKEY_BYTES]>::try_from(value).map(Self::from)
    }
}

#[cfg(feature = "solana-pubkey-convert-v2")]
impl From<SolanaPubkeyV2> for Pubkey {
    /// Convert a [`solana_pubkey::Pubkey`] into a [`Pubkey`] by taking ownership of its bytes.
    #[inline]
    fn from(value: SolanaPubkeyV2) -> Self {
        Self(value.to_bytes())
    }
}

#[cfg(feature = "solana-pubkey-convert-v2")]
impl From<Pubkey> for SolanaPubkeyV2 {
    /// Convert a [`Pubkey`] into a [`solana_pubkey::Pubkey`] by taking ownership of its bytes.
    #[inline]
    fn from(value: Pubkey) -> Self {
        Self::new_from_array(value.0)
    }
}

#[cfg(feature = "solana-address-convert-v1")]
impl From<Pubkey> for solana_address_v1::Address {
    fn from(value: Pubkey) -> Self {
        Self::new_from_array(value.0)
    }
}

#[cfg(feature = "solana-address-convert-v1")]
impl From<solana_address_v1::Address> for Pubkey {
    fn from(value: solana_address_v1::Address) -> Self {
        Self::new_from_array(value.to_bytes())
    }
}

#[cfg(feature = "solana-address-convert-v2")]
impl From<Pubkey> for solana_address_v2::Address {
    fn from(value: Pubkey) -> Self {
        Self::new_from_array(value.0)
    }
}

#[cfg(feature = "solana-address-convert-v2")]
impl From<solana_address_v2::Address> for Pubkey {
    fn from(value: solana_address_v2::Address) -> Self {
        Self::new_from_array(value.to_bytes())
    }
}

/// Error that occured while attempting to parse a base58 string into a [`Pubkey`].
///
/// Contains the same information as [`five8::DecodeError`], but has additional derives for easier
/// usage in downstream code.
#[cfg(feature = "pubkey-str")]
#[derive(Debug, Clone, Copy, PartialEq, thiserror::Error)]
pub enum DecodeError {
    /// A character not in the Base58 character set was encountered.
    #[error("Invalid base58 character encountered: char number {0}")]
    InvalidChar(u8),
    /// The input string was too long to be a valid 32-byte value.
    #[error("Base58 string too long")]
    TooLong,
    /// The input string was too short to be a valid 32-byte value.
    #[error("Base58 string too short")]
    TooShort,
    /// The largest term in the base58 encoded value is too large.
    #[error("Largest term greater than 2^32")]
    LargestTermTooHigh,
    /// The base58 string decoded to more than 32 bytes.
    #[error("Decoded output has too many bytes")]
    OutputTooLong,
}

/// Convert from the five8 error type to our version.
#[cfg(feature = "pubkey-str")]
impl From<five8::DecodeError> for DecodeError {
    fn from(value: five8::DecodeError) -> Self {
        match value {
            five8::DecodeError::InvalidChar(c) => Self::InvalidChar(c),
            five8::DecodeError::TooLong => Self::TooLong,
            five8::DecodeError::TooShort => Self::TooShort,
            five8::DecodeError::LargestTermTooHigh => Self::LargestTermTooHigh,
            five8::DecodeError::OutputTooLong => Self::OutputTooLong,
        }
    }
}

/// Allow conversion back to the five8 error type if needed.
#[cfg(feature = "pubkey-str")]
impl From<DecodeError> for five8::DecodeError {
    fn from(value: DecodeError) -> Self {
        match value {
            DecodeError::InvalidChar(c) => Self::InvalidChar(c),
            DecodeError::TooLong => Self::TooLong,
            DecodeError::TooShort => Self::TooShort,
            DecodeError::LargestTermTooHigh => Self::LargestTermTooHigh,
            DecodeError::OutputTooLong => Self::OutputTooLong,
        }
    }
}

#[cfg(feature = "pubkey-str")]
impl FromStr for Pubkey {
    type Err = DecodeError;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let mut out = [0u8; PUBKEY_BYTES];
        five8::decode_32(s, &mut out)?;
        Ok(Self::new_from_array(out))
    }
}

#[cfg(feature = "pubkey-str")]
impl core::fmt::Display for Pubkey {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let mut buf = [0u8; five8::BASE58_ENCODED_32_MAX_LEN];
        let encoded_len = five8::encode_32(self.as_array(), &mut buf);
        let as_str = std::str::from_utf8(&buf[..encoded_len as usize])
            .expect("base58 string should be valid utf8");
        f.write_str(as_str)
    }
}

#[cfg(test)]
mod test {
    use super::{Pubkey, PUBKEY_BYTES};

    const SOL_ADDR_BYTES: [u8; PUBKEY_BYTES] = [
        6, 155, 136, 87, 254, 171, 129, 132, 251, 104, 127, 99, 70, 24, 192, 53, 218, 196, 57, 220,
        26, 235, 59, 85, 152, 160, 240, 0, 0, 0, 0, 0,
    ];
    const SOL_ADDR_BASE58: &str = "So11111111111111111111111111111111111111111";

    #[test]
    fn test_construct() {
        let pubkey = Pubkey::new_from_array(SOL_ADDR_BYTES);
        assert_eq!(pubkey.as_array().len(), PUBKEY_BYTES);
        assert_eq!(pubkey.as_array(), &SOL_ADDR_BYTES);
    }

    #[cfg(feature = "pubkey-fromstr-const")]
    #[test]
    fn test_base58_const() {
        let expected = Pubkey::new_from_array(SOL_ADDR_BYTES);
        let decoded = Pubkey::from_str_const(SOL_ADDR_BASE58);

        assert_eq!(decoded, expected);

        let decoded_macro = pubkey!("So11111111111111111111111111111111111111111");
        assert_eq!(decoded_macro, expected);
    }

    #[cfg(feature = "pubkey-str")]
    #[test]
    fn test_from_str() {
        use std::str::FromStr;

        let expected = Pubkey::new_from_array(SOL_ADDR_BYTES);
        let decoded = Pubkey::from_str(SOL_ADDR_BASE58).expect("should decode from base58");

        assert_eq!(decoded, expected);
    }

    #[cfg(feature = "solana-pubkey-convert-v2")]
    #[test]
    fn test_convert_to_solana_pubkey_v2() {
        use solana_pubkey_v2::Pubkey as SolanaPubkey;

        let orig = SolanaPubkey::new_from_array(SOL_ADDR_BYTES);
        let ours = Pubkey::from(orig);
        assert_eq!(ours.as_array(), &SOL_ADDR_BYTES);
    }

    #[cfg(feature = "solana-address-convert-v1")]
    #[test]
    fn test_convert_to_solana_address_v1() {
        use solana_address_v1::Address;

        let orig = Address::new_from_array(SOL_ADDR_BYTES);
        let ours = Pubkey::from(orig);
        assert_eq!(ours.as_array(), &SOL_ADDR_BYTES);
    }

    #[cfg(feature = "serde")]
    #[test]
    fn test_serialize_messagepack() {
        let pubkey = Pubkey::new_from_array(SOL_ADDR_BYTES);
        let encoded = rmp_serde::encode::to_vec(&pubkey).expect("should encode via messagepack");

        // Should be encoded to binary with 2 header bytes for type and length.
        assert_eq!(encoded.len(), PUBKEY_BYTES + 2);
        // Header should be 0xc4 0x20 (bin 8 format with length 32)
        assert_eq!(encoded[0..2], [0xc4u8, 0x20u8]);
        assert_eq!(encoded[2..], SOL_ADDR_BYTES)
    }

    #[cfg(feature = "serde")]
    #[test]
    fn test_deserialize_messagepack() {
        let pubkey = Pubkey::new_from_array(SOL_ADDR_BYTES);
        let mut encoded = [0u8; PUBKEY_BYTES + 2];
        encoded[0] = 0xc4;
        encoded[1] = 32;
        encoded[2..].copy_from_slice(&SOL_ADDR_BYTES);

        let decoded = rmp_serde::decode::from_slice::<Pubkey>(&encoded).expect("should decode");
        assert_eq!(decoded, pubkey);
    }

    #[cfg(feature = "serde-pubkey-human-base58")]
    #[test]
    fn test_serde_pubkey_human() {
        use serde::{Deserialize, Serialize};

        #[derive(Serialize, Deserialize, PartialEq, Eq, Debug)]
        struct Test {
            key: Pubkey,
        }

        let value = Test {
            key: Pubkey::new_from_array(SOL_ADDR_BYTES),
        };
        let expected_json = serde_json::json!({
            "key": "So11111111111111111111111111111111111111111"
        });
        let as_json = serde_json::to_value(&value).expect("should encode to JSON");
        assert_eq!(as_json, expected_json);

        let from_json: Test =
            serde_json::from_value(expected_json).expect("should decode from JSON with string key");
        assert_eq!(from_json, value);

        let json_key_as_array = serde_json::json!({
            "key": SOL_ADDR_BYTES,
        });
        let from_json: Test = serde_json::from_value(json_key_as_array)
            .expect("should decode from JSON with array key");
        assert_eq!(from_json, value);
    }
}
