use binparse_derive::AllVariants;
use binparse_derive::ParseBin;
use std::fmt;
use std::error;
use std::ffi::CString;

/***** BTRFS structures *****/

pub const NODE_SIZE: usize = 0x4000;

#[allow(unused)]
#[derive(Debug,Clone,Copy,AllVariants,PartialEq,Eq,PartialOrd,Ord)]
#[repr(u8)]
pub enum ItemType {
	Invalid = 0x00,           // invalid
	Inode = 0x01,             // implemented
	Ref = 0x0c,               // implemented
	ExtRef = 0x0d,
	XAttr = 0x18,
	VerityDesc = 0x24,
	VerityMerkle = 0x25,
	Orphan = 0x30,
	DirLog = 0x3c,
	DirLogIndex = 0x48,
	Dir = 0x54,               // implemented (better with len feature; allow multiple?)
	DirIndex = 0x60,          // implemented
	ExtentData = 0x6c,        // implemented
	ExtentCsum = 0x80,
	Root = 0x84,              // implemented
	RootBackRef = 0x90,
	RootRef = 0x9c,
	Extent = 0xa8,            // implemented (with only one version of extra data)
	Metadata = 0xa9,          // implemented (with only one version of extra data)
	TreeBlockRef = 0xb0,
	ExtentDataRef = 0xb2,
	ExtentRefV0 = 0xb4,
	SharedBlockRef = 0xb6,
	SharedDataRef = 0xb8,
	BlockGroup = 0xc0,        // implemented
	FreeSpaceInfo = 0xc6,     // implemented
	FreeSpaceExtent = 0xc7,   // implemented
	FreeSpaceBitmap = 0xc8,
	DevExtent = 0xcc,         // implemented
	Dev = 0xd8,               // implemented
	Chunk = 0xe4,             // implemented
	QGroupStatus = 0xf0,
	QGroupInfo = 0xf2,
	QGroupLimit = 0xf4,
	QGroupRelation = 0xf6,
	Temporary = 0xf8,
	Persistent = 0xf9,        // TODO?
	DevReplace = 0xfa,
	UUIDSubvol = 0xfb,        // implemented
	UUIDReceivedSubvol = 0xfc,
	String = 0xfd,
}

#[allow(unused)]
#[derive(Debug,Clone,Copy,ParseBin,PartialEq,Eq,PartialOrd,Ord)]
pub struct Key {
	pub key_id: u64,
	pub key_type: ItemType,
	pub key_offset: u64,
}

impl Key {
	pub fn new(key_id: u64, key_type: ItemType, key_offset: u64) -> Key {
		Key { key_id, key_type, key_offset }
	}

	pub fn id(key_id: u64) -> Key {
		Key { key_id, key_type: ItemType::Invalid, key_offset: 0 }
	}
}

pub const ZERO_KEY: Key = Key {key_id: 0, key_type: ItemType::Invalid, key_offset: 0};

#[allow(unused)]
#[derive(Debug,Clone)]
pub enum Value {
	Extent(ExtentItem),
	BlockGroup(BlockGroupItem),
	Inode(InodeItem),
	Chunk(ChunkItem),
	Root(RootItem),
	Dir(DirItem),
	DirIndex(DirItem),
	FreeSpaceInfo(FreeSpaceInfoItem),
	FreeSpaceExtent,
	UUIDSubvol(UUIDSubvolItem),
	Dev(DevItem),
	DevExtent(DevExtentItem),
	ExtentData(ExtentDataItem),
	Ref(RefItem),
	Unknown(Vec<u8>),
}

#[allow(unused)]
#[allow(unused)]
#[derive(Debug,Clone)]
pub struct Item {
	pub key: Key,
	pub value: Value,
}

#[allow(unused)]
#[derive(Clone,ParseBin)]
pub struct Checksum([u8; 32]);

#[allow(unused)]
#[derive(Clone,ParseBin)]
pub struct UUID([u8; 16]);

