//! Device hierachy management code
use core::{fmt, ops::Range};

use alloc::{
    sync::{Arc, Weak},
    vec::Vec,
};
use device_api::{
    bus::Bus,
    clock::ClockController,
    device::{Device, DeviceInitContext},
    interrupt::{ExternalInterruptController, FullIrq, MessageInterruptController},
};
use fdt_rs::spec::Phandle;
use libk::dma::DummyDmaAllocator;
use libk_mm::address::PhysicalAddress;
use libk_util::OneTimeInit;
use yggdrasil_abi::error::Error;

use crate::{DeviceTree, DeviceTreeNodeExt, DeviceTreePropertyRead, TNode, TProp};

use super::{
    lookup_phandle, map_interrupt,
    registry::{register_phandle, DEVICE_TREE, DRIVERS, ROOT},
    DeviceTreeInterruptController, Driver, ProbeContext,
};

/// Represents a single node in the device tree, which may or may not:
///
/// * Have a device probed (see [Node::probe])
/// * Have a device initialized (see [Node::lazy_init])
/// * Be a bus
/// * Be an interrupt-controller
pub struct Node {
    // Info from device tree
    dt_node: TNode<'static>,
    bus_address_cells: usize,
    bus_size_cells: usize,
    interrupt_parent: Option<Phandle>,
    name: Option<&'static str>,
    compatible: Option<&'static str>,

    // Hierachy info
    children: OneTimeInit<Vec<Arc<Node>>>,
    parent: Option<Weak<Node>>,

    // Driver/device info
    device: OneTimeInit<NodeDevice>,
    init_token: OneTimeInit<()>,
    pub(crate) interrupt_controller: OneTimeInit<Arc<dyn DeviceTreeInterruptController>>,
    pub(crate) msi_controller: OneTimeInit<Arc<dyn MessageInterruptController>>,
}

enum NodeDevice {
    // Node probed, no device found
    Missing,
    // Node probed and driver found
    Present {
        driver: &'static dyn Driver,
        device: Arc<dyn Device>,
    },
}

// struct NodeDevice {
//     driver: &'static dyn Driver,
//     device: Arc<dyn Device>,
// }

struct EnumerationContext {
    address_cells: usize,
    size_cells: usize,

    interrupt_parent: Option<Phandle>,
}

impl NodeDevice {
    fn as_device(&self) -> Option<Arc<dyn Device>> {
        match self {
            Self::Missing => None,
            Self::Present { device, .. } => Some(device.clone()),
        }
    }

    fn driver(&self) -> Option<&'static dyn Driver> {
        match self {
            Self::Missing => None,
            Self::Present { driver, .. } => Some(*driver),
        }
    }
}

impl Node {
    fn probe_upwards(self: Arc<Self>) -> (Option<Arc<dyn Device>>, Option<Weak<dyn Bus>>) {
        let mut parent_bus = None;
        if let Some(parent) = self.parent.as_ref().and_then(Weak::upgrade) {
            let (_, bus) = parent.probe_upwards();
            parent_bus = bus;
        }

        // Don't even try if `status = "disabled"`
        if let Some(status) = self.property("status") {
            let status = status.as_str().unwrap_or("");
            if status == "disabled" {
                return (None, parent_bus);
            }
        }

        let cx = ProbeContext {
            bus: parent_bus.clone(),
        };

        let inner = self.device.or_init_with_opt(|| {
            let compatible = self.compatible?;
            let drivers = DRIVERS.read();
            let driver = drivers.iter().find(|d| d.matches(compatible));
            if driver.is_none() {
                // log::warn!("No driver for {compatible:?}");
            }
            let driver = driver?;

            let device = driver.imp.probe(&self, &cx);

            let slot = match device {
                Some(device) => NodeDevice::Present {
                    driver: driver.imp,
                    device,
                },
                None => NodeDevice::Missing,
            };
            Some(slot)
        });

        let device = inner.and_then(|d| d.as_device());

        let bus = if let Some(device) = device.as_ref() {
            device.clone().as_bus().as_ref().map(Arc::downgrade)
        } else {
            parent_bus
        };

        (device, bus)
    }

