// LNP/BP lLibraries implementing LNPBP specifications & standards
// Written in 2021-2022 by
//     Dr. Maxim Orlovsky <orlovsky@pandoraprime.ch>
//
// To the extent possible under law, the author(s) have dedicated all
// copyright and related and neighboring rights to this software to
// the public domain worldwide. This software is distributed without
// any warranty.
//
// You should have received a copy of the MIT License
// along with this software.
// If not, see <https://opensource.org/licenses/MIT>.

// Coding conventions
#![recursion_limit = "256"]
#![deny(dead_code, missing_docs, warnings)]

//! Library implementing LNPBP-14 standard: Bech32 encoding for
//! client-side-validated data.
//!
//! Types that need to have `data1...` and `z1...` bech 32 implementation
//! according to LNPBP-14 must implement [`ToBech32Payload`] and
//! [`FromBech32Payload`] traits.
//!
//! Bech32 `id1...` representation is provided automatically only for hash types
//! implementing [`bitcoin_hashes::Hash`] trait

#[macro_use]
extern crate amplify;
#[macro_use]
extern crate strict_encoding;
#[cfg(feature = "serde")]
#[macro_use]
extern crate serde_crate as serde;

use std::convert::{Infallible, TryFrom};
use std::fmt::{self, Debug, Formatter};
use std::str::FromStr;

use amplify::hex::ToHex;
use bech32::{FromBase32, ToBase32, Variant};
use bitcoin_hashes::{sha256t, Hash};
#[cfg(feature = "zip")]
use deflate::{write::DeflateEncoder, Compression};
#[cfg(feature = "serde")]
use serde::{
    de::{Error as SerdeError, Unexpected, Visitor},
    Deserializer, Serializer,
};
#[cfg(feature = "serde")]
use serde_with::{hex::Hex, As};

/// Bech32 HRP used in generic identifiers
pub const HRP_ID: &str = "id";

/// Bech32 HRP used for representation of arbitrary data blobs in their raw
/// (uncompressed) form
pub const HRP_DATA: &str = "data";

#[cfg(feature = "zip")]
/// Bech32 HRP used for representation of zip-compressed blobs
pub const HRP_ZIP: &str = "z";

/// Constant specifying default compression algorithm ("deflate")
#[cfg(feature = "zip")]
pub const RAW_DATA_ENCODING_DEFLATE: u8 = 1u8;

/// Errors generated by Bech32 conversion functions (both parsing and
/// type-specific conversion errors)
#[derive(Clone, PartialEq, Eq, Display, Debug, From, Error)]
#[display(doc_comments)]
pub enum Error {
    /// bech32 string parse error - {0}
    #[from]
    Bech32Error(::bech32::Error),

    /// payload data are not strictly encoded - {0}
    #[from]
    NotStrictEncoded(strict_encoding::Error),

    /// payload data are not a bitcoin hash - {0}
    #[from]
    NotBitcoinHash(bitcoin_hashes::Error),

    /// Requested object type does not match used Bech32 HRP
    WrongPrefix,

    /// bech32m encoding must be used instead of legacy bech32
    WrongVariant,

    /// payload must start with encoding prefix
    NoEncodingPrefix,

    /// provided raw data use unknown encoding version {0}
    UnknownRawDataEncoding(u8),

    /// can not encode raw data with DEFLATE algorithm
    DeflateEncoding,

    /// error inflating compressed data from payload: {0}
    InflateError(String),
}

impl From<Infallible> for Error {
    fn from(_: Infallible) -> Self {
        unreachable!("infalliable error in lnpbp_bech32 blob")
    }
}