#[allow(unused)]
#[derive(Debug,ParseBin)]
pub struct Superblock {
	pub csum: Checksum,
	pub fsid: UUID,
	pub bytenr: u64,
	pub flags: u64,
	pub magic: u64,
	pub generation: u64,
	pub root: u64,
	pub chunk_root: u64,
	pub log_root: u64,
	#[skip_bytes = 8]
	pub total_bytes: u64,
	pub bytes_used: u64,

	pub root_dir_objectid: u64,
	pub num_devices: u64,
	pub sectorsize: u32,
	pub nodesize: u32,
	#[skip_bytes = 4]
	pub stripesize: u32,
	pub sys_chunk_array_size: u32,
	pub chunk_root_generation: u64,
	pub compat_flags: u64,
	pub compat_ro_flags: u64,
	pub incompat_flags: u64,
	pub csum_type: u16,
	pub root_level: u8,
	pub chunk_root_level: u8,
	pub log_root_level: u8,

	pub dev_item: DevItem,

	pub label: [u8; 0x100],

	pub cache_generation: u64,
	pub uuid_tree_generation: u64,

	pub metadata_uuid: UUID,
	pub nr_global_roots: u64,

	#[skip_bytes = 216]
	pub sys_chunk_array: [u8; 0x800],

	// next up would be the root backups, but we don't read them for now
}

#[allow(unused)]
#[derive(Debug,Clone,ParseBin)]
pub struct NodeHeader {
	pub csum: Checksum,
	pub fs_uid: UUID,
	pub bytenr: u64,
	pub flags: u64,
	pub chunk_tree_uid: UUID,
	pub generation: u64,
	pub owner: u64,
	pub nritems: u32,
	pub level: u8,
}

#[allow(unused)]
#[derive(Debug,Clone)]
pub struct Leaf {
	pub header: NodeHeader,
	pub items: Vec<Item>,
}

#[allow(unused)]
#[derive(Debug,Clone,ParseBin)]
pub struct KeyPointer {
	pub key: Key,
	pub ptr: u64,
	pub generation: u64,
}

#[allow(unused)]
#[derive(Debug,Clone)]
pub struct InteriorNode {
	pub header: NodeHeader,
	pub children: Vec<KeyPointer>,
}

#[allow(unused)]
#[derive(Debug,Clone)]
pub enum Node {
	Interior(InteriorNode),
	Leaf(Leaf),
}

#[allow(unused)]
#[derive(Debug,Clone,ParseBin)]
pub struct BlockGroupItem {
	pub used: u64,
	pub chunk_objectid: u64,
	pub flags: u64,
}

#[allow(unused)]
#[derive(Debug,Clone)]
pub struct ExtentItem {
	pub refs: u64,
	pub generation: u64,
	pub flags: u64,
	//	pub data: Vec<u8>,

	// this is only correct if flags == 2, fix later!
	pub block_refs: Vec<(ItemType, u64)>,
//	pub tree_block_key_type: ItemType,
//	pub tree_block_key_id: u64,
}

#[allow(unused)]
#[derive(Debug,Clone,ParseBin)]
pub struct Time {
	pub sec: u64,
	pub nsec: u32,
}

#[allow(unused)]
#[derive(Debug,Clone,ParseBin)]
pub struct InodeItem {
	pub generation: u64,
	pub transid: u64,
	pub size: u64,
	pub nbytes: u64,
	pub block_group: u64,
	pub nlink: u32,
	pub uid: u32,
	pub gid: u32,
	pub mode: u32,
	pub rdev: u64,
	pub flags: u64,
	pub sequence: u64,
	#[skip_bytes = 32]
	pub atime: Time,
	pub ctime: Time,
	pub mtime: Time,
	pub otime: Time,
}

#[allow(unused)]
#[derive(Debug,Clone,ParseBin)]
pub struct ChunkItem {
	pub size: u64,
	pub root: u64,
	pub stripelen: u64,
	pub chunktype: u64,
	pub align: u32,
	pub width: u32,
	pub sectorsize: u32,
	pub nrstripes: u16,
	pub substripes: u16,