    /// Performs a "probe" of the node by looking up a matching driver. If the node
    /// has a parent, its parent is probed first.
    ///
    /// If the node has already been probed, the device reference is just returned instead.
    pub fn probe(self: Arc<Self>) -> Option<Arc<dyn Device>> {
        let (device, _) = self.probe_upwards();
        device
    }

    /// Returns the parent node of this node
    pub fn parent(&self) -> Option<Arc<Node>> {
        self.parent.as_ref().and_then(Weak::upgrade)
    }

    /// When called from an interrupt-controller driver, informs the node of its capability as
    /// an interrupt controller, allowing any other nodes which refer to it to map their
    /// interrupts.
    ///
    /// This function has to be called in the driver's `probe` impl to function properly.
    pub fn make_interrupt_controller(&self, intc: Arc<dyn DeviceTreeInterruptController>) {
        self.interrupt_controller.init(intc);
    }

    /// When called from an msi-controller driver, informs the node of its capability as an MSI
    /// controller.
    pub fn make_msi_controller(&self, intc: Arc<dyn MessageInterruptController>) {
        self.msi_controller.init(intc);
    }

    /// Returns the device driver associated with this node, if any was probed.
    pub fn driver(&self) -> Option<&'static dyn Driver> {
        self.device.try_get()?.driver()
    }

    /// Performs a lazy initialization of the node:
    ///
    /// 1. If the node hasn't yet been probed for a device, a driver is looked up
    ///    (see [Node::probe]).
    /// 2. If the node hasn't yet been initialized, initialization is performed.
    ///
    /// Returns:
    ///
    /// * `Some(Ok(()))` - device was probed and initialized (now or before).
    /// * `Some(Err(...))` - device was probed, but initialization failed.
    /// * `None` - no driver matched the device.
    pub fn lazy_init(self: Arc<Self>) -> Option<Result<(), Error>> {
        let device = self.clone().probe()?;
        let result = self.init_token.or_try_init_with(|| {
            let cx = self.make_init_context();
            unsafe { device.init(cx) }?;
            Ok(())
        });
        match result {
            Ok(()) => Some(Ok(())),
            Err(error) => Some(Err(error)),
        }
    }

    /// Performs an IRQ initialization for the node's device:
    ///
    /// 1. Does the same thing that [Node::lazy_init] does.
    /// 2. If [Node::lazy_init] succeeds and a device exists, calls [Device::init_irq].
    ///
    /// Return value is the same as in [Node::lazy_init].
    pub fn init_irqs(self: Arc<Self>) -> Option<Result<(), Error>> {
        match self.clone().lazy_init() {
            Some(Ok(())) => {
                let device = self.device.get();
                let status = unsafe { device.as_device()?.init_irq() };
                Some(status)
            }
            Some(Err(_)) | None => None,
        }
    }

    /// "Forces" an initialization of the device. This is a stricter version of [Node::lazy_init]
    /// which enforces that:
    ///
    /// 1. A device/driver actually exists for this node.
    /// 2. The node's device hasn't been initialized before.
    ///
    /// Returns:
    ///
    /// * `Ok(())` - if probe/init succeeded.
    /// * `Err(Error::DoesNotExist)` - couldn't find a device/driver for this node.
    /// * `Err(other)` - initialization failed.
    pub fn force_init(self: Arc<Self>) -> Result<(), Error> {
        let device = self
            .clone()
            .probe()
            .ok_or(Error::DoesNotExist)
            .inspect_err(|_| log::error!("Does not exist: probe({:?})", self.name))?;

        self.init_token.try_init_with_opt(|| {
            let cx = self.make_init_context();
            unsafe { device.init(cx) }?;
            Ok(())
        })?;
        Ok(())
    }

    fn make_init_context(&self) -> DeviceInitContext {
        let cx = DeviceInitContext {
            dma_allocator: Arc::new(DummyDmaAllocator),
        };

        cx
    }

    /// Returns an iterator over the node's children
    pub fn children(&self) -> impl Iterator<Item = &Arc<Node>> {
        self.children.get().iter()
    }

    // Properties

    /// Returns the node's name string
    pub fn name(&self) -> Option<&'static str> {
        self.name
    }

    /// Returns the addresses by which some part of the device is represented on its parent bus.
    /// The bus address is taken from `reg[index]` property and is expected to be a pair
    /// of `(#address-cells, #size-cells)`-sized elements.
    pub fn bus_range(&self, index: usize) -> Option<Range<u64>> {
        let reg = self.dt_node.property("reg")?;
        let (base, size) = reg.read_cells(index, (self.bus_address_cells, self.bus_size_cells))?;
        Some(base..base + size)
    }

    /// Returns the physical-address range used by some part of the device.
    ///
    /// This function is a result of [Node::bus_range], passed through its' parent bus mapper
    /// (or returned as-is if the device has no parent bus).
    pub fn map_range(&self, context: &ProbeContext, index: usize) -> Option<Range<u64>> {
        let bus_range = self.bus_range(index)?;
        context.map_range(bus_range)
    }

    /// Same as [Node::map_range], but only returns the base of the range, cast to
    /// [PhysicalAddress].
    pub fn map_base(&self, context: &ProbeContext, index: usize) -> Option<PhysicalAddress> {
        self.map_range(context, index)
            .map(|range| PhysicalAddress::from_u64(range.start))
    }

    /// Reads interrupt information from `interrupts[index]` property, mapped by the node's
    /// `interrupt-parent`.
    pub fn interrupt(&self, index: usize) -> Option<FullIrq> {
        let interrupts = self.property("interrupts")?;
        let phandle = self.interrupt_parent?;
        map_interrupt(phandle, &interrupts, index)
    }

    /// Same as [Node::interrupt], but allows specifying other property to read the information
    /// from.
    pub fn interrupt_from(&self, property: &TProp, index: usize) -> Option<FullIrq> {
        let phandle = self.interrupt_parent?;
        map_interrupt(phandle, property, index)
    }

    /// Returns a node reference to the device's `interrupt-parent`.
    pub fn interrupt_controller(&self) -> Option<Arc<Node>> {
        lookup_phandle(self.interrupt_parent?, true)
    }

    /// Attempts to get a clock controller represented by this node, if any
    pub fn as_clock_controller(&self) -> Option<Arc<dyn ClockController>> {
        let device = self.device.try_get()?;
        device.as_device()?.as_clock_controller()
    }

    /// Attempts to get an interrupt controller represented by this node, if any
    pub fn as_interrupt_controller(&self) -> Option<Arc<dyn ExternalInterruptController>> {
        self.interrupt_controller
            .try_get()
            .map(|e| e.clone().as_interrupt_controller())
    }

    /// Returns the `#address-cells` value of the node's parent bus
    pub fn bus_address_cells(&self) -> usize {
        self.bus_address_cells
    }

    /// Returns the `#interrupt-cells` of this node, if any
    pub fn self_interrupt_cells(&self) -> Option<usize> {
        self.prop_usize("#interrupt-cells")
    }

    /// Returns the `#address-cells` of this node, if any
    pub fn self_address_cells(&self) -> Option<usize> {
        self.prop_usize("#address-cells")
    }

    /// Returns the `#size-cells` of this node, if any
    pub fn self_size_cells(&self) -> Option<usize> {
        self.prop_usize("#size-cells")
    }

    /// Looks up a property `name` in this node
    pub fn property(&self, name: &str) -> Option<TProp<'static>> {
        self.dt_node.property(name)
    }

    /// Interprets property `name` as a single cell and casts it to usize
    pub fn prop_usize(&self, name: &str) -> Option<usize> {
        self.dt_node.prop_cell_usize(name)
    }
}