/// Type for wrapping Vec<u8> data in cases you need to do a convenient
/// enum variant display derives with `#[display(inner)]`
#[cfg_attr(
    feature = "serde",
    derive(Serialize, Deserialize),
    serde(crate = "serde_crate", transparent)
)]
#[derive(
    Wrapper, Clone, Ord, PartialOrd, Eq, PartialEq, Hash, Default, Display,
    From
)]
#[derive(StrictEncode, StrictDecode)]
#[wrap(
    Index,
    IndexMut,
    IndexRange,
    IndexFull,
    IndexFrom,
    IndexTo,
    IndexInclusive
)]
#[display(Vec::bech32_data_string)]
// We get `(To)Bech32DataString` and `FromBech32DataString` for free b/c
// the wrapper creates `From<Vec<u8>>` impl for us, which with rust stdlib
// implies `TryFrom<Vec<u8>>`, for which we have auto trait derivation
// `FromBech32Payload`, for which the traits above are automatically derived
pub struct Blob(
    #[cfg_attr(feature = "serde", serde(with = "As::<Hex>"))] Vec<u8>,
);

impl AsRef<[u8]> for Blob {
    fn as_ref(&self) -> &[u8] { &self.0 }
}

impl Debug for Blob {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        write!(f, "Blob({})", self.0.to_hex())
    }
}

impl FromStr for Blob {
    type Err = Error;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        Blob::from_bech32_data_str(s)
    }
}

/// Convertor trait for extracting data from a given type which will be part of
/// Bech32 payload
pub trait ToBech32Payload {
    /// Must return a vector with Bech32 payload data
    fn to_bech32_payload(&self) -> Vec<u8>;
}

/// Extracts data representation from a given type which will be part of Bech32
/// payload
pub trait AsBech32Payload {
    /// Must return a reference to a slice representing Bech32 payload data
    fn as_bech32_payload(&self) -> &[u8];
}

impl<T> AsBech32Payload for T
where
    T: AsRef<[u8]>,
{
    fn as_bech32_payload(&self) -> &[u8] { self.as_ref() }
}

/// Convertor which constructs a given type from Bech32 payload data
pub trait FromBech32Payload
where
    Self: Sized,
{
    /// Construct type from Bech32 payload data
    fn from_bech32_payload(payload: Vec<u8>) -> Result<Self, Error>;
}

impl<T> FromBech32Payload for T
where
    T: TryFrom<Vec<u8>>,
    Error: From<T::Error>,
{
    fn from_bech32_payload(payload: Vec<u8>) -> Result<T, Error> {
        Ok(T::try_from(payload)?)
    }
}

// -- Common (non-LNPBP-39) traits

/// Creates Bech32 string with appropriate type data representation.
/// Depending on the specific type, this may be `id`-string, `data`-string,
/// `z`-string or other type of HRP.
pub trait ToBech32String {
    /// Creates Bech32 string with appropriate type data representation
    fn to_bech32_string(&self) -> String;
}

/// Constructs type from the provided Bech32 string, or fails with
/// [`enum@Error`]
pub trait FromBech32Str {
    /// Specifies which HRP is used by Bech32 string representing this data type
    const HRP: &'static str;

    /// Constructs type from the provided Bech32 string, or fails with
    /// [`enum@Error`]
    fn from_bech32_str(s: &str) -> Result<Self, Error>
    where
        Self: Sized;
}

/// Strategies for automatic implementation of the Bech32 traits
pub mod strategies {
    use amplify::{Holder, Wrapper};
    use strict_encoding::{StrictDecode, StrictEncode};

    use super::*;

    /// Strategy for Bech32 representation as uncompressed data (starting from
    /// `data1...` HRP). The data are takken by using [`StrictEncode`]
    /// implementation defined for the type.
    pub struct UsingStrictEncoding;

    /// Strategy for Bech32 representation of the newtypes wrapping other types.
    /// The strategy simply inherits Bech32 representation from the inner type.
    pub struct Wrapped;

    #[cfg(feature = "zip")]
    /// Strategy for Bech32 representation as compressed data (starting from
    /// `z1...` HRP). The data are takken by using [`StrictEncode`]
    /// implementation defined for the type.
    pub struct CompressedStrictEncoding;

    /// Helper trait for implementing specific strategy for Bech32 construction
    pub trait Strategy {
        /// Bech32 HRP prefix used by a type
        const HRP: &'static str;
        /// Specific strategy used for automatic implementation of all
        /// Bech32-related traits.
        type Strategy;
    }

    impl<T> ToBech32String for T
    where
        T: Strategy + Clone,
        Holder<T, <T as Strategy>::Strategy>: ToBech32String,
    {
        #[inline]
        fn to_bech32_string(&self) -> String {
            Holder::new(self.clone()).to_bech32_string()
        }
    }