	#[len = "nrstripes"]
	pub stripes: Vec<ChunkItemStripe>,
}

#[allow(unused)]
#[derive(Debug,Clone,ParseBin)]
pub struct ChunkItemStripe {
	pub devid: u64,
	pub offset: u64,
	pub devuuid: UUID,
}

#[allow(unused)]
#[derive(Debug,Clone,ParseBin)]
pub struct RootItem {
	pub inode: InodeItem,
	pub generation: u64,
	pub root_dirid: u64,
	pub bytenr: u64,
	pub byte_limit: u64,
	pub bytes_used: u64,
	pub last_snapshot: u64,
	pub flags: u64,
	pub refs: u32,
	pub drop_progress: Key,
	pub drop_level: u8,
	pub level: u8,

	pub generation_v2: u64,
	pub uuid: UUID,
	pub parent_uuid: UUID,
	pub received_uuid: UUID,
	pub ctransid: u64,
	pub otransid: u64,
	pub stransid: u64,
	pub rtransid: u64,
	pub ctime: Time,
	pub otime: Time,
	pub stime: Time,
	pub rtime: Time,
}

#[allow(unused)]
#[derive(Debug,Clone,ParseBin)]
pub struct DirItem {
	pub location: Key,
	pub transid: u64,
	pub data_len: u16,
	pub name_len: u16,
	pub dir_type: u8,

//	#[len = "name_len"]
	pub name: CString,
}

#[allow(unused)]
#[derive(Debug,Clone,ParseBin)]
pub struct FreeSpaceInfoItem {
	pub extent_count: u32,
	pub flags: u32,
}

#[allow(unused)]
#[derive(Debug,Clone,ParseBin)]
pub struct UUIDSubvolItem {
	pub subvol_id: u64,
}

#[allow(unused)]
#[derive(Debug,Clone,ParseBin)]
pub struct DevItem {
	pub devid: u64,
	pub total_bytes: u64,
	pub bytes_used: u64,
	pub io_align: u32,
	pub io_width: u32,
	pub sector_size: u32,
	pub dev_type: u64,
	pub generation: u64,
	pub start_offset: u64,
	pub dev_group: u32,
	pub seek_speed: u8,
	pub bandwidth: u8,
	pub uuid: UUID,
	pub fsid: UUID,
}

#[allow(unused)]
#[derive(Debug,Clone,ParseBin)]
pub struct DevExtentItem {
	pub chunk_tree: u64,
	pub chunk_objectid: u64,
	pub chunk_offset: u64,
	pub length: u64,
	pub chunk_tree_uuid: UUID,
}

#[allow(unused)]
#[derive(Debug,Clone)]
pub struct ExtentDataItem {
	pub header: ExtentDataHeader,
	pub data: ExtentDataBody,
}

#[allow(unused)]
#[derive(Debug,Clone)]
pub enum ExtentDataBody {
	Inline(Vec<u8>),
	External(ExternalExtent),
}

#[allow(unused)]
#[derive(Debug,Clone,ParseBin)]
pub struct ExternalExtent {
	pub disk_bytenr: u64,
	pub disk_num_bytes: u64,
	pub offset: u64,
	pub num_bytes: u64,
}

#[allow(unused)]
#[derive(Debug,Clone,ParseBin)]
pub struct ExtentDataHeader {
	pub generation: u64,
	pub ram_bytes: u64,
	pub compression: u8,
	pub encryption: u8,
	pub other_encoding: u16,
	pub extent_type: u8,
}

#[allow(unused)]
#[derive(Debug,Clone,ParseBin)]
pub struct RefItem {
	pub index: u64,
	pub name_len: u16,

	//	#[len = "name_len"]
	pub name: CString,
}

