use alloc::{borrow::Cow, vec::Vec};

/// A JAM-specific program blob.
pub struct ProgramBlob<'a> {
	pub ro_data: Cow<'a, [u8]>,
	pub rw_data: Cow<'a, [u8]>,
	pub code_blob: Cow<'a, [u8]>,
	pub rw_data_padding: u32,
	pub stack_size: u32,
}

fn read_u24(bytes: &mut &[u8]) -> Option<u32> {
	let xs = bytes.get(..3)?;
	*bytes = &bytes[3..];
	Some(u32::from_le_bytes([xs[0], xs[1], xs[2], 0]))
}

fn write_u24(value: u32, output: &mut Vec<u8>) -> Result<(), ()> {
	if value >= (1 << 24) {
		return Err(());
	}

	output.extend_from_slice(&value.to_le_bytes()[0..3]);
	Ok(())
}

#[test]
fn read_write_u24() {
	let mut output = Vec::new();
	write_u24(0x00345678, &mut output).unwrap();
	assert_eq!(read_u24(&mut &output[..]), Some(0x00345678));

	assert!(write_u24(0x00ffffff, &mut output).is_ok());
	assert!(write_u24(0x01000000, &mut output).is_err());
}

fn read_u32(bytes: &mut &[u8]) -> Option<u32> {
	let xs = bytes.get(..4)?;
	*bytes = &bytes[4..];
	Some(u32::from_le_bytes([xs[0], xs[1], xs[2], xs[3]]))
}

fn read_cow<'a>(bytes: &mut &'a [u8], length: u32) -> Option<Cow<'a, [u8]>> {
	let length = length as usize;
	let cow = bytes.get(..length)?;
	*bytes = &bytes[length..];
	Some(Cow::Borrowed(cow))
}

impl<'a> ProgramBlob<'a> {
	pub fn from_bytes(mut bytes: &'a [u8]) -> Option<Self> {
		let ro_data_len = read_u24(&mut bytes)?;
		let rw_data_len = read_u24(&mut bytes)?;
		let rw_data_padding = read_u24(&mut bytes)?;
		let stack_size = read_u24(&mut bytes)?;
		let ro_data = read_cow(&mut bytes, ro_data_len)?;
		let rw_data = read_cow(&mut bytes, rw_data_len)?;
		let code_blob_len = read_u32(&mut bytes)?;
		let code_blob = read_cow(&mut bytes, code_blob_len)?;

		if !bytes.is_empty() {
			return None;
		}

		Some(ProgramBlob { rw_data_padding, stack_size, ro_data, rw_data, code_blob })
	}

	pub fn to_vec(&self) -> Result<Vec<u8>, &'static str> {
		let mut output = Vec::new();
		write_u24(u32::try_from(self.ro_data.len()).map_err(|_| "too large RO data")?, &mut output)
			.map_err(|_| "too large RO data")?;
		write_u24(u32::try_from(self.rw_data.len()).map_err(|_| "too large RW data")?, &mut output)
			.map_err(|_| "too large RW data")?;
		write_u24(self.rw_data_padding, &mut output).map_err(|_| "too large RW data padding")?;
		write_u24(self.stack_size, &mut output).map_err(|_| "too large stack size")?;
		output.extend_from_slice(&self.ro_data);
		output.extend_from_slice(&self.rw_data);
		output.extend_from_slice(
			&u32::try_from(self.code_blob.len()).map_err(|_| "too large code")?.to_le_bytes(),
		);
		output.extend_from_slice(&self.code_blob);
		Ok(output)
	}
}
