//! A higher-level abstraction over [`crate::base`].
//!
//! Provides a safe, high-level interface for working with FDT nodes and properties, and parsing common properties.
//!
//! There are no docs for what properties represent see the devicetree specification for details.
//!
//! NOTE: This will panic if the FDT is corrupted, also expect any specific property parsing method to panic if parsing fails and the property is corrupted.
mod cell;
mod ranges;
mod str;

use core::fmt::Debug;
use core::sync::atomic::AtomicIsize;

pub use cell::*;
pub use str::*;

use crate::{
    RawFDTTokenData, base,
    node::ranges::{NCellsIter, RangesIter},
};

/// A node in the FDT, can be a subnode.
pub struct Node<'a> {
    parent_offset: AtomicIsize,
    fdt_offset: u32,
    name: &'a str,
    fdt: &'a Fdt<'a>,
}

impl Debug for Node<'_> {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("Node")
            .field("fdt_offset", &self.fdt_offset)
            .field("name", &self.name)
            .field(
                "parent_offset",
                &self
                    .parent_offset
                    .load(core::sync::atomic::Ordering::Relaxed),
            )
            .finish()
    }
}

impl PartialEq for Node<'_> {
    fn eq(&self, other: &Self) -> bool {
        self.fdt_offset == other.fdt_offset && self.fdt == other.fdt
    }
}

impl Clone for Node<'_> {
    fn clone(&self) -> Self {
        Self {
            parent_offset: AtomicIsize::new(
                self.parent_offset
                    .load(core::sync::atomic::Ordering::Relaxed),
            ),
            fdt_offset: self.fdt_offset,
            name: self.name,
            fdt: self.fdt,
        }
    }
}

impl<'a> Node<'a> {
    /// Returns the parent node of this node, if it has one.
    ///
    /// May lookup the parent node in the FDT if it has not been cached yet.
    pub fn parent(&self) -> Option<Node<'a>> {
        let parent_offset = self
            .parent_offset
            .load(core::sync::atomic::Ordering::Relaxed);

