//! Utility functions for device tree handling
use alloc::string::String;
use fdt_rs::index::iters::DevTreeIndexNodeSiblingIter;
use libk::{
    error::Error,
    fs::sysfs::{
        self,
        attribute::{BytesAttribute, BytesAttributeOps, StringAttribute, StringAttributeOps},
        object::KObject,
    },
};
use libk_mm::{address::PhysicalAddress, phys::PhysicalMemoryRegion};

use crate::{node::DeviceTreeNodeExt, property::DeviceTreePropertyRead, tree::DeviceTree};

/// Iterator for physical memory regions present in the device tree
#[derive(Clone)]
pub struct DeviceTreeMemoryRegionIter<'a> {
    inner: DevTreeIndexNodeSiblingIter<'a, 'a, 'a>,
}

/// Iterator for reserved physical memory regions
#[derive(Clone)]
pub struct DeviceTreeReservedRegionIter<'a> {
    inner: Option<DevTreeIndexNodeSiblingIter<'a, 'a, 'a>>,
}

impl<'a> DeviceTreeMemoryRegionIter<'a> {
    pub(crate) fn new(dt: &'a DeviceTree) -> Self {
        let inner = dt.root().children();
        Self { inner }
    }
}

impl Iterator for DeviceTreeMemoryRegionIter<'_> {
    type Item = PhysicalMemoryRegion;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            let Some(node) = self.inner.next() else {
                break None;
            };

            let name = node.name().unwrap_or("");
            if !name.starts_with("memory@") && name != "memory" {
                continue;
            }

            if let Some(reg) = node.property("reg") {
                let address_cells = node.parent_address_cells();
                let size_cells = node.parent_size_cells();

                if let Some((base, size)) = reg.read_cells(0, (address_cells, size_cells)) {
                    let base = PhysicalAddress::from_u64(base);
                    let size = size as usize;
                    break Some(PhysicalMemoryRegion { base, size });
                }
            }
        }
    }
}

impl<'a> DeviceTreeReservedRegionIter<'a> {
    pub(crate) fn new(dt: &'a DeviceTree) -> Self {
        let inner = dt.root().child("reserved-memory").map(|r| r.children());
        Self { inner }
    }
}

impl Iterator for DeviceTreeReservedRegionIter<'_> {
    type Item = PhysicalMemoryRegion;

    fn next(&mut self) -> Option<Self::Item> {
        let inner = self.inner.as_mut()?;

        loop {
            let Some(node) = inner.next() else {
                break None;
            };

            if let Some(reg) = node.property("reg") {
                let address_cells = node.parent_address_cells();
                let size_cells = node.parent_size_cells();

                if let Some((base, size)) = reg.read_cells(0, (address_cells, size_cells)) {
                    let base = PhysicalAddress::from_u64(base);
                    let size = size as usize;
                    break Some(PhysicalMemoryRegion { base, size });
                }
            }
        }
    }
}

/// Registers sysfs objects related to the device tree
pub fn create_sysfs_nodes(dt: &'static DeviceTree) {
    struct Raw;
    struct Compatible;
    struct Model;

    struct MachineNodeInfo {
        compatible: Option<&'static str>,
        model: Option<&'static str>,
    }

    impl BytesAttributeOps for Raw {
        type Data = &'static DeviceTree<'static>;
        const NAME: &'static str = "raw";

        fn read(state: &Self::Data, pos: usize, buffer: &mut [u8]) -> Result<usize, Error> {
            let data = state.as_bytes();
            if pos >= data.len() {
                return Ok(0);
            }
            let amount = (data.len() - pos).min(buffer.len());
            buffer[..amount].copy_from_slice(&data[pos..pos + amount]);
            Ok(amount)
        }

        fn size(state: &Self::Data) -> usize {
            state.size()
        }
    }

    impl StringAttributeOps for Compatible {
        type Data = MachineNodeInfo;
        const NAME: &'static str = "compatible";

        fn read(state: &Self::Data) -> Result<String, Error> {
            Ok(state.compatible.unwrap_or("").into())
        }
    }

    impl StringAttributeOps for Model {
        type Data = MachineNodeInfo;
        const NAME: &'static str = "model";

        fn read(state: &Self::Data) -> Result<String, Error> {
            Ok(state.model.unwrap_or("").into())
        }
    }

    if let Some(device) = sysfs::device() {
        let object = KObject::new(dt);

        object.add_attribute(BytesAttribute::from(Raw)).ok();

        let compatible = dt.root().prop_string("compatible");
        let model = dt.root().prop_string("model");

        let machine = KObject::new(MachineNodeInfo { compatible, model });

        machine
            .add_attribute(StringAttribute::from(Compatible))
            .ok();
        machine.add_attribute(StringAttribute::from(Model)).ok();

        object.add_object("machine", machine).ok();

        device.add_object("device-tree", object).ok();
    }
}