impl fmt::Debug for Node {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self.name() {
            Some(name) => write!(f, "<{name}>"),
            None => f.write_str("<unknown>"),
        }
    }
}

unsafe impl Sync for Node {}
unsafe impl Send for Node {}

fn unflatten_node(
    dt: TNode<'static>,
    parent: Option<Weak<Node>>,
    ctx: &EnumerationContext,
) -> Arc<Node> {
    let name = dt.name().ok();
    // TODO <stringlist>
    let compatible = dt.prop_string("compatible");

    // Extract #...-cells for children
    let address_cells = dt
        .prop_cell_usize("#address-cells")
        .unwrap_or(ctx.address_cells);
    let size_cells = dt.prop_cell_usize("#size-cells").unwrap_or(ctx.size_cells);

    // Extract interrupt-parent for children
    let interrupt_parent = dt.prop_phandle("interrupt-parent").or(ctx.interrupt_parent);

    let phandle = dt.prop_phandle("phandle");

    // Setup this node
    let node = Arc::new(Node {
        dt_node: dt.clone(),
        interrupt_parent,
        bus_address_cells: ctx.address_cells,
        bus_size_cells: ctx.size_cells,
        name,
        compatible,

        children: OneTimeInit::new(),
        parent,

        device: OneTimeInit::new(),
        init_token: OneTimeInit::new(),
        interrupt_controller: OneTimeInit::new(),
        msi_controller: OneTimeInit::new(),
    });

    if let Some(phandle) = phandle {
        register_phandle(phandle, node.clone());
    }

    // Unflatten children
    let mut children = Vec::new();
    let child_ctx = EnumerationContext {
        address_cells,
        size_cells,
        interrupt_parent,
    };
    for child in dt.children() {
        let child_node = unflatten_node(child, Some(Arc::downgrade(&node)), &child_ctx);
        children.push(child_node);
    }

    node.children.init(children);
    node
}