    impl<T> FromBech32Str for T
    where
        T: Strategy,
        Holder<T, <T as Strategy>::Strategy>: FromBech32Str,
    {
        const HRP: &'static str = T::HRP;

        #[inline]
        fn from_bech32_str(s: &str) -> Result<Self, Error> {
            Ok(Holder::from_bech32_str(s)?.into_inner())
        }
    }

    impl<T> ToBech32String for Holder<T, Wrapped>
    where
        T: Wrapper,
        T::Inner: ToBech32String,
    {
        #[inline]
        fn to_bech32_string(&self) -> String {
            self.as_inner().as_inner().to_bech32_string()
        }
    }

    impl<T> FromBech32Str for Holder<T, Wrapped>
    where
        T: Wrapper + Strategy,
        T::Inner: FromBech32Str,
    {
        const HRP: &'static str = T::HRP;

        #[inline]
        fn from_bech32_str(s: &str) -> Result<Self, Error> {
            Ok(Self::new(T::from_inner(T::Inner::from_bech32_str(s)?)))
        }
    }

    impl<T> ToBech32String for Holder<T, UsingStrictEncoding>
    where
        T: StrictEncode + Strategy,
    {
        #[inline]
        fn to_bech32_string(&self) -> String {
            let data = self
                .as_inner()
                .strict_serialize()
                .expect("in-memory strict encoding failure");
            ::bech32::encode(T::HRP, data.to_base32(), Variant::Bech32m)
                .unwrap_or_else(|_| s!("Error: wrong bech32 prefix"))
        }
    }

    impl<T> FromBech32Str for Holder<T, UsingStrictEncoding>
    where
        T: StrictDecode + Strategy,
    {
        const HRP: &'static str = T::HRP;

        #[inline]
        fn from_bech32_str(s: &str) -> Result<Self, Error> {
            let (hrp, data, variant) = ::bech32::decode(s)?;
            if hrp.as_str() != Self::HRP {
                return Err(Error::WrongPrefix);
            }
            if variant != Variant::Bech32m {
                return Err(Error::WrongVariant);
            }
            Ok(Self::new(T::strict_deserialize(Vec::<u8>::from_base32(
                &data,
            )?)?))
        }
    }
}
pub use strategies::Strategy;

// -- Sealed traits & their implementation

/// Special trait for preventing implementation of [`FromBech32DataStr`] and
/// others from outside of this crate. For details see
/// <https://rust-lang.github.io/api-guidelines/future-proofing.html#sealed-traits-protect-against-downstream-implementations-c-sealed>
mod sealed {
    use amplify::Wrapper;

    use super::*;

    pub trait HashType<Tag>: Wrapper<Inner = sha256t::Hash<Tag>>
    where
        Tag: sha256t::Tag,
    {
    }
    pub trait ToPayload: ToBech32Payload {}
    pub trait AsPayload: AsBech32Payload {}
    pub trait FromPayload: FromBech32Payload {}

    impl<T, Tag> HashType<Tag> for T
    where
        T: Wrapper<Inner = sha256t::Hash<Tag>>,
        Tag: sha256t::Tag,
    {
    }
    impl<T> ToPayload for T where T: ToBech32Payload {}
    impl<T> AsPayload for T where T: AsBech32Payload {}
    impl<T> FromPayload for T where T: FromBech32Payload {}
}

/// Trait for creating `data1...` Bech32 representation of a given type
pub trait ToBech32DataString: sealed::ToPayload {
    /// Returns `data1...` Bech32 representation of a given type
    fn to_bech32_data_string(&self) -> String {
        ::bech32::encode(
            HRP_DATA,
            self.to_bech32_payload().to_base32(),
            Variant::Bech32m,
        )
        .expect("HRP is hardcoded and can't fail")
    }
}

impl<T> ToBech32DataString for T where T: sealed::ToPayload {}

