//! Raw device tree handling structures
use core::ptr;

use fdt_rs::{
    base::DevTree,
    index::{DevTreeIndex, DevTreeIndexNode, DevTreeIndexProp},
    prelude::PropReader,
};
use libk_mm::address::PhysicalAddress;
use yggdrasil_abi::error::Error;

use crate::{
    node::DeviceTreeNodeExt,
    property::DeviceTreePropertyRead,
    util::{DeviceTreeMemoryRegionIter, DeviceTreeReservedRegionIter},
};

const FDT_INDEX_BUFFER_SIZE: usize = 65536;

/// Helper alias for a device tree node.
pub type TNode<'a> = DevTreeIndexNode<'a, 'a, 'a>;
/// Helper alias for a device tree node's property.
pub type TProp<'a> = DevTreeIndexProp<'a, 'a, 'a>;

/// Main data structure for interacting with a raw device tree
pub struct DeviceTree<'a> {
    tree: DevTree<'a>,
    index: DevTreeIndex<'a, 'a>,
}

#[repr(C, align(0x10))]
struct FdtIndexBuffer([u8; FDT_INDEX_BUFFER_SIZE]);

impl<'a> DeviceTree<'a> {
    /// Minimum possible FDT header size
    pub const MIN_HEADER_SIZE: usize = DevTree::MIN_HEADER_SIZE;
    /// Default `#address-cells` value assumed if one is absent
    pub const DEFAULT_ADDRESS_CELLS: usize = 2;
    /// Default `#size-cells` value assumed if one is absent
    pub const DEFAULT_SIZE_CELLS: usize = 1;

    /// Constructs a device tree from a raw tree and an index
    pub fn from_parts(tree: DevTree<'a>, index: DevTreeIndex<'a, 'a>) -> Self {
        Self { tree, index }
    }

    /// Constructs a device tree representation from raw FDT address + index buffer.
    ///
    /// # Safety
    ///
    /// Accepts a raw address, unsafe.
    pub unsafe fn from_raw_with_index(address: usize, index: &'a mut [u8]) -> Result<Self, Error> {
        let tree = DevTree::from_raw_pointer(ptr::with_exposed_provenance(address))
            .map_err(|_| Error::InvalidArgument)?;
        let index = DevTreeIndex::new(tree, index).map_err(|_| Error::InvalidArgument)?;
        Ok(Self { tree, index })
    }

    /// Constructs a device tree representation from raw FDT address + a **static** index buffer.
    ///
    /// # Safety
    ///
    /// Accepts a raw address, uses a statically mutable index buffer, unsafe.
    pub unsafe fn from_raw(address: usize) -> Result<Self, Error> {
        static mut BUFFER: FdtIndexBuffer = FdtIndexBuffer([0; FDT_INDEX_BUFFER_SIZE]);
        #[allow(static_mut_refs)]
        Self::from_raw_with_index(address, &mut BUFFER.0[..])
    }

    /// Returns the total size in bytes of the underlying FDT blob
    pub fn size(&self) -> usize {
        self.tree.totalsize()
    }

