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

use alloc::{
    sync::{Arc, Weak},
    vec::Vec,
};
use device_api::{
    bus::Bus,
    clock::{ClockController, ClockHandle, ResetHandle},
    device::{Device, DeviceInitContext},
    gpio::PinHandle,
    interrupt::{ExternalInterruptController, IrqHandle, 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::{
    driver::{
        controller::GpioIter, traits::DeviceTreePinController, DeviceTreeGpioPins, DeviceTreeSyscon,
    },
    tree, DeviceTree, DeviceTreeNodeExt, DeviceTreePropertyRead, TNode, TProp,
};

use super::{
    controller::{ClockIter, ResetIter},
    lookup_phandle, map_interrupt,
    registry::{register_phandle, DEVICE_TREE, DRIVERS, ROOT},
    DeviceTreeClockController, DeviceTreeInterruptController, DeviceTreeResetController, Driver,
    InitSequence, 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<TProp<'static>>,

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

    // Driver/device info
    device: OneTimeInit<NodeDevice>,
    pub(crate) interrupt_controller: OneTimeInit<Arc<dyn DeviceTreeInterruptController>>,
    pub(crate) msi_controller: OneTimeInit<Arc<dyn MessageInterruptController>>,
    pub(crate) clock_controler: OneTimeInit<Arc<dyn DeviceTreeClockController>>,
    pub(crate) reset_controller: OneTimeInit<Arc<dyn DeviceTreeResetController>>,
    pub(crate) system_controller: OneTimeInit<Arc<dyn DeviceTreeSyscon>>,
    pub(crate) pin_controller: OneTimeInit<Arc<dyn DeviceTreePinController>>,
}

pub(crate) struct ProbedDevice {
    driver: &'static dyn Driver,
    sequence: InitSequence,
    init_token: OneTimeInit<Result<(), Error>>,
    irq_init_token: OneTimeInit<Result<(), Error>>,
    device: Arc<dyn Device>,
}

enum NodeDevice {
    // Node probed, no device found
    Missing,
    // Node probed and driver found
    Present(ProbedDevice),
}

struct EnumerationContext {
    address_cells: usize,
    size_cells: usize,

    interrupt_parent: Option<Phandle>,
}

impl NodeDevice {
    fn as_probed(&self) -> Option<&ProbedDevice> {
        match self {
            Self::Missing => None,
            Self::Present(probed) => Some(probed),
        }
    }
}

impl Node {
    fn probe_single(node: &Arc<Node>, cx: &mut ProbeContext) -> Option<ProbedDevice> {
        let name = node.name();
        let compatible = node.compatible.as_ref()?;
        let drivers = DRIVERS.read();

        // Find a matching driver
        let driver = compatible
            .as_str_list()
            .find_map(|c| drivers.iter().find(|d| d.matches(c)));

        if libk::config::get().device_tree.log_missing {
            if driver.is_none() {
                // FIXME don't spam virtio missing stuff
                if !name.map_or(false, |n| n.starts_with("virtio_mmio")) {
                    for (i, compatible) in compatible.as_str_list().enumerate() {
                        if i == 0 {
                            log::warn!("No driver for {name:?} ({compatible:?})");
                        } else {
                            log::warn!("  also {compatible:?}");
                        }
                    }
                }
            }
        }
        let driver = driver?;

        // Initialize default pinctrl before probing
        node.setup_pins()
            .inspect_err(|e| log::error!("{name:?}: pinctrl init error {e:?}"))
            .ok()?;

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

        // Move to early init unless driver wants to sleep
        let sequence = cx.sequence.unwrap_or(InitSequence::Early);

        device.map(|device| ProbedDevice {
            driver: driver.imp,
            sequence,
            device,
            init_token: OneTimeInit::new(),
            irq_init_token: OneTimeInit::new(),
        })
    }

    fn probe_upwards(self: &Arc<Self>) -> (Option<&ProbedDevice>, 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 mut cx = ProbeContext {
            bus: parent_bus.clone(),
            sequence: None,
        };

        let inner = self
            .device
            .or_init_with(|| match Self::probe_single(self, &mut cx) {
                Some(probed) => NodeDevice::Present(probed),
                None => NodeDevice::Missing,
            });
        let probed = inner.as_probed();

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

        (probed, 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(crate) fn probe(self: &Arc<Self>) -> Option<&ProbedDevice> {
        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);
    }

    /// Informs the node of its capability as a reset controller, allowing `resets` to be mapped
    /// through this device.
    pub fn make_reset_controller(&self, rstc: Arc<dyn DeviceTreeResetController>) {
        self.reset_controller.init(rstc);
    }

    /// Informs the node of its capability as a clock controller, allowing `clocks` to be mapped
    /// through this device.
    pub fn make_clock_controller(&self, clkc: Arc<dyn DeviceTreeClockController>) {
        self.clock_controler.init(clkc);
    }

    /// Informs the node of its capability as a system controller.
    pub fn make_system_controller(&self, syscon: Arc<dyn DeviceTreeSyscon>) {
        self.system_controller.init(syscon);
    }

    /// Informs the node of its capability as a pin controller.
    pub fn make_pin_controller(&self, pinctrl: Arc<dyn DeviceTreePinController>) {
        self.pin_controller.init(pinctrl);
    }

    /// Returns the device driver associated with this node, if any was probed.
    pub fn driver(&self) -> Option<&'static dyn Driver> {
        let probed = self.device.try_get()?.as_probed()?;
        Some(probed.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>, sequence: InitSequence) -> Option<Result<(), Error>> {
        let probed = self.probe()?;

        // Needs to be initialized later
        if probed.sequence > sequence {
            return None;
        }

        let state = probed.init_token.or_init_with(|| {
            let cx = self.make_init_context();
            unsafe { probed.device.clone().init(cx) }
        });

        Some(*state)
    }

    /// 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>, sequence: InitSequence) -> Option<Result<(), Error>> {
        let probed = match self.clone().lazy_init(sequence) {
            Some(Ok(())) => {
                let device = self.device.get();
                device.as_probed()?
            }
            Some(Err(_)) | None => return None,
        };

        let state = probed
            .irq_init_token
            .or_init_with(|| unsafe { probed.device.clone().init_irq() });

        Some(*state)
    }

    /// "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>, sequence: InitSequence) -> Result<(), Error> {
        let probed = self.probe().ok_or(Error::DoesNotExist)?;

        // Needs to be initialized later
        if probed.sequence > sequence {
            return Ok(());
        }

        let state = probed.init_token.try_init_with(|| {
            let cx = self.make_init_context();
            unsafe { probed.device.clone().init(cx) }
        })?;

        *state
    }

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

    /// 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))
    }

    /// Returns an iterator over the node's defined `gpios`
    pub fn gpios<'a>(
        &self,
        options: &'a DeviceTreeGpioPins,
    ) -> Option<impl Iterator<Item = PinHandle> + 'a> {
        let gpios = self.property("gpios")?;
        Some(GpioIter {
            gpios,
            options,
            offset: 0,
        })
    }

    /// Returns a `gpios` entry with given index
    pub fn gpio(&self, options: &DeviceTreeGpioPins, index: usize) -> Option<PinHandle> {
        self.gpios(options)?.nth(index)
    }

    /// Returns an input `clock` handle for a given reset name
    pub fn named_clock(&self, name: &str) -> Option<ClockHandle> {
        self.property("clock-names")?
            .as_str_list()
            .position(|n| n == name)
            .and_then(|i| self.clock(i))
    }

    /// Returns a `reset` handle for a given reset name
    pub fn named_reset(&self, name: &str) -> Option<ResetHandle> {
        self.property("reset-names")?
            .as_str_list()
            .position(|n| n == name)
            .and_then(|i| self.reset(i))
    }

    /// Returns an iterator over the node's input clocks
    pub fn clocks(&self) -> Option<impl Iterator<Item = ClockHandle>> {
        let clocks = self.property("clocks")?;
        Some(ClockIter { clocks, offset: 0 })
    }

    /// Returns an iterator over the node's resets
    pub fn resets(&self) -> Option<impl Iterator<Item = ResetHandle>> {
        let resets = self.property("resets")?;
        Some(ResetIter { resets, offset: 0 })
    }

    /// Returns the `index`th input clock of the node
    pub fn clock(&self, index: usize) -> Option<ClockHandle> {
        self.clocks()?.nth(index)
    }

    /// Returns the `index`th reset of the node
    pub fn reset(&self, index: usize) -> Option<ResetHandle> {
        self.resets()?.nth(index)
    }

    /// Reads interrupt information from `interrupts[index]` property, mapped by the node's
    /// `interrupt-parent`.
    pub fn interrupt(&self, index: usize) -> Option<IrqHandle> {
        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<IrqHandle> {
        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>> {
        todo!()
        // 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())
    }

    /// Attempts to get a system controller represented by this node, if any
    pub fn as_system_controller(self: &Arc<Self>) -> Option<Arc<dyn DeviceTreeSyscon>> {
        self.system_controller.try_get().cloned()
    }

    /// Attempts to get a pin controller represented by this node, if any
    pub fn as_pin_controller(&self) -> Option<Arc<dyn DeviceTreePinController>> {
        self.pin_controller.try_get().cloned()
    }

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

    /// Returns the `#size-cells` value of the node's parent bus
    pub fn bus_size_cells(&self) -> usize {
        self.bus_size_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)
    }

    /// Returns `true` if the node has a property `name` inside
    pub fn has_property(&self, name: &str) -> bool {
        self.dt_node.has_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)
    }

    /// Interprets property `name` as a string value
    pub fn prop_str(&self, name: &str) -> Option<&'static str> {
        self.dt_node.prop_string(name)
    }

    /// Returns `true` if the node contains `query` in its `compatible` property.
    pub fn is_compatible(&self, query: &str) -> bool {
        let Some(compatible) = self.compatible.as_ref() else {
            return false;
        };

        compatible.as_str_list().any(|c| c == query)
    }

    /// Looks up the first node matching `compatible` query.
    pub fn by_compatible(compatible: &str, ensure_probed: bool) -> Option<Arc<Self>> {
        let node = walk_device_tree(|node| {
            if node.is_compatible(compatible) {
                Some(node.clone())
            } else {
                None
            }
        })?;

        if ensure_probed {
            node.probe()?;
        }

        Some(node)
    }

    /// Dump the contents of the device tree node into Info-level log
    pub fn dump(&self) {
        tree::dump(self.dt_node.clone(), 0);
    }

    fn setup_pins(&self) -> Result<(), Error> {
        // TODO lookup pin state by name
        let Some(pinctrl0) = self.prop_usize("pinctrl-0") else {
            return Ok(());
        };
        let pinctrl0 = lookup_phandle(pinctrl0 as Phandle, false).ok_or(Error::DoesNotExist)?;

        // Find pinctrl-group's parent pin controller
        let pin_controller = pinctrl0.parent().ok_or(Error::DoesNotExist)?;
        pin_controller.probe().ok_or(Error::DoesNotExist)?;
        let pin_controller = pin_controller
            .as_pin_controller()
            .ok_or(Error::DoesNotExist)?;

        pin_controller.configure_pin_group(&pinctrl0)
    }
}

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();
    let compatible = dt.property("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(),

        interrupt_controller: OneTimeInit::new(),
        msi_controller: OneTimeInit::new(),
        clock_controler: OneTimeInit::new(),
        reset_controller: OneTimeInit::new(),
        system_controller: OneTimeInit::new(),
        pin_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();
    }
}