/// Trait for creating `data1...` Bech32 representation of a given type
pub trait Bech32DataString: sealed::AsPayload {
    /// Returns `data1...` Bech32 representation of a given type
    fn bech32_data_string(&self) -> String {
        ::bech32::encode(
            HRP_DATA,
            self.as_bech32_payload().to_base32(),
            Variant::Bech32m,
        )
        .expect("HRP is hardcoded and can't fail")
    }
}

impl<T> Bech32DataString for T where T: sealed::AsPayload {}

/// Trait for reconstruction type data from `data1...` Bech32 string
pub trait FromBech32DataStr
where
    Self: Sized + sealed::FromPayload,
{
    /// Reconstructs type data from `data1...` Bech32 string
    fn from_bech32_data_str(s: &str) -> Result<Self, Error> {
        let (hrp, data, variant) = bech32::decode(s)?;
        if hrp != HRP_DATA {
            return Err(Error::WrongPrefix);
        }
        if variant != Variant::Bech32m {
            return Err(Error::WrongVariant);
        }
        Self::from_bech32_payload(Vec::<u8>::from_base32(&data)?)
    }
}

impl<T> FromBech32DataStr for T where T: sealed::FromPayload {}

#[doc(hidden)]
#[cfg(feature = "zip")]
pub mod zip {
    use amplify::Holder;
    use strict_encoding::{StrictDecode, StrictEncode};

    use super::*;

    fn payload_to_bech32_zip_string(hrp: &str, payload: &[u8]) -> String {
        use std::io::Write;

        // We initialize writer with a version byte, indicating deflation
        // algorithm used
        let writer = vec![RAW_DATA_ENCODING_DEFLATE];
        let mut encoder = DeflateEncoder::new(writer, Compression::Best);
        encoder
            .write_all(payload)
            .expect("in-memory strict encoder failure");
        let data = encoder.finish().expect("zip algorithm failure");

        ::bech32::encode(hrp, data.to_base32(), Variant::Bech32m)
            .expect("HRP is hardcoded and can't fail")
    }

    fn bech32_zip_str_to_payload(hrp: &str, s: &str) -> Result<Vec<u8>, Error> {
        let (prefix, data, version) = bech32::decode(s)?;
        if prefix != hrp {
            return Err(Error::WrongPrefix);
        }
        if version != Variant::Bech32m {
            return Err(Error::WrongVariant);
        }
        let data = Vec::<u8>::from_base32(&data)?;
        match *data[..].first().ok_or(Error::NoEncodingPrefix)? {
            RAW_DATA_ENCODING_DEFLATE => {
                let decoded = inflate::inflate_bytes(&data[1..])
                    .map_err(Error::InflateError)?;
                Ok(decoded)
            }
            unknown_ver => Err(Error::UnknownRawDataEncoding(unknown_ver)),
        }
    }

    /// Trait for creating `z1...` (compressed binary data blob) Bech32
    /// representation of a given type
    pub trait ToBech32ZipString: sealed::ToPayload {
        /// Returns `z1...` (compressed binary data blob) Bech32 representation
        /// of a given type
        fn to_bech32_zip_string(&self) -> String {
            payload_to_bech32_zip_string(HRP_ZIP, &self.to_bech32_payload())
        }
    }

    impl<T> ToBech32ZipString for T where T: sealed::ToPayload {}

    /// Trait for creating `z1...` (compressed binary data blob) Bech32
    /// representation of a given type
    pub trait Bech32ZipString: sealed::AsPayload {
        /// Returns `z1...` (compressed binary data blob) Bech32 representation
        /// of a given type
        fn bech32_zip_string(&self) -> String {
            payload_to_bech32_zip_string(HRP_ZIP, self.as_bech32_payload())
        }
    }

    impl<T> Bech32ZipString for T where T: sealed::AsPayload {}

    /// Trait for reconstruction type data from `z1...` (compressed binary data
    /// blob) Bech32 string
    pub trait FromBech32ZipStr: sealed::FromPayload {
        /// Reconstructs type data from `z1...` (compressed binary data blob)
        /// Bech32 string
        fn from_bech32_zip_str(s: &str) -> Result<Self, Error> {
            Self::from_bech32_payload(bech32_zip_str_to_payload(HRP_ZIP, s)?)
        }
    }

    impl<T> FromBech32ZipStr for T where T: sealed::FromPayload {}