        if parent_offset.is_positive() && parent_offset != 0 {
            let parent_offset = (parent_offset - 1) as u32;
            let Some(RawFDTTokenData::NodeBegin(name)) = self
                .raw_fdt()
                .iter_off(parent_offset)
                .next()
                .map(|t| t.data)
            else {
                panic!("Failed to find parent node at offset {parent_offset}");
            };
            Some(Node {
                parent_offset: AtomicIsize::new(-1),
                fdt_offset: parent_offset,
                name: name,
                fdt: self.fdt,
            })
        } else if parent_offset == -1 {
            let node = self
                .fdt
                .find_node(|node| node.subnodes().any(|s| s.fdt_offset == self.fdt_offset));

            if let Some(node) = node {
                _ = self.parent_offset.compare_exchange(
                    -1,
                    (node.fdt_offset + 1) as isize,
                    core::sync::atomic::Ordering::Acquire,
                    core::sync::atomic::Ordering::Relaxed,
                );
                Some(node)
            } else {
                _ = self.parent_offset.compare_exchange(
                    -1,
                    0,
                    core::sync::atomic::Ordering::Acquire,
                    core::sync::atomic::Ordering::Relaxed,
                );
                None
            }
        } else {
            None
        }
    }

    pub fn compatible(&self) -> Option<StrIter<'a>> {
        self.property("compatible").map(|compt| compt.as_str_iter())
    }

    /// Returns the phandle of the node, if it exists.
    pub fn phandle(&self) -> Option<u32> {
        self.property("phandle")
            .map(|phandle| phandle.as_u32().expect("Invalid phandle"))
    }

    /// Returns the value of the `interrupts` property, if it exists, as an iterator over interrupts_parent.[`Self::interrupt_cells`] [`Cells`] each.
    ///
    /// Get the interrupt-parent with [`Self::interrupt_parent()`].
    pub fn interrupts(
        &self,
        interrupt_parent: &Node<'a>,
    ) -> Option<impl Iterator<Item = Cells<'a>> + use<'a> + Clone> {
        let interrupts_cells = interrupt_parent.interrupt_cells();
        self.property("interrupts").map(|interrupts| {
            interrupts
                .as_cells_groups([interrupts_cells.unwrap_or(1)])
                .map(|[c]| c)
        })
    }

    /// Returns the value of the `interrupt-map-mask` property, if it exists.
    pub fn interrupt_map_mask(&self) -> Option<Cells<'a>> {
        self.property("interrupt-map-mask")
            .map(|mask| mask.as_cells())
    }

    /// Returns the value of the `interrupt-map` property, if it exists.
    ///
    /// See [`Self::parse_interrupt_map`]..
    pub fn interrupt_map(&self) -> Option<Cells<'a>> {
        self.property("interrupt-map").map(|i| i.as_cells())
    }

    /// Parses the [`Self::interrupt_map`] property if it exists.
    ///
    /// Returns an iterator over the parsed interrupts map entries.
    ///
    /// Panics if any of the phandles in the map are invalid.
    /// See [`Self::interrupt_map`], for a non-panicking alternative.
    pub fn parse_interrupt_map(
        &self,
    ) -> Option<
        impl Iterator<Item = (Cells<'a>, Cells<'a>, Node<'a>, Cells<'a>, Cells<'a>)> + use<'a> + Clone,
    > {
        let interrupts_map = self.interrupt_map()?;
        let interrupts_cells = self.interrupt_cells().unwrap_or(1);
        let address_cells = self.address_cells();

        #[derive(Debug, Clone)]
        struct InterruptsMapIter<'a> {
            fdt: &'a Fdt<'a>,
            cells: Cells<'a>,
            interrupts_cells: u32,
            address_cells: u32,
            last_node: Option<(Node<'a>, u32)>,
        }

        impl<'a> Iterator for InterruptsMapIter<'a> {
            type Item = (Cells<'a>, Cells<'a>, Node<'a>, Cells<'a>, Cells<'a>);

            fn next(&mut self) -> Option<Self::Item> {
                let address = self.cells.take_n(self.address_cells)?;
                let interrupts = self.cells.take_n(self.interrupts_cells)?;
                let phandle = self.cells.next_cell()?;
                let node = self
                    .last_node
                    .clone()
                    .filter(|(_, p)| *p == phandle)
                    .map(|(n, _)| n)
                    .unwrap_or_else(|| {
                        self.fdt
                            .find_by_phandle(phandle)
                            .expect("Failed to find the crossponding interrupts-map node")
                    });

                let node_address_cells = node.address_cells();
                let node_interrupt_cells = node.interrupt_cells().unwrap_or(1);

                let node_address = self.cells.take_n(node_address_cells)?;
                let node_interrupts = self.cells.take_n(node_interrupt_cells)?;

                self.last_node = Some((node.clone(), phandle));
                Some((address, interrupts, node, node_address, node_interrupts))
            }
        }

        Some(InterruptsMapIter {
            fdt: self.fdt,
            cells: interrupts_map,
            interrupts_cells,
            address_cells,
            last_node: None,
        })
    }

    /// Returns the value of the `interrupt-parent` property, if it exists.
    pub fn interrupt_parent_phandle(&self) -> Option<u32> {
        self.property("interrupt-parent")
            .map(|phandle| phandle.as_u32().expect("Invalid phandle"))
    }

    /// Returns the interrupt parent node, if it exists.
    ///
    /// Uses [`Self::interrupt_parent_phandle`] to find the parent node by phandle,
    /// and falls back to [`Self::parent`] if no phandle is found.
    pub fn interrupt_parent(&self) -> Option<Node<'a>> {
        self.interrupt_parent_phandle()
            .and_then(|phandle| self.fdt.find_by_phandle(phandle))
            .or_else(|| self.parent())
    }

    /// Returns `true` if the node is an interrupt controller.
    pub fn is_interrupt_controller(&self) -> bool {
        self.has("interrupt-controller")
    }

    /// Parses the `#interrupt-cells` property, if it exists.
    ///
    /// Returns the number of cells in an interrupt specifier.
    pub fn interrupt_cells(&self) -> Option<u32> {
        self.property("#interrupt-cells")
            .map(|cells| cells.as_u32().expect("Failed to parse <interrupt-cells>"))
    }

    /// Attempts to read the `#address-cells` property, returning `None` if not found.
    ///
    /// Panicks if the property cannot be parsed as a `u32`.
    pub fn try_address_cells(&self) -> Option<u32> {
        self.property("#address-cells")
            .map(|cells| cells.as_u32().expect("Failed to parse <#address-cells>"))
    }

    /// Attempts to read the `#size-cells` property, returning `None` if not found.
    ///
    /// Panicks if the property cannot be parsed as a `u32`.
    pub fn try_size_cells(&self) -> Option<u32> {
        self.property("#size-cells")
            .map(|cells| cells.as_u32().expect("Failed to parse <#size-cells>"))
    }

    /// Attempts to parse the `#address-cells` property, returning `2` if not found.
    ///
    /// If necessary to find out whether the property exists, use [`Node::try_address_cells`] instead.
    pub fn address_cells(&self) -> u32 {
        self.try_address_cells().unwrap_or(2)
    }

    /// Returns the `#address-cells` property of the parent node, or `2` if not found.
    pub fn parent_address_cells(&self) -> u32 {
        self.parent().map(|p| p.address_cells()).unwrap_or(2)
    }

    /// Returns the `#size-cells` property of the parent node, or `1` if not found.
    pub fn parent_size_cells(&self) -> u32 {
        self.parent().map(|p| p.size_cells()).unwrap_or(1)
    }

    /// Attempts to parse the `reg` property, returning `None` if not found or if #address-cells or #size-cells of the parent node are bigger than the amount of cells `usize` can hold (i.e. 2 on 64-bit systems).
    ///
    /// See [`Self::reg`].
    /// Doesn't panic, results may be truncated in case of corrupted data.
    pub fn reg_addresses(&self) -> Option<impl Iterator<Item = (usize, usize)> + use<'a> + Clone> {
        self.property("reg").and_then(|prop| {
            prop.as_address_size_iter(self.parent_address_cells(), self.parent_size_cells())
        })
    }

    /// Returns the `reg` property as a [`Cells`], if one exists.
    /// Doesn't panic, results may be truncated in case of corrupted data.
    pub fn reg(&self) -> Option<Cells<'a>> {
        self.property("reg").map(|prop| prop.as_cells())
    }

    /// Attempts to parse the `ranges` property, returning `None` if not found or if #address-cells of the parent node or #size-cells of this node are bigger than the amount of cells `usize` can hold (i.e. 2 on 64-bit systems).
    ///
    /// See [`Self::ranges`].
    /// Doesn't panic, results may be truncated in case of corrupted data.
    pub fn ranges_addresses(
        &self,
    ) -> Option<impl Iterator<Item = (Cells<'a>, usize, usize)> + use<'a> + Clone> {
        self.property("ranges").and_then(|prop| {
            prop.as_ranges(
                self.address_cells(),
                self.parent_address_cells(),
                self.size_cells(),
            )
        })
    }

    /// Returns the `ranges` property as a [`Cells`], if one exists.
    pub fn ranges(&self) -> Option<Cells<'a>> {
        self.property("ranges").map(|prop| prop.as_cells())
    }

    /// Attempts to parse the `#size-cells` property, returning `1` if not found.
    ///
    /// If necessary to find out whether the property exists, use [`Node::try_size_cells`] instead.
    pub fn size_cells(&self) -> u32 {
        self.try_size_cells().unwrap_or(1)
    }

    /// Returns the property with the given name, if one exists.
    ///
    /// Requires Iterating all the prior properties so it is O(N).
    ///
    /// Panics if the fdt is corrupted.
    pub fn property(&self, name: &str) -> Option<Property<'a>> {
        self.properties().find(|prop| prop.name == name)
    }

    /// Checks if the node has a property with the given name.
    ///
    /// Requires iterating and serializing all the prior properties.
    #[inline(always)]
    pub fn has(&self, name: &str) -> bool {
        self.properties().any(|prop| prop.name == name)
    }

    /// Returns the name of the node.
    pub const fn name(&self) -> &str {
        self.name
    }

    /// Returns the raw FDT structure for this node.
    pub const fn raw_fdt(&self) -> &'a base::RawFDT<'a> {
        &self.fdt.raw_fdt
    }

    /// Returns the FDT structure for this node.
    pub const fn fdt(&self) -> &'a Fdt<'a> {
        &self.fdt
    }

    /// Returns an iterator over the subnodes of this node.
    ///
    /// The iterator would panic if the fdt is corrupted.
    pub fn subnodes(&self) -> impl Iterator<Item = Node<'a>> + 'a {
        self.iter().filter_map(|child| match child {
            Child::Subnode(node) => Some(node),
            Child::Property(_) => None,
        })
    }

    /// Returns an iterator over the properties of this node.
    ///
    /// The iterator would panic if the fdt is corrupted.
    pub fn properties(&self) -> impl Iterator<Item = Property<'a>> + 'a {
        self.iter().filter_map(|child| match child {
            Child::Subnode(_) => None,
            Child::Property(prop) => Some(prop),
        })
    }

    /// Returns an iterator over the children of this node, properties + subnodes in order.
    ///
    /// The iterator would panic if the fdt is corrupted.
    pub fn iter(&self) -> impl Iterator<Item = Child<'a>> + 'a {
        let mut iter = self.raw_fdt().iter_off(self.fdt_offset);
        assert_eq!(
            iter.next().map(|t| t.data),
            Some(RawFDTTokenData::NodeBegin(self.name))
        );

        NodeChildrenIter {
            parent_offset: Some(self.fdt_offset),
            fdt: self.fdt,
            fdt_iter: iter,
            is_end: false,
        }
    }
}