#[allow(unused)]
#[repr(u64)]
#[derive(Clone,Copy,Debug)]
pub enum TreeID {
	Root = 1,
	Extent = 2,
	Chunk = 3,
	Dev = 4,
	FS = 5,
	RootDir = 6,
	CSum = 7,
	Quota = 8,
	UUID = 9,
	FreeSpace = 10,
	BlockGroup = 11,
}

impl From<TreeID> for u64 {
	fn from(value: TreeID) -> u64 {
		value as u64
	}
}

/***** trait for parsing, and implementations for basic types *****/
// most of the more complex types will be parsed using derive macros

#[derive(Debug)]
pub struct ParseError(String);

impl error::Error for ParseError {}

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

impl From<String> for ParseError {
	fn from(value: String) -> ParseError {
		ParseError(value)
	}
}

impl From<&str> for ParseError {
	fn from(value: &str) -> ParseError {
		ParseError::from(String::from(value))
	}
}

pub trait ParseBin where Self: Sized {
	fn parse_len(bytes: &[u8]) -> Result<(Self, usize), ParseError>;

	fn parse(bytes: &[u8]) -> Result<Self, ParseError> {
		Self::parse_len(bytes).map(|x|x.0)
	}
}

impl ParseBin for u8 {
	fn parse_len(bytes: &[u8]) -> Result<(Self, usize), ParseError> {
		if bytes.len() < 1 {
			err!("not enough data")
		} else {
			Ok((bytes[0], 1))
		}
	}
}

impl ParseBin for u16 {
	fn parse_len(bytes: &[u8]) -> Result<(Self, usize), ParseError> {
		if bytes.len() < 2 {
			err!("not enough data")
		} else {
			let result = u16::from_le_bytes(bytes[0..2].try_into().unwrap());
			Ok((result, 2))
		}
	}
}

impl ParseBin for u32 {
	fn parse_len(bytes: &[u8]) -> Result<(Self, usize), ParseError> {
		if bytes.len() < 4 {
			err!("not enough data")
		} else {
			let result = u32::from_le_bytes(bytes[0..4].try_into().unwrap());
			Ok((result, 4))
		}
	}
}

impl ParseBin for u64 {
	fn parse_len(bytes: &[u8]) -> Result<(Self, usize), ParseError> {
		if bytes.len() < 8 {
			err!("not enough data")
		} else {
			let result = u64::from_le_bytes(bytes[0..8].try_into().unwrap());
			Ok((result, 8))
		}
	}
}

impl<const N: usize> ParseBin for [u8; N] {
	fn parse_len(bytes: &[u8]) -> Result<([u8; N], usize), ParseError> {
		if bytes.len() < N {
			err!("not enough data")
		} else {
			Ok((bytes[0..N].try_into().unwrap(), N))
		}
	}
}

// we use Vec<u8> for "unknown extra data", so just eat up everything
impl ParseBin for Vec<u8> {
	fn parse_len(bytes: &[u8]) -> Result<(Self, usize), ParseError> {
		Ok((Vec::from(bytes), bytes.len()))
	}
}

impl ParseBin for CString {
	fn parse_len(bytes: &[u8]) -> Result<(Self, usize), ParseError> {
		let mut chars = Vec::from(bytes);
		chars.push(0);
		Ok((CString::from_vec_with_nul(chars).unwrap(), bytes.len()))
	}
}

/***** parsing ItemType and Leaf *****/

impl From<ItemType> for u8 {
	fn from(value: ItemType) -> u8 {
		value as u8
	}
}

impl From<u8> for ItemType {
	fn from(value: u8) -> ItemType {
		let variants = ItemType::all_variants();
		match variants.binary_search_by_key(&value, |x|u8::from(*x)) {
			Ok(idx) => variants[idx],
			Err(_) => ItemType::Invalid,
		}
	}
}

impl ParseBin for ItemType {
	fn parse_len(bytes: &[u8]) -> Result<(ItemType, usize), ParseError> {
		u8::parse(bytes).map(|x|(ItemType::from(x),1))
	}
}

