//! ARM device tree utlities

use fdt_rs::{
    base::DevTree,
    index::{iters::DevTreeIndexNodeSiblingIter, DevTreeIndex, DevTreeIndexNode, DevTreeIndexProp},
    prelude::PropReader,
    spec::Phandle,
};
use libk_mm::{address::PhysicalAddress, phys::PhysicalMemoryRegion};
use yggdrasil_abi::error::Error;

const INDEX_BUFFER_SIZE: usize = 65536;

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

static mut FDT_INDEX_BUFFER: FdtIndexBuffer = FdtIndexBuffer::zeroed();

impl FdtIndexBuffer {
    const fn zeroed() -> Self {
        Self([0; INDEX_BUFFER_SIZE])
    }
}

/// Device tree node
pub type TNode<'a> = DevTreeIndexNode<'a, 'a, 'a>;
/// Device tree property
pub type TProp<'a> = DevTreeIndexProp<'a, 'a, 'a>;

/// Extension trait for [DevTreeIndexNode] to obtain typed property values
pub trait DevTreeIndexNodePropGet<T> {
    /// Returns a property value of given type, if it exists
    fn prop(&self, name: &str) -> Option<T>;
}

/// Device tree wrapper struct
pub struct DeviceTree<'a> {
    tree: DevTree<'a>,
    index: DevTreeIndex<'a, 'a>,
}

impl<'a> DeviceTree<'a> {
    pub fn find(&'a self, name: &str) -> Option<TNode<'a>> {
        let path = self.resolve_alias(name)?;
        self.find_absolute(path)
    }
}

impl DevTreeIndexNodePropGet<u32> for DevTreeIndexNode<'_, '_, '_> {
    fn prop(&self, name: &str) -> Option<u32> {
        self.props().find_map(|prop| {
            if prop.name().ok()? == name {
                prop.u32(0).ok()
            } else {
                None
            }
        })
    }
}

impl<'a> DevTreeIndexNodePropGet<&'a str> for DevTreeIndexNode<'a, '_, '_> {
    fn prop(&self, name: &str) -> Option<&'a str> {
        self.props().find_map(|prop| {
            if prop.name().ok()? == name {
                prop.str().ok()
            } else {
                None
            }
        })
    }
}

pub fn dump_node<F: Fn(&str)>(node: &TNode, print: &F, indent: usize) {
    fn tab<F: Fn(&str)>(print: &F, indent: usize) {
        for _ in 0..indent {
            print("  ");
        }
    }

    let node_name = node.name().unwrap_or("<unknown>");
    tab(print, indent);
    print(node_name);
    print(" {\r\n");
    for prop in node.props() {
        let name = prop.name().unwrap_or("<unknown>");
        tab(print, indent + 1);
        print(name);
        print(": ");
        match name {
            _ if node_name == "aliases" => {
                let model = prop.str().unwrap_or("<invalid>");
                print("\"");
                print(model);
                print("\"");
            }
            "model" => {
                let model = prop.str().unwrap_or("<invalid>");
                print("\"");
                print(model);
                print("\"");
            }
            _ => {
                print("???");
            }
        }
        print("\r\n");
    }
    for child in node.children() {
        dump_node(&child, print, indent + 1);
    }
    tab(print, indent);
    print("}\r\n");
}

/// Looks up a property with given name in the node
pub fn find_prop<'a>(node: &TNode<'a>, name: &str) -> Option<TProp<'a>> {
    node.props().find(|p| p.name().unwrap_or("") == name)
}

fn find_interrupt_controller<'a>(of: &TNode<'a>) -> Option<TNode<'a>> {
    // Step 1. Find interrupt-parent property in the nearest ancestor
    let interrupt_parent: Option<Phandle> = of.prop("interrupt-parent");

    if let Some(phandle) = interrupt_parent {
        of.index().nodes().find(|node| {
            let value: Option<u32> = node.prop("phandle");
            value.map(|value| value == phandle).unwrap_or(false)
        })
    } else if let Some(parent) = of.parent() {
        find_interrupt_controller(&parent)
    } else {
        // No parents, no interrupt-controller property found
        None
    }
}

#[cfg(test)]
mod tests {
    use alloc::{boxed::Box, vec::Vec};

    use crate::dt::DevTreeIndexNodeExt;

    use super::DeviceTree;

    #[repr(C, align(0x1000))]
    struct Buffer<A, D: ?Sized> {
        _align: [A; 0],
        data: D,
    }

    #[repr(C, align(0x1000))]
    struct PageAlign;

    static TEST1_DTB: &'static Buffer<PageAlign, [u8]> = &Buffer {
        _align: [],
        data: *include_bytes!("../dtbs/test1.dtb"),
    };

    fn test1_dtb() -> DeviceTree<'static> {
        let buffer = &mut Box::leak(Box::new(Buffer::<PageAlign, _> {
            _align: [],
            data: [0; 65536],
        }))
        .data;

        unsafe { DeviceTree::from_addr_with_index(TEST1_DTB.data.as_ptr().addr(), buffer) }
    }

    #[test]
    fn test_absolute_path_lookup() {
        let dtb = test1_dtb();

        // Non-aliased lookups

        // Root
        let root = dtb.node_by_path("/").unwrap();
        assert!(root.parent().is_none());
        assert_eq!(root.name().unwrap(), "");

        // Nested lookup
        let soc = dtb.node_by_path("/soc").unwrap();
        let uart = dtb.node_by_path("/soc/serial@10000000").unwrap();
        assert!(soc.is_parent_of(&uart));
        assert!(root.is_parent_of(&soc));
        assert_eq!(soc.name().unwrap(), "soc");
    }

    #[test]
    fn test_aliased_lookup() {
        let dtb = test1_dtb();

        let serial = dtb.node_by_path("serial").unwrap();
        let uart = dtb.node_by_path("/soc/serial@10000000").unwrap();
        assert!(serial == uart);
        assert_eq!(serial.name().unwrap(), "serial@10000000");
    }

    #[test]
    fn test_chosen_stdout() {
        let dtb = test1_dtb();

        let (stdout, settings) = dtb.chosen_stdout().unwrap();
        assert_eq!(stdout, "serial");
        assert_eq!(settings, Some("115200n8"));
    }

    #[test]
    fn test_find_interrupt_controller() {
        let dtb = test1_dtb();

        let serial = dtb.node_by_path("serial").unwrap();
        let serial_interrupt_parent = serial.interrupt_parent().unwrap();
        let intc = dtb
            .node_by_path("/soc/interrupt-controller@10001000")
            .unwrap();

        assert!(intc == serial_interrupt_parent);
    }
}