/// A property of a node in the FDT.
#[derive(Debug, Clone, PartialEq)]
pub struct Property<'a> {
    name: &'a str,
    value: &'a [u8],
    fdt_offset: u32,
}

impl<'a> Property<'a> {
    /// Returns the name of this property.
    pub const fn name(&self) -> &'a str {
        self.name
    }

    /// Returns the value of this property as a byte slice.
    pub const fn as_bytes(&self) -> &'a [u8] {
        self.value
    }

    /// Attempts to interpret the value of this property as a UTF-8 string.
    pub fn as_str(&self) -> Option<&'a str> {
        match core::str::from_utf8(self.value) {
            Ok(s) => Some(s.trim_end_matches('\0')),
            Err(_) => None,
        }
    }

    /// Attempts to interpret the value of this property as a u32 value.
    pub const fn as_u32(&self) -> Option<u32> {
        match self.value.len() {
            4 => Some(u32::from_be_bytes([
                self.value[0],
                self.value[1],
                self.value[2],
                self.value[3],
            ])),
            _ => None,
        }
    }

    /// Returns the value of this property as a slice of u32 values, to read the values you have to call `u32::from_be` on each value.
    ///
    /// NOTE: If the value is not a multiple of 4 bytes, the last value will be entirely skipped.
    pub const fn as_u32be_cells(&self) -> &'a [u32] {
        // SAFETY: The pointer is guaranteed to be valid and aligned for u32,
        unsafe {
            core::slice::from_raw_parts(self.value.as_ptr().cast::<u32>(), self.value.len() / 4)
        }
    }

    /// Returns the value of this property as a [`Cells`] of u32 values.
    ///
    /// You can use [`Self::as_u32be_cells`] to get the raw slice of u32 values, but they'd be in big-endian byte order.
    /// NOTE: If the value is not a multiple of 4 bytes, the last value will be entirely skipped.
    pub const fn as_cells(&self) -> Cells<'a> {
        Cells::new(self.as_u32be_cells())
    }

    /// Same as [`Self::as_ranges`], but with bus_address_cells=0, i.e. the bus address is not included in the result.
    ///
    /// Can be used for the <reg> property.
    ///
    /// Doesn't panic results may be truncated in case of corrupted data.
    pub fn as_address_size_iter(
        &self,
        address_cells: u32,
        size_cells: u32,
    ) -> Option<impl Iterator<Item = (usize, usize)> + Clone + use<'a>> {
        RangesIter::new(self.as_cells(), 0, address_cells, size_cells)
            .map(|i| i.map(|(_, a, b)| (a, b)))
    }

    /// Returns an iterator over the value of this property as `N` cell groups.
    ///
    /// Each group is a [`Cells`] slice of `cells_count[i]` cells.
    ///
    /// Doesn't panic, results may be truncated in case of corrupted data.
    pub fn as_cells_groups<const N: usize>(
        &self,
        cells_count: [u32; N],
    ) -> impl Iterator<Item = [Cells<'a>; N]> + use<'a, N> + Copy {
        NCellsIter::new(self.as_cells(), cells_count)
    }

    /// Attempts to interpret the value of this property as an iterator of ([`Cells`], usize, usize) triples, where each triple represents a bus address, a physical address, and a size.
    ///
    /// I.e. the <ranges> or <dma-ranges> properties.
    ///
    /// Returns `None` if `bus_size_cells` or `p_address_cells` is too much cells to fit in a `usize`.
    /// `p_address_cells` is the parent's address cells count.
    ///
    /// Doesn't panic, results may be truncated in case of corrupted data.
    pub const fn as_ranges(
        &self,
        bus_address_cells: u32,
        p_address_cells: u32,
        bus_size_cells: u32,
    ) -> Option<impl Iterator<Item = (Cells<'a>, usize, usize)> + use<'a> + Copy> {
        RangesIter::new(
            self.as_cells(),
            bus_address_cells,
            p_address_cells,
            bus_size_cells,
        )
    }

    /// Returns an iterator over the property value as an iterator of UTF-8 strings.
    ///
    /// Panics if any of the strings are not valid UTF-8.
    pub const fn as_str_iter(&self) -> StrIter<'a> {
        StrIter::new(self.value)
    }
}

