use bitvec::{mem::bits_of, view::AsBits};
use num_bigint::BigUint;
use tlb_ton::{
    Cell, EitherInlineOrRef, Error, MsgAddress, ParseFully, Ref, Same,
    bits::{
        NoArgs, Remainder, VarInt,
        de::{BitReader, BitReaderExt, BitUnpack},
        integer::ConstU32,
        ser::{BitPack, BitWriter, BitWriterExt},
    },
    de::{CellDeserialize, CellParser, CellParserError},
    either::Either,
    ser::{CellBuilder, CellBuilderError, CellSerialize},
};

/// Jetton Transfer message from [TEP-74](https://github.com/ton-blockchain/TEPs/blob/master/text/0074-jettons-standard.md#tl-b-schema)
/// ```tlb
/// transfer#0f8a7ea5 query_id:uint64 amount:(VarUInteger 16) destination:MsgAddress
/// response_destination:MsgAddress custom_payload:(Maybe ^Cell)
/// forward_ton_amount:(VarUInteger 16) forward_payload:(Either Cell ^Cell)
/// = InternalMsgBody;
/// ```
#[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct JettonTransfer<P = Cell, F = Cell> {
    pub query_id: u64,
    pub amount: BigUint,
    pub dst: MsgAddress,
    pub response_dst: MsgAddress,
    pub custom_payload: Option<P>,
    pub forward_ton_amount: BigUint,
    pub forward_payload: ForwardPayload<F>,
}

const JETTON_TRANSFER_TAG: u32 = 0x0f8a7ea5;

impl<P, F> CellSerialize for JettonTransfer<P, F>
where
    P: CellSerialize<Args: NoArgs>,
    F: CellSerialize<Args: NoArgs>,
{
    type Args = ();

    fn store(&self, builder: &mut CellBuilder, _: Self::Args) -> Result<(), CellBuilderError> {
        builder
            // transfer#0f8a7ea5
            .pack(JETTON_TRANSFER_TAG, ())?
            // query_id:uint64
            .pack(self.query_id, ())?
            // amount:(VarUInteger 16)
            .pack_as::<_, &VarInt<4>>(&self.amount, ())?
            // destination:MsgAddress
            .pack(self.dst, ())?
            // response_destination:MsgAddress
            .pack(self.response_dst, ())?
            // custom_payload:(Maybe ^Cell)
            .store_as::<_, Option<Ref>>(self.custom_payload.as_ref(), NoArgs::EMPTY)?
            // forward_ton_amount:(VarUInteger 16)
            .pack_as::<_, &VarInt<4>>(&self.forward_ton_amount, ())?
            // forward_payload:(Either Cell ^Cell)
            .store_as::<_, EitherInlineOrRef>(&self.forward_payload, NoArgs::EMPTY)?;

        Ok(())
    }
}

impl<'de, P, F> CellDeserialize<'de> for JettonTransfer<P, F>
where
    P: CellDeserialize<'de, Args: NoArgs>,
    F: CellDeserialize<'de, Args: NoArgs>,
{
    type Args = ();

    fn parse(parser: &mut CellParser<'de>, _: Self::Args) -> Result<Self, CellParserError<'de>> {
        // transfer#0f8a7ea5
        parser.unpack::<ConstU32<JETTON_TRANSFER_TAG>>(())?;
        Ok(Self {
            // query_id:uint64
            query_id: parser.unpack(())?,
            // amount:(VarUInteger 16)
            amount: parser.unpack_as::<_, VarInt<4>>(())?,
            // destination:MsgAddress
            dst: parser.unpack(())?,
            // response_destination:MsgAddress
            response_dst: parser.unpack(())?,
            // custom_payload:(Maybe ^Cell)
            custom_payload: parser.parse_as::<_, Option<Ref<ParseFully>>>(P::Args::EMPTY)?,
            // forward_ton_amount:(VarUInteger 16)
            forward_ton_amount: parser.unpack_as::<_, VarInt<4>>(())?,
            // forward_payload:(Either Cell ^Cell)
            forward_payload: parser
                .parse_as::<Either<ForwardPayload<F>, ForwardPayload<F>>, Either<ParseFully, Ref<ParseFully>>>(NoArgs::EMPTY)?
                .into_inner(),
        })
    }
}

#[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ForwardPayload<T = Cell> {
    Data(T),
    Comment(ForwardPayloadComment),
}

impl<T> ForwardPayload<T> {
    const COMMENT_PREFIX: u32 = 0x00000000;
}

impl<T> CellSerialize for ForwardPayload<T>
where
    T: CellSerialize<Args: NoArgs>,
{
    type Args = ();

    #[inline]
    fn store(&self, builder: &mut CellBuilder, _: Self::Args) -> Result<(), CellBuilderError> {
        match self {
            Self::Data(data) => builder.store(data, NoArgs::EMPTY)?,
            Self::Comment(comment) => builder.pack(Self::COMMENT_PREFIX, ())?.pack(comment, ())?,
        };
        Ok(())
    }
}

impl<'de, F> CellDeserialize<'de> for ForwardPayload<F>
where
    F: CellDeserialize<'de, Args: NoArgs>,
{
    type Args = ();

    fn parse(parser: &mut CellParser<'de>, _: Self::Args) -> Result<Self, CellParserError<'de>> {
        if parser.bits_left() >= bits_of::<u32>()
            // clone, so we don't advance original parser
            && parser.clone().unpack::<u32>(())? == Self::COMMENT_PREFIX
        {
            // skip the prefix
            let _ = parser.unpack::<u32>(())?;
            return parser.unpack(()).map(Self::Comment);
        }
        parser.parse(NoArgs::EMPTY).map(Self::Data)
    }
}

#[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ForwardPayloadComment {
    Text(String),
    Binary(Vec<u8>),
}