impl ParseBin for Node {
	fn parse_len(bytes: &[u8]) -> Result<(Node, usize), ParseError> {
		if bytes.len() < 0x65 {
			return err!("Not enough data to parse node header");
		}

		let header = NodeHeader::parse(&bytes[0..0x65])?;

		if header.level > 0 {
			// interior node
			let mut children = Vec::new();
			let num = header.nritems as usize;

			for i in 0 .. num {
				children.push(KeyPointer::parse(&bytes[0x65 + i*0x21 .. 0x86 + i*0x21])?);
			}

			Ok((Node::Interior(InteriorNode { header, children }), NODE_SIZE))
		} else {
			// leaf node
			let mut items = Vec::new();

			for i in 0..header.nritems as usize {
				let key =    Key::parse(&bytes[i*0x19 + 0x65        .. i*0x19 + 0x65 + 0x11])?;
				let offset = u32::parse(&bytes[i*0x19 + 0x65 + 0x11 .. i*0x19 + 0x65 + 0x15])?;
				let size =   u32::parse(&bytes[i*0x19 + 0x65 + 0x15 .. i*0x19 + 0x65 + 0x19])?;

				let data_slice = &bytes[0x65 + offset as usize ..
										0x65 + offset as usize + size as usize];

				let value = match key.key_type {
					ItemType::BlockGroup =>
						Value::BlockGroup(BlockGroupItem::parse(data_slice)?),
					ItemType::Metadata => {
						let item = ExtentItem::parse(data_slice)?;
						if item.flags != 2 || item.refs > 1 {
							println!("Metadata item with refs = {}, flags = {}, data = {:x?}", item.refs, item.flags, &data_slice[0x18..]);
						}
						Value::Extent(item)
					},
					ItemType::Extent =>
						Value::Extent(ExtentItem::parse(data_slice)?),
					ItemType::Inode =>
						Value::Inode(InodeItem::parse(data_slice)?),
					ItemType::Root =>
						Value::Root(RootItem::parse(data_slice)?),
					ItemType::Dir =>
						Value::Dir(DirItem::parse(data_slice)?),
					ItemType::DirIndex =>
						Value::DirIndex(DirItem::parse(data_slice)?),
					ItemType::Chunk =>
						Value::Chunk(ChunkItem::parse(data_slice)?),
					ItemType::FreeSpaceInfo =>
						Value::FreeSpaceInfo(FreeSpaceInfoItem::parse(data_slice)?),
					ItemType::FreeSpaceExtent =>
						Value::FreeSpaceExtent,
					ItemType::UUIDSubvol =>
						Value::UUIDSubvol(UUIDSubvolItem::parse(data_slice)?),
					ItemType::Dev =>
						Value::Dev(DevItem::parse(data_slice)?),
					ItemType::DevExtent =>
						Value::DevExtent(DevExtentItem::parse(data_slice)?),
					ItemType::ExtentData =>
						Value::ExtentData(ExtentDataItem::parse(data_slice)?),
					ItemType::Ref =>
						Value::Ref(RefItem::parse(data_slice)?),
					_ =>
						Value::Unknown(Vec::from(data_slice)),
				};

				items.push(Item { key, value });
			}

			Ok((Node::Leaf(Leaf { header, items }), NODE_SIZE))
		}
	}
}

impl ParseBin for InteriorNode {
	fn parse_len(bytes: &[u8]) -> Result<(Self, usize), ParseError> {
		match Node::parse_len(bytes)? {
			(Node::Interior(int_node), len) => Ok((int_node, len)),
			_ => err!("Expected interior node, found leaf"),
		}
	}
}

impl ParseBin for Leaf {
	fn parse_len(bytes: &[u8]) -> Result<(Self, usize), ParseError> {
		match Node::parse_len(bytes)? {
			(Node::Leaf(leaf_node), len) => Ok((leaf_node, len)),
			_ => err!("Expected leaf, found interior node"),
		}
	}
}