/// A child of a node in the FDT, can be a property or a subnode.
#[derive(Debug, Clone, PartialEq)]
pub enum Child<'a> {
    Property(Property<'a>),
    Subnode(Node<'a>),
}

struct NodeChildrenIter<'a> {
    parent_offset: Option<u32>,
    fdt: &'a Fdt<'a>,
    fdt_iter: base::RawFDTIterator<'a>,
    is_end: bool,
}

impl<'a> Iterator for NodeChildrenIter<'a> {
    type Item = Child<'a>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.is_end {
            return None;
        }

        loop {
            let next = self.fdt_iter.next()?;
            match next.data {
                RawFDTTokenData::NodeBegin(name) => {
                    let mut nesting = 0;
                    while let Some(next) = self.fdt_iter.next() {
                        match next.data {
                            RawFDTTokenData::NodeBegin(_) => nesting += 1,
                            RawFDTTokenData::NodeEnd => {
                                if nesting == 0 {
                                    break;
                                }
                                nesting -= 1;
                            }
                            RawFDTTokenData::FDTEnd => {
                                self.is_end = true;
                                break;
                            }
                            _ => {}
                        }
                    }
                    break Some(Child::Subnode(Node {
                        parent_offset: AtomicIsize::new(
                            self.parent_offset
                                .map(|off| (off + 1) as isize)
                                .unwrap_or(0),
                        ),
                        fdt_offset: next.offset,
                        name,
                        fdt: self.fdt,
                    }));
                }
                RawFDTTokenData::NodeEnd | RawFDTTokenData::FDTEnd => {
                    self.is_end = true;
                    break None;
                }
                RawFDTTokenData::Prop { name, value } => {
                    break Some(Child::Property(Property {
                        name,
                        value,
                        fdt_offset: next.offset,
                    }));
                }
            }
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
/// A wrapper around [`base::RawFDT`] that provides a higher-level API for accessing FDT nodes.
pub struct Fdt<'a> {
    raw_fdt: base::RawFDT<'a>,
}