impl ForwardPayloadComment {
    const BINARY_PREFIX: u8 = 0xff;
}

impl BitPack for ForwardPayloadComment {
    type Args = ();

    #[inline]
    fn pack<W>(&self, writer: &mut W, _: Self::Args) -> Result<(), W::Error>
    where
        W: BitWriter + ?Sized,
    {
        match self {
            Self::Text(comment) => writer.write_bitslice(comment.as_bytes().as_bits())?,
            Self::Binary(comment) => writer
                .pack(Self::BINARY_PREFIX, ())?
                .write_bitslice(comment.as_bits())?,
        };
        Ok(())
    }
}

impl<'de> BitUnpack<'de> for ForwardPayloadComment {
    type Args = ();

    #[inline]
    fn unpack<R>(reader: &mut R, _: Self::Args) -> Result<Self, R::Error>
    where
        R: BitReader<'de> + ?Sized,
    {
        let mut r = reader.checkpoint();
        if r.bits_left() >= bits_of::<u8>() && r.unpack::<u8>(())? == Self::BINARY_PREFIX {
            return r.unpack_as::<_, Remainder>(()).map(Self::Binary);
        }
        r.restore()
            .unpack_as::<_, Remainder>(())
            .map(Self::Text)
            .map_err(Error::custom)
    }
}

/// Jetton Transfer Notification message from[TEP-74](https://github.com/ton-blockchain/TEPs/blob/master/text/0074-jettons-standard.md#tl-b-schema)
/// ```tlb
/// transfer_notification#7362d09c query_id:uint64 amount:(VarUInteger 16)
/// sender:MsgAddress forward_payload:(Either Cell ^Cell)
/// = InternalMsgBody;
/// ```
#[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct JettonTransferNotification<P = Cell> {
    pub query_id: u64,
    pub amount: BigUint,
    pub sender: MsgAddress,
    pub forward_payload: ForwardPayload<P>,
}

const JETTON_TRANSFER_NOTIFICATION_TAG: u32 = 0x7362d09c;

impl<P> CellSerialize for JettonTransferNotification<P>
where
    P: CellSerialize<Args: NoArgs>,
{
    type Args = ();

    fn store(&self, builder: &mut CellBuilder, _: Self::Args) -> Result<(), CellBuilderError> {
        builder
            .pack(JETTON_TRANSFER_NOTIFICATION_TAG, ())?
            .pack(self.query_id, ())?
            .pack_as::<_, &VarInt<4>>(&self.amount, ())?
            .pack(self.sender, ())?
            .store_as::<Either<(), _>, Either<Same, Ref>>(
                Either::Right(&self.forward_payload),
                NoArgs::EMPTY,
            )?;
        Ok(())
    }
}

impl<'de, P> CellDeserialize<'de> for JettonTransferNotification<P>
where
    P: CellDeserialize<'de, Args: NoArgs>,
{
    type Args = ();

    fn parse(parser: &mut CellParser<'de>, _: Self::Args) -> Result<Self, CellParserError<'de>> {
        parser.unpack::<ConstU32<JETTON_TRANSFER_NOTIFICATION_TAG>>(())?;
        Ok(Self {
            query_id: parser.unpack(())?,
            amount: parser.unpack_as::<_, VarInt<4>>(())?,
            sender: parser.unpack(())?,
            forward_payload: parser
                .parse_as::<Either<ForwardPayload<P>, ForwardPayload<P>>, Either<Same, Ref<ParseFully>>>(NoArgs::EMPTY)?
                .into_inner(),
        })
    }
}

/// Jetton Burn message from [TEP-74](https://github.com/ton-blockchain/TEPs/blob/master/text/0074-jettons-standard.md#tl-b-schema)
/// ```tlb
/// burn#595f07bc query_id:uint64 amount:(VarUInteger 16)
/// response_destination:MsgAddress custom_payload:(Maybe ^Cell)
/// = InternalMsgBody;
/// ```
#[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct JettonBurn<P = Cell> {
    pub query_id: u64,
    pub amount: BigUint,
    pub response_dst: MsgAddress,
    pub custom_payload: Option<P>,
}

const JETTON_BURN_TAG: u32 = 0x595f07bc;

impl<P> CellSerialize for JettonBurn<P>
where
    P: CellSerialize<Args: NoArgs>,
{
    type Args = ();

    fn store(&self, builder: &mut CellBuilder, _: Self::Args) -> Result<(), CellBuilderError> {
        builder
            .pack(JETTON_BURN_TAG, ())?
            .pack_as::<_, &VarInt<4>>(&self.amount, ())?
            .pack(self.response_dst, ())?
            .store_as::<_, Option<Ref>>(self.custom_payload.as_ref(), NoArgs::EMPTY)?;
        Ok(())
    }
}

impl<'de, P> CellDeserialize<'de> for JettonBurn<P>
where
    P: CellDeserialize<'de, Args: NoArgs>,
{
    type Args = ();

    fn parse(parser: &mut CellParser<'de>, _: Self::Args) -> Result<Self, CellParserError<'de>> {
        parser.unpack::<ConstU32<JETTON_BURN_TAG>>(())?;
        Ok(Self {
            query_id: parser.unpack(())?,
            amount: parser.unpack_as::<_, VarInt<4>>(())?,
            response_dst: parser.unpack(())?,
            custom_payload: parser.parse_as::<_, Option<Ref<ParseFully>>>(NoArgs::EMPTY)?,
        })
    }
}