// ExtentDataItem needs a custom implementation because it can have inline or external data
impl ParseBin for ExtentDataItem {
	fn parse_len(bytes: &[u8]) -> Result<(Self, usize), ParseError> {
		let (header, header_size) = ExtentDataHeader::parse_len(bytes)?;
		if header.extent_type == 1 { // external extent
			let (body, body_size) = ExternalExtent::parse_len(&bytes[header_size..])?;
			return Ok((ExtentDataItem { header: header, data: ExtentDataBody::External(body)},
					   header_size + body_size))
		} else { // inline extent
			let data_slice = &bytes[header_size..];
			return Ok((ExtentDataItem {
				header: header,
				data: ExtentDataBody::Inline(Vec::from(data_slice))
			}, header_size + data_slice.len()))
		}
	}
}

impl ParseBin for ExtentItem {
	fn parse_len(bytes: &[u8]) -> Result<(Self, usize), ParseError> {
		let refs = u64::parse(bytes)?;
		let generation = u64::parse(&bytes[0x08..])?;
		let flags = u64::parse(&bytes[0x10..])?;

		let mut block_refs = Vec::new();

		if flags & 0x03 == 0x02 {
			for i in 0 .. refs as usize {
				let key_type = ItemType::parse(&bytes[0x18 + i*0x09 .. ])?;
				let key_id = u64::parse(&bytes[0x19 + i*0x09 .. ])?;
				block_refs.push((key_type, key_id));
			}
		}

		Ok((ExtentItem { refs, generation, flags, block_refs }, 0x18 + refs as usize * 0x09))
	}
}

/***** prettier debug output for UUIDs and checksums *****/

impl fmt::Debug for UUID {
	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
		let UUID([x0, x1, x2,  x3,  x4,  x5,  x6,  x7,
				  x8, x9, x10, x11, x12, x13, x14, x15]) = self;
		write!(f, "{:02x}{:02x}{:02x}{:02x}-{:02x}{:02x}-{:02x}{:02x}-\
				   {:02x}{:02x}-{:02x}{:02x}{:02x}{:02x}{:02x}{:02x}",
			   x0, x1, x2,  x3,  x4,  x5,  x6,  x7,
			   x8, x9, x10, x11, x12, x13, x14, x15)
	}
}

impl fmt::Debug for Checksum {
	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
		let Checksum(
			[x0,  x1,  x2,  x3,  x4,  x5,  x6,  x7,
			 x8,  x9,  x10, x11, x12, x13, x14, x15,
			 x16, x17, x18, x19, x20, x21, x22, x23,
			 x24, x25, x26, x27, x28, x29, x30, x31]) = self;
		write!(f, "{:02x}{:02x}{:02x}{:02x}{:02x}{:02x}{:02x}{:02x}\
				   {:02x}{:02x}{:02x}{:02x}{:02x}{:02x}{:02x}{:02x}\
				   {:02x}{:02x}{:02x}{:02x}{:02x}{:02x}{:02x}{:02x}\
				   {:02x}{:02x}{:02x}{:02x}{:02x}{:02x}{:02x}{:02x}",
			   x0,  x1,  x2,  x3,  x4,  x5,  x6,  x7,
			   x8,  x9,  x10, x11, x12, x13, x14, x15,
			   x16, x17, x18, x19, x20, x21, x22, x23,
			   x24, x25, x26, x27, x28, x29, x30, x31)
	}
}

#[macro_export]
macro_rules! key {
	($arg1:expr) => {
		btrfs_structs::Key { key_id: $arg1, key_type: btrfs_structs::ItemType::Invalid, key_offset: 0 }
	};
	($arg1:expr, $arg2:expr) => {
		btrfs_structs::Key { key_id: $arg1, key_type: $arg2, key_offset: 0 }
	};
	($arg1:expr, $arg2:expr, $arg3:expr) => {
		btrfs_structs::Key { key_id: $arg1, key_type: $arg2, key_offset: $arg3 }
	};
}

//pub(crate) use key;