rostrum 14.0.1

use bitcoincash::consensus::encode::CheckedData;
use bitcoincash::consensus::serialize;
use bitcoincash::consensus::{encode, Decodable, Encodable};
use bitcoincash::network::message::CommandString;
use bitcoincash::network::Address;
use bitcoincash::VarInt;

use super::block::{Block, BlockHeader};
use super::transaction::Transaction;

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct RawNetworkMessage {
    /// Magic bytes to identify the network these messages are meant for
    pub magic: u32,
    /// The actual message data
    pub payload: NetworkMessage,
}

impl RawNetworkMessage {
    /// Return the message command as a static string reference.
    ///
    /// This returns `"unknown"` for [NetworkMessage::Unknown],
    /// regardless of the actual command in the unknown message.
    /// Use the [Self::command] method to get the command for unknown messages.
    pub fn cmd(&self) -> &'static str {
        self.payload.cmd()
    }

    /// Return the CommandString for the message command.
    pub fn command(&self) -> CommandString {
        self.payload.command()
    }
}

struct HeaderSerializationWrapper<'a>(&'a Vec<BlockHeader>);

impl Encodable for HeaderSerializationWrapper<'_> {
    #[inline]
    fn consensus_encode<W: std::io::Write + ?Sized>(
        &self,
        w: &mut W,
    ) -> Result<usize, std::io::Error> {
        let mut len = 0;
        len += VarInt(self.0.len() as u64).consensus_encode(w)?;
        for header in self.0.iter() {
            len += header.consensus_encode(w)?;
            len += 0u8.consensus_encode(w)?;
        }
        Ok(len)
    }
}

impl Encodable for RawNetworkMessage {
    fn consensus_encode<W: std::io::Write + ?Sized>(
        &self,
        w: &mut W,
    ) -> Result<usize, std::io::Error> {
        let mut len = 0;
        len += self.magic.consensus_encode(w)?;
        len += self.command().consensus_encode(w)?;
        len += CheckedData(match &self.payload {
            NetworkMessage::Version(ref dat) => serialize(dat),
            NetworkMessage::Inv(ref dat) => serialize(dat),
            NetworkMessage::GetData(ref dat) => serialize(dat),
            NetworkMessage::GetHeaders(ref dat) => serialize(dat),
            NetworkMessage::Tx(ref dat) => serialize(dat),
            NetworkMessage::Block(ref dat) => serialize(dat),
            NetworkMessage::Headers(dat) => serialize(&HeaderSerializationWrapper(dat)),
            NetworkMessage::Ping(ref dat) => serialize(dat),
            NetworkMessage::Pong(ref dat) => serialize(dat),
            NetworkMessage::Reject(ref dat) => serialize(dat),
            NetworkMessage::Addr(ref dat) => serialize(dat),
            NetworkMessage::Verack => vec![],
            NetworkMessage::Unknown {
                payload: ref data, ..
            } => serialize(data),
        })
        .consensus_encode(w)?;
        Ok(len)
    }
}

impl Decodable for RawNetworkMessage {
    fn consensus_decode_from_finite_reader<R: std::io::Read + ?Sized>(
        r: &mut R,
    ) -> Result<Self, encode::Error> {
        let magic = Decodable::consensus_decode_from_finite_reader(r)?;
        let cmd = CommandString::consensus_decode_from_finite_reader(r)?;
        let raw_payload = CheckedData::consensus_decode_from_finite_reader(r)?.0;

        let mut mem_d = std::io::Cursor::new(raw_payload);
        let payload = match cmd.as_ref() {
            "version" => {
                NetworkMessage::Version(Decodable::consensus_decode_from_finite_reader(&mut mem_d)?)
            }
            "verack" => NetworkMessage::Verack,
            "inv" => {
                NetworkMessage::Inv(Decodable::consensus_decode_from_finite_reader(&mut mem_d)?)
            }
            "getdata" => {
                NetworkMessage::GetData(Decodable::consensus_decode_from_finite_reader(&mut mem_d)?)
            }
            "block" => {
                NetworkMessage::Block(Decodable::consensus_decode_from_finite_reader(&mut mem_d)?)
            }
            "ping" => {
                NetworkMessage::Ping(Decodable::consensus_decode_from_finite_reader(&mut mem_d)?)
            }
            "pong" => {
                NetworkMessage::Pong(Decodable::consensus_decode_from_finite_reader(&mut mem_d)?)
            }
            "tx" => NetworkMessage::Tx(Decodable::consensus_decode_from_finite_reader(&mut mem_d)?),
            "reject" => {
                NetworkMessage::Reject(Decodable::consensus_decode_from_finite_reader(&mut mem_d)?)
            }
            _ => NetworkMessage::Unknown {
                command: cmd,
                payload: mem_d.into_inner(),
            },
        };
        Ok(RawNetworkMessage { magic, payload })
    }

    // #[inline]
    // fn consensus_decode<R: std::io::Read + ?Sized>(r: &mut R) -> Result<Self, encode::Error> {
    //     Self::consensus_decode_from_finite_reader(r.take(MAX_MSG_SIZE as u64).by_ref())
    // }
}

#[derive(Clone, PartialEq, Eq, Debug)]
pub enum NetworkMessage {
    /// `version`
    Version(bitcoincash::network::message_network::VersionMessage),
    /// `verack`
    Verack,
    /// `inv`
    Inv(Vec<bitcoincash::network::message_blockdata::Inventory>),
    /// `getdata`
    GetData(Vec<bitcoincash::network::message_blockdata::Inventory>),
    // `getheaders`
    GetHeaders(bitcoincash::network::message_blockdata::GetHeadersMessage),
    /// `notfound`
    Tx(Transaction),
    /// `block`
    Block(Block),
    /// `headers`
    Headers(Vec<BlockHeader>),
    /// `ping`
    Ping(u64),
    /// `pong`
    Pong(u64),
    /// `reject`
    Reject(bitcoincash::network::message_network::Reject),
    /// `feefilter`
    Addr(Vec<(u32, Address)>),
    /// Any other message.
    Unknown {
        /// The command of this message.
        command: bitcoincash::network::message::CommandString,
        /// The payload of this message.
        payload: Vec<u8>,
    },
}

impl NetworkMessage {
    /// Return the message command as a static string reference.
    ///
    /// This returns `"unknown"` for [NetworkMessage::Unknown],
    /// regardless of the actual command in the unknown message.
    /// Use the [Self::command] method to get the command for unknown messages.
    pub fn cmd(&self) -> &'static str {
        match *self {
            NetworkMessage::Version(_) => "version",
            NetworkMessage::Verack => "verack",
            NetworkMessage::Inv(_) => "inv",
            NetworkMessage::GetData(_) => "getdata",
            NetworkMessage::GetHeaders(_) => "getheaders",
            NetworkMessage::Tx(_) => "tx",
            NetworkMessage::Block(_) => "block",
            NetworkMessage::Headers(_) => "headers",
            NetworkMessage::Ping(_) => "ping",
            NetworkMessage::Pong(_) => "pong",
            NetworkMessage::Reject(_) => "reject",
            NetworkMessage::Addr(_) => "addr",
            NetworkMessage::Unknown { .. } => "unknown",
        }
    }

    /// Return the CommandString for the message command.
    pub fn command(&self) -> CommandString {
        match *self {
            NetworkMessage::Unknown { command: ref c, .. } => c.clone(),
            _ => CommandString::try_from_static(self.cmd()).expect("cmd returns valid commands"),
        }
    }
}