    /// Returns the underlying FDT blob bytes
    pub fn as_bytes(&self) -> &'a [u8] {
        self.tree.buf()
    }

    /// Returns the root node of this device tree
    pub fn root(&self) -> TNode<'_> {
        self.index.root()
    }

    /// Dump the device tree into the debug log
    pub fn dump(&self) {
        dump(self.root(), 0)
    }

    /// Looks up a node by its `/absolute/path`.
    pub fn find_absolute(&'a self, path: &str) -> Option<TNode<'a>> {
        let mut path = path.trim_start_matches('/');
        let mut current = self.root();

        loop {
            if path.is_empty() {
                return Some(current);
            }

            let (head, tail) = match path.split_once('/') {
                Some((head, tail)) => (head, tail.trim_start_matches('/')),
                None => (path, ""),
            };
            path = tail;

            current = current
                .children()
                .find(|child| child.name().map(|n| n == head).unwrap_or(false))?;
        }
    }

    /// Resolves an alias to its real path by looking it up in the `/aliases` node of the
    /// device tree.
    pub fn resolve_alias(&'a self, name: &'a str) -> Option<&'a str> {
        if name.starts_with('/') {
            Some(name)
        } else {
            let aliases = self.find_absolute("/aliases")?;
            let alias = aliases.property(name)?;
            alias.str().ok()
        }
    }

    // Convenience methods

    /// Returns the /chosen.stdout-path value.
    /// The value is in the following form:
    ///
    /// ```ignore
    /// stdout-path = "<device-name>[:<settings>]";
    /// ```
    pub fn chosen_stdout(&self) -> Option<(&str, Option<&str>)> {
        let chosen = self.find_absolute("/chosen")?;
        let value = chosen.property("stdout-path")?.str().ok()?;

        match value.split_once(':') {
            Some((left, right)) => Some((left, Some(right))),
            None => Some((value, None)),
        }
    }

    /// Returns the start and end address of the initrd specified by this device tree.
    ///
    /// The initrd info is read from:
    ///
    /// * `/chosen/linux,initrd-start`
    /// * `/chosen/linux,initrd-end`
    pub fn chosen_initrd(&self) -> Option<(PhysicalAddress, PhysicalAddress)> {
        let root = self.root();
        let chosen = root.child("chosen")?;

        let address_cells = root
            .property("#address-cells")
            .and_then(|prop| prop.read_cell(0, 1))
            .unwrap_or(Self::DEFAULT_ADDRESS_CELLS as u64) as usize;

        let initrd_start = chosen.property("linux,initrd-start")?; // find_prop(&chosen, "linux,initrd-start")?;
        let initrd_end = chosen.property("linux,initrd-end")?;

        let initrd_start = initrd_start.read_cell(0, address_cells)?;
        let initrd_end = initrd_end.read_cell(0, address_cells)?;

        Some((
            PhysicalAddress::from_u64(initrd_start),
            PhysicalAddress::from_u64(initrd_end),
        ))
    }

    /// Returns the value of `/chosen/bootargs` interpreted as a string
    pub fn chosen_bootargs(&self) -> Option<&str> {
        self.root().child("chosen")?.prop_string("bootargs")
    }

    /// Returns an iterator over the memory regions specified by this device tree
    pub fn memory_regions(&self) -> DeviceTreeMemoryRegionIter<'_> {
        DeviceTreeMemoryRegionIter::new(self)
    }

    /// Returns an iterator over the reserved memory regions specified by this device tree
    pub fn reserved_regions(&self) -> DeviceTreeReservedRegionIter<'_> {
        DeviceTreeReservedRegionIter::new(self)
    }

    /// Returns the length of the header provided as a slice of bytes.
    ///
    /// # Safety
    ///
    /// Preconditions:
    ///
    /// * `header` is at least u32-aligned;
    /// * `header` points to a slice of at least [DevTree::MIN_HEADER_SIZE] bytes.
    pub unsafe fn read_totalsize(header: &[u8]) -> Result<usize, Error> {
        DevTree::read_totalsize(header).map_err(|_| Error::InvalidArgument)
    }
}

unsafe impl Sync for DeviceTree<'_> {}

fn indent(amount: usize) {
    log::info!(target: ":raw", "{0:1$}", "", amount * 2);
}

fn dump_property(node: &TNode, prop: TProp, level: usize) {
    enum Type {
        String,
        Cells,
        Bytes,
    }

    let node_name = node.name().unwrap_or("");
    let name = prop.name().unwrap_or("<unknown-prop>");

    indent(level);
    log::info!(target: ":raw", "{name}");

    if prop.len() == 0 {
        log::info!(target: ":raw", ";\n");
        return;
    }

    log::info!(target: ":raw", " = ");

    let ty = match name {
        "model" | "compatible" | "clock-output-names" | "clock-names" | "device_type"
        | "enable-method" | "status" | "label" | "method" => Type::String,
        "stdout-path" | "bootargs" if node_name == "chosen" => Type::String,
        _ if node_name == "aliases" => Type::String,
        _ if prop.len() % size_of::<u32>() == 0 => Type::Cells,
        _ => Type::Bytes,
    };

    match ty {
        Type::String => {
            // "..."
            for (i, string) in prop.as_str_list().enumerate() {
                if i != 0 {
                    log::info!(target: ":raw", ", ");
                }
                log::info!(target: ":raw", "{string:?}");
            }
        }
        Type::Cells => {
            // <...>
            log::info!(target: ":raw", "<");
            let mut i = 0;
            while let Some(value) = prop.read_cell(i, 1) {
                if i != 0 {
                    log::info!(target: ":raw", ", ");
                }
                log::info!(target: ":raw", "{value:#x}");
                i += 1;
            }
            log::info!(target: ":raw", ">");
        }
        Type::Bytes => {
            // <...>
            log::info!(target: ":raw", "<");
            for (i, byte) in prop.raw().iter().enumerate() {
                if i != 0 {
                    log::info!(target: ":raw", ", ");
                }
                log::info!(target: ":raw", "{byte:#x}");
            }
            log::info!(target: ":raw", ">");
        }
    }

    log::info!(target: ":raw", ";\n");
}

pub(crate) fn dump(node: TNode, level: usize) {
    let name = node.name().unwrap_or("<unknown>");

    indent(level);
    if name.is_empty() {
        log::info!(target: ":raw", "{{\n");
    } else {
        log::info!(target: ":raw", "{name}: {{\n");
    }

    let mut do_break = false;
    for property in node.props() {
        dump_property(&node, property, level + 1);
        do_break = true;
    }

    for child in node.children() {
        if do_break {
            log::info!(target: ":raw", "\n");
            do_break = false;
        }
        dump(child, level + 1);
    }

    indent(level);
    log::info!(target: ":raw", "}}\n");
}
