// A crummy ad-hoc TLV cause I'm lazy and these are fun

#[derive (Debug, PartialEq)]
pub enum Message {
	IAmOnline {
		peer_id: u32,
	},
	IHaveMessage {
		peer_id: u32,
		msg_index: u32,
		body: Vec <u8>,
	},
	IAcknowledgeYourMessage {
		msg_index: u32,
	},
}

#[derive (Debug, thiserror::Error)]
pub enum Error {
	#[error ("Bad magic number")]
	BadMagic,
	#[error ("Data is too big to encode / decode")]
	DataTooBig,
	#[error ("Required field missing while decoding")]
	MissingField,
	#[error ("No type in TLV")]
	NoType,
	#[error ("Unrecognized TLV type")]
	UnrecognizedType,
}

const MAGIC: [u8; 4] = [55, 11, 101, 38];

pub fn decode (src: &[u8]) -> Result <Message, Error>
{
	let (magic, src) = take_bytes (src, 4).ok_or (Error::BadMagic)?;
	if magic != MAGIC {
		return Err (Error::BadMagic);
	}
	
	let (msg_type, src) = take_bytes(src, 1).ok_or(Error::NoType)?;
	
	Ok (match msg_type [0] {
		1 => {
			let (peer_id, src) = take_u32 (src).ok_or (Error::MissingField)?;
			Message::IAmOnline { peer_id, }
		},
		2 => {
			let (peer_id, src) = take_u32 (src).ok_or (Error::MissingField)?;
			let (msg_index, src) = take_u32 (src).ok_or (Error::MissingField)?;
			let (body_len, src) = take_u32 (src).ok_or (Error::MissingField)?;
			let (body_bytes, src) = take_bytes(src, body_len as usize).ok_or(Error::MissingField)?;
			let body = body_bytes.into ();
			Message::IHaveMessage { peer_id, msg_index, body, }
		},
		3 => {
			let (msg_index, src) = take_u32 (src).ok_or (Error::MissingField)?;
			Message::IAcknowledgeYourMessage { msg_index, }
		},
		_ => return Err (Error::UnrecognizedType),
	})
}

fn take_u32 (src: &[u8]) -> Option <(u32, &[u8])>
{
	let (a, b) = take_bytes (src, 4)?;
	Some ((u32::from_le_bytes([a [0], a [1], a [2], a [3]]), b))
}

fn take_bytes (src: &[u8], n: usize) -> Option <(&[u8], &[u8])>
{
	if src.len () < n {
		return None;
	}
	
	Some (src.split_at (n))
}

pub fn encode (msg: &Message) -> Result <Vec <u8>, Error>
{
	let mut buf = Vec::from (MAGIC);
	
	match msg {
		Message::IAmOnline { peer_id } => {
			buf.push (1);
			buf.extend_from_slice (&peer_id.to_le_bytes());
		},
		Message::IHaveMessage { peer_id, msg_index, body } => {
			buf.push (2);
			buf.extend_from_slice (&peer_id.to_le_bytes());
			buf.extend_from_slice (&msg_index.to_le_bytes());
			
			let body_len = u32::try_from (body.len ()).or (Err (Error::DataTooBig))?;
			buf.extend_from_slice (&body_len.to_le_bytes());
			buf.extend_from_slice(body.as_slice());
		},
		Message::IAcknowledgeYourMessage { msg_index } => {
			buf.push (3);
			buf.extend_from_slice (&msg_index.to_le_bytes());
		},
	}
	
	Ok (buf)
}

#[cfg (test)]
mod test {
	use super::*;
	
	#[test]
	fn roundtrip () {
		for msg in [
			Message::IAmOnline { peer_id: 93 },
			Message::IHaveMessage { peer_id: 93, msg_index: 1000, body: Vec::from (b":V".as_slice()) },
			Message::IAcknowledgeYourMessage { msg_index: 1000 },
		] {
			let encoded = encode (&msg).unwrap ();
			let roundtripped = decode (&encoded).unwrap ();
			assert_eq!(roundtripped, msg);
		}
	}
}