/// Sets up node hierarchy from a [DeviceTree] reference. This is a requirement to be able to
/// probe/initialize/manage devices using a device-tree approach.
pub fn unflatten_device_tree(dt: &'static DeviceTree) {
    const DEFAULT_ADDRESS_CELLS: u32 = 2;
    const DEFAULT_SIZE_CELLS: u32 = 1;

    DEVICE_TREE.init(dt);

    let ctx = EnumerationContext {
        address_cells: DEFAULT_ADDRESS_CELLS as usize,
        size_cells: DEFAULT_SIZE_CELLS as usize,

        interrupt_parent: None,
    };

    let root = unflatten_node(dt.root(), None, &ctx);

    ROOT.init(root);
}

fn walk_node<T, V: FnMut(&Arc<Node>) -> Option<T>>(node: &Arc<Node>, visitor: &mut V) -> Option<T> {
    if let Some(value) = visitor(node) {
        return Some(value);
    }

    for child in node.children.get() {
        if let Some(value) = walk_node(child, visitor) {
            return Some(value);
        }
    }

    None
}

/// Performs a walk of the device (in no particular order), passing each node to the visitor and
/// aborting if `Some(...)` is returned by it.
pub fn walk_device_tree<T, V: FnMut(&Arc<Node>) -> Option<T>>(mut visitor: V) -> Option<T> {
    if let Some(root) = ROOT.try_get() {
        walk_node(root, &mut visitor)
    } else {
        None
    }
}

/// Finds a node in the device tree by its absolute path or alias.
pub fn find_node(name: &str) -> Option<Arc<Node>> {
    let dt = *DEVICE_TREE.try_get()?;
    let mut path = dt.resolve_alias(name)?.trim_matches('/');
    let mut current = ROOT.try_get()?.clone();

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

        let (left, rest) = match path.split_once('/') {
            Some(elem) => elem,
            None => (path, ""),
        };
        path = rest.trim_start_matches('/');

        let child = current
            .children()
            .find(|node| node.name.map(|name| name == left).unwrap_or(false))?;
        current = child.clone();
    }
}