    impl<T> ToBech32String for Holder<T, strategies::CompressedStrictEncoding>
    where
        T: StrictEncode + Strategy,
    {
        #[inline]
        fn to_bech32_string(&self) -> String {
            let data = self
                .as_inner()
                .strict_serialize()
                .expect("in-memory strict encoding failure");
            payload_to_bech32_zip_string(T::HRP, &data)
        }
    }

    impl<T> FromBech32Str for Holder<T, strategies::CompressedStrictEncoding>
    where
        T: StrictDecode + Strategy,
    {
        const HRP: &'static str = T::HRP;

        #[inline]
        fn from_bech32_str(s: &str) -> Result<Self, Error> {
            Ok(Self::new(T::strict_deserialize(
                bech32_zip_str_to_payload(Self::HRP, s)?,
            )?))
        }
    }
}
#[cfg(feature = "zip")]
pub use zip::*;

/// Trait representing given bitcoin hash type as a Bech32 `id1...` value
pub trait ToBech32IdString<Tag>
where
    Self: sealed::HashType<Tag>,
    Tag: sha256t::Tag,
{
    /// Returns Bech32-encoded string in form of `id1...` representing the type
    fn to_bech32_id_string(&self) -> String;
}

/// Trait that can generate the type from a given Bech32 `id1...` value
pub trait FromBech32IdStr<Tag>
where
    Self: sealed::HashType<Tag> + Sized,
    Tag: sha256t::Tag,
{
    /// Reconstructs the identifier type from the provided Bech32 `id1...`
    /// string
    fn from_bech32_id_str(s: &str) -> Result<Self, Error>;
}

impl<T, Tag> ToBech32IdString<Tag> for T
where
    Self: sealed::HashType<Tag>,
    Tag: sha256t::Tag,
{
    fn to_bech32_id_string(&self) -> String {
        ::bech32::encode(HRP_ID, self.to_inner().to_base32(), Variant::Bech32m)
            .expect("HRP is hardcoded and can't fail")
    }
}

impl<T, Tag> FromBech32IdStr<Tag> for T
where
    Self: sealed::HashType<Tag>,
    Tag: sha256t::Tag,
{
    fn from_bech32_id_str(s: &str) -> Result<T, Error> {
        let (hrp, id, variant) = ::bech32::decode(s)?;
        if hrp != HRP_ID {
            return Err(Error::WrongPrefix);
        }
        if variant != Variant::Bech32m {
            return Err(Error::WrongVariant);
        }
        let vec = Vec::<u8>::from_base32(&id)?;
        Ok(Self::from_inner(Self::Inner::from_slice(&vec)?))
    }
}

/// Helper method for serde serialization of types supporting Bech32
/// representation
#[cfg(feature = "serde")]
pub fn serialize<T, S>(data: &T, serializer: S) -> Result<S::Ok, S::Error>
where
    S: Serializer,
    T: ToBech32String,
{
    serializer.serialize_str(&data.to_bech32_string())
}

/// Helper method for serde deserialization of types supporting Bech32
/// representation
#[cfg(feature = "serde")]
pub fn deserialize<'de, T, D>(deserializer: D) -> Result<T, D::Error>
where
    D: Deserializer<'de>,
    T: FromBech32Str,
{
    deserializer.deserialize_str(Bech32Visitor::<T>(std::marker::PhantomData))
}

#[cfg(feature = "serde")]
struct Bech32Visitor<Value>(std::marker::PhantomData<Value>);

#[cfg(feature = "serde")]
impl<'de, ValueT> Visitor<'de> for Bech32Visitor<ValueT>
where
    ValueT: FromBech32Str,
{
    type Value = ValueT;

    fn expecting(
        &self,
        formatter: &mut std::fmt::Formatter,
    ) -> std::fmt::Result {
        formatter.write_str("a bech32m-encoded string")
    }

    fn visit_str<E>(self, v: &str) -> Result<Self::Value, E>
    where
        E: SerdeError,
    {
        Self::Value::from_bech32_str(v).map_err(|_| {
            E::invalid_value(Unexpected::Str(v), &"valid bech32 string")
        })
    }
}