//! Provides a low-level parser for the FDT, you can build your own interface on top of this.

use core::ffi::CStr;

/// The FDT magic number.
pub const MAGIC: u32 = 0xd00dfeed;
/// The FDT version, compatible with this crate.
pub const DTB_VERSION: u32 = 17;

#[repr(C)]
pub struct RawHeader {
    magic: u32,
    totalsize: u32,
    off_dt_struct: u32,
    off_dt_strings: u32,
    off_mem_rvsmap: u32,
    version: u32,
    last_comp_version: u32,
    boot_cpuid_phys: u32,
    size_dt_strings: u32,
    size_dt_structs: u32,
}

impl RawHeader {
    pub const fn magic(&self) -> u32 {
        u32::from_be(self.magic)
    }

    pub const fn totalsize(&self) -> u32 {
        u32::from_be(self.totalsize)
    }
    pub const fn off_dt_struct(&self) -> u32 {
        u32::from_be(self.off_dt_struct)
    }
    pub const fn off_dt_strings(&self) -> u32 {
        u32::from_be(self.off_dt_strings)
    }

    pub const fn off_mem_rvsmap(&self) -> u32 {
        u32::from_be(self.off_mem_rvsmap)
    }

    /// Returns a slice of the DT Struct's bytes.
    pub const fn dt_structs_bytes(&self) -> *const [u8] {
        unsafe {
            core::slice::from_raw_parts(
                (self as *const Self as *const u8).add(self.off_dt_struct() as usize),
                self.size_dt_structs() as usize,
            )
        }
    }

    pub const fn dt_strings_bytes(&self) -> *const [u8] {
        unsafe {
            core::slice::from_raw_parts(
                (self as *const Self as *const u8).add(self.off_dt_strings() as usize),
                self.size_dt_strings() as usize,
            )
        }
    }

    /// Returns the number of entries in the memory revisions map.
    ///
    /// O(N) as it iterates through the map until a zero entry is found.
    pub const fn mem_rvsmap_len(&self) -> usize {
        unsafe {
            let mut curr: *const RawMemRvsEntry = (self as *const Self as *const u8)
                .add(self.off_mem_rvsmap() as usize)
                .cast();

            let mut count = 0;
            while (*curr).address != 0 && (*curr).size != 0 {
                curr = curr.add(1);
                count += 1;
            }
            count
        }
    }

    /// Returns a pointer to the memory revisions map.
    ///
    /// O(N) as it iterates through the map until a zero entry is found.
    /// # Safety
    ///
    /// that pointer is an offset from the header reference.
    pub const fn mem_rvsmap(&self) -> *const [RawMemRvsEntry] {
        unsafe {
            core::slice::from_raw_parts(
                (self as *const Self as *const u8)
                    .add(self.off_mem_rvsmap() as usize)
                    .cast(),
                self.mem_rvsmap_len(),
            )
        }
    }

    pub const fn version(&self) -> u32 {
        u32::from_be(self.version)
    }

    pub const fn last_comp_version(&self) -> u32 {
        u32::from_be(self.last_comp_version)
    }
    pub const fn boot_cpuid_phys(&self) -> u32 {
        u32::from_be(self.boot_cpuid_phys)
    }
    pub const fn size_dt_strings(&self) -> u32 {
        u32::from_be(self.size_dt_strings)
    }
    pub const fn size_dt_structs(&self) -> u32 {
        u32::from_be(self.size_dt_structs)
    }

    fn is_compatible(&self) -> bool {
        self.magic() == MAGIC && self.version() == DTB_VERSION
    }

    /// Given a raw pointer to a DTB header, returns a compatible [`RawHeader`] if possible.
    pub unsafe fn from_raw<'a>(data: *const u8) -> Option<&'a Self> {
        let p = unsafe { &*(data as *const Self) };
        if p.is_compatible() { Some(p) } else { None }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(C, packed)]
pub struct RawMemRvsEntry {
    address: u64,
    size: u64,
}

impl RawMemRvsEntry {
    pub const fn address(&self) -> u64 {
        u64::from_be(self.address)
    }
    pub const fn size(&self) -> u64 {
        u64::from_be(self.size)
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct RawFDT<'a> {
    pub version: u32,
    pub mem_rvsmap: &'a [RawMemRvsEntry],
    pub dt_structs_bytes: &'a [u8],
    pub dt_strings_bytes: &'a [u8],
}

impl<'a> RawFDT<'a> {
    pub fn get_cstr(&self, offset: usize) -> &CStr {
        CStr::from_bytes_until_nul(&self.dt_strings_bytes[offset..])
            .expect("Corrupted Device Tree Nul not found")
    }