impl<'a> Fdt<'a> {
    pub const fn from_base(raw_fdt: base::RawFDT<'a>) -> Self {
        Self { raw_fdt }
    }

    pub unsafe fn from_raw(raw_ptr: *const u8) -> Option<Self> {
        Some(Self {
            raw_fdt: unsafe { base::RawFDT::from_raw(raw_ptr)? },
        })
    }

    #[inline]
    pub fn find_compatible(&'a self, with: &[&str]) -> Option<Node<'a>> {
        self.find_node(|node| {
            node.compatible()
                .is_some_and(|mut comp| comp.any(|c| with.contains(&c)))
        })
    }

    /// Finds a node by its phandle.
    pub fn find_by_phandle(&'a self, phandle: u32) -> Option<Node<'a>> {
        self.find_node(|node| node.phandle().is_some_and(|n_phandle| n_phandle == phandle))
    }

    /// Iterates over all the nodes and nests in the FDT until the given predicate returns true.
    pub fn find_node<F>(&'a self, f: F) -> Option<Node<'a>>
    where
        F: Fn(&Node<'a>) -> bool,
    {
        fn iter_subnodes<'c, F: Copy + Fn(&Node<'c>) -> bool>(
            f: F,
            node: Node<'c>,
        ) -> Option<Node<'c>> {
            for subnode in node.subnodes() {
                if f(&subnode) {
                    return Some(subnode);
                }
                if let Some(result) = iter_subnodes(f, subnode) {
                    return Some(result);
                }
            }

            None
        }

        for node in self.nodes() {
            if f(&node) {
                return Some(node);
            }

            if let Some(result) = iter_subnodes(&f, node) {
                return Some(result);
            }
        }
        None
    }

    /// Returns the root node of the FDT.
    ///
    /// Panics if the fdt is corrupted.
    pub fn root_node(&'a self) -> Node<'a> {
        self.nodes().next().expect("No root node in the FDT")
    }

    /// Returns an iterator over the nodes in the FDT.
    ///
    /// Doesn't nest over subnodes, so this iterator will only yield top-level nodes.
    fn nodes(&'a self) -> impl Iterator<Item = Node<'a>> + 'a {
        NodeChildrenIter {
            parent_offset: None,
            fdt: self,
            fdt_iter: self.raw_fdt.iter(),
            is_end: false,
        }
        .filter_map(|child| match child {
            Child::Subnode(node) => Some(node),
            _ => None,
        })
    }
}