    pub fn get_str(&self, offset: usize) -> &str {
        self.get_cstr(offset)
            .to_str()
            .expect("Invalid UTF-8 in Device Tree string")
    }

    /// Parses a raw FDT from the given pointer, returning `None` if the data is invalid.
    pub unsafe fn from_raw(data: *const u8) -> Option<Self> {
        let header = unsafe { RawHeader::from_raw(data) }?;
        unsafe {
            Some(Self {
                version: header.version(),
                mem_rvsmap: &*header.mem_rvsmap(),
                dt_structs_bytes: &*header.dt_structs_bytes(),
                dt_strings_bytes: &*header.dt_strings_bytes(),
            })
        }
    }

    /// Returns an iterator over the tokens within the FDT's structures section.
    pub const fn iter(&'a self) -> RawFDTIterator<'a> {
        self.iter_off(0)
    }

    /// Returns an iterator over the tokens within the FDT starting at the given offset from within the FDT's structures section.
    ///
    /// # Safety
    ///
    /// - May panic if offset doesn't point to a start of a token or is out of bounds.
    pub const fn iter_off(&'a self, offset: u32) -> RawFDTIterator<'a> {
        RawFDTIterator {
            fdt: self,
            curr_offset: offset,
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
/// Represents a raw FDT token's data.
///
/// Such as a node begin name, property name&value, or end of node/FDT.
pub enum RawFDTTokenData<'a> {
    /// A node's beginning, contains the node's name.
    NodeBegin(&'a str),
    /// A property, contains the property's name and value.
    Prop { name: &'a str, value: &'a [u8] },
    /// A node's end.
    NodeEnd,
    /// The FDT's end.
    FDTEnd,
}

#[derive(Debug, Clone, Copy)]
/// Represents a raw FDT token, including its offset in bytes within the FDT's structures section and data.
pub struct RawFDTToken<'a> {
    pub offset: u32,
    pub data: RawFDTTokenData<'a>,
}

pub const FDT_BEGIN_NODE: u32 = 0x00000001;
pub const FDT_END_NODE: u32 = 0x00000002;
pub const FDT_PROP: u32 = 0x00000003;
pub const FDT_NOP: u32 = 0x00000004;
pub const FDT_END: u32 = 0x00000009;

macro_rules! align_to {
    ($val:expr, $align:expr) => {{
        const _: () = assert!($align.is_power_of_two());
        ($val + ($align - 1)) & !($align - 1)
    }};
}

/// An Iterator over the raw FDT tokens [`RawFDTToken`].
#[derive(Debug, Clone, Copy)]
pub struct RawFDTIterator<'a> {
    fdt: &'a RawFDT<'a>,
    curr_offset: u32,
}

impl<'a> Iterator for RawFDTIterator<'a> {
    type Item = RawFDTToken<'a>;

    fn next(&mut self) -> Option<Self::Item> {
        let bytes = self.fdt.dt_structs_bytes;

        if self.curr_offset >= bytes.len() as u32 {
            return None;
        }
        macro_rules! read_unit {
            () => {{
                let bytes = &bytes[self.curr_offset as usize..(self.curr_offset + 4) as usize];
                self.curr_offset += 4;
                u32::from_be_bytes(bytes.try_into().unwrap())
            }};
        }

        let mut off;
        let data = loop {
            off = self.curr_offset;

            let token = read_unit!();
            match token {
                FDT_NOP => continue,
                FDT_END => {
                    break RawFDTTokenData::FDTEnd;
                }
                FDT_END_NODE => {
                    break RawFDTTokenData::NodeEnd;
                }
                FDT_BEGIN_NODE => {
                    let c_str = CStr::from_bytes_until_nul(&bytes[self.curr_offset as usize..])
                        .expect("While iterating FDT: Node name wasn't null-terminated");
                    let str = c_str
                        .to_str()
                        .expect("While iterating FDT: Node name wasn't a string");

                    let aligned_len = align_to!(str.len() + 1, size_of::<u32>());
                    self.curr_offset += aligned_len as u32;
                    break RawFDTTokenData::NodeBegin(str);
                }
                FDT_PROP => {
                    let len = read_unit!();
                    let name_off = read_unit!();
                    let aligned_len = align_to!(len, 4u32);

                    let name = self.fdt.get_str(name_off as usize);
                    let value = &self.fdt.dt_structs_bytes
                        [self.curr_offset as usize..(self.curr_offset + len) as usize];
                    self.curr_offset += aligned_len;
                    break RawFDTTokenData::Prop { name, value };
                }
                _ => panic!(
                    "While iterating FDT: unexpected token 0x{:08x}, offset={}",
                    token, self.curr_offset
                ),
            }
        };
        Some(RawFDTToken { offset: off, data })
    }
}
