//! PCI bus device wrapper for device tree

use alloc::{collections::btree_map::BTreeMap, sync::Arc, vec::Vec};
use device_api::{device::Device, interrupt::Irq};
use device_tree::{
    driver::{
        device_tree_driver,
        util::{pcie_interrupt_map, pcie_msi_map},
        Node, ProbeContext,
    },
    DeviceTreePropertyRead,
};
use libk_mm::address::PhysicalAddress;
use ygg_driver_pci::{
    device::{PciInterrupt, PciInterruptPin, PciInterruptRoute},
    interrupt::{PciFixedMsiMap, PciFixedMsiMapping, PciInterruptMap, PciMsiMap},
    PciAddress, PciAddressRange, PciBusManager, PciRangeType,
};

#[derive(Debug)]
struct RangeProperties {
    space: PciRangeType,
    bus: u8,
}

impl From<u32> for RangeProperties {
    fn from(value: u32) -> Self {
        let s = (value >> 24) & 0x3;
        let bus = (value >> 16) as u8;
        let space = match s {
            0 => PciRangeType::Configuration,
            1 => PciRangeType::Io,
            2 => PciRangeType::Memory32,
            3 => PciRangeType::Memory64,
            _ => unreachable!(),
        };
        Self { space, bus }
    }
}

fn make_range(
    context: &ProbeContext,
    options: RangeProperties,
    pci_base: u64,
    host_base: u64,
    size: u64,
) -> Option<PciAddressRange> {
    let host_range = host_base..host_base + size;
    let host_base = PhysicalAddress::from_u64(context.map_range(host_range)?.start);
    let size = size.try_into().ok()?;

    Some(PciAddressRange {
        ty: options.space,
        bus_number: options.bus,
        host_base,
        pci_base,
        size,
    })
}

fn make_interrupt_map(node: &Arc<Node>) -> Option<PciInterruptMap> {
    let interrupt_map = pcie_interrupt_map(node)?;
    let mut imap = BTreeMap::new();

    for entry in interrupt_map {
        let address = PciAddress::for_function(0, entry.bus, entry.device, entry.function);
        let pin = PciInterruptPin::try_from(entry.pin as u32).ok()?;
        let options = entry.irq.options;
        let Irq::External(irq) = entry.irq.irq else {
            continue;
        };

        let src = PciInterrupt { address, pin };
        // TODO specify interrupt controller for PCIe interrupts
        let dst = PciInterruptRoute {
            number: irq,
            options,
        };
        imap.insert(src, dst);
    }

    Some(PciInterruptMap::Fixed(imap))
}

fn make_msi_map(node: &Arc<Node>) -> Option<PciMsiMap> {
    let msi_map = pcie_msi_map(node)?;
    let mut entries = Vec::new();

    for entry in msi_map {
        let (bus, device, function) = entry.base_rid.as_parts();
        let start_address = PciAddress {
            // TODO
            segment: 0,
            bus,
            device,
            function,
        };
        let (bus, device, function) = entry.base_rid.offset(entry.length).as_parts();
        let end_address = PciAddress {
            // TODO
            segment: 0,
            bus,
            device,
            function,
        };

        entries.push(PciFixedMsiMapping {
            start_address,
            end_address,
            controller: entry.msi_controller,
        });
    }

    if entries.is_empty() {
        None
    } else {
        Some(PciMsiMap::Fixed(PciFixedMsiMap { entries }))
    }
}

device_tree_driver! {
    compatible: ["pci-host-ecam-generic"],
    driver: {
        fn probe(&self, node: &Arc<Node>, context: &mut ProbeContext) -> Option<Arc<dyn Device>> {
            let ecam_base = node.map_base(context, 0)?;
            let ranges = node.property("ranges")?;

            let parent_address_cells = node.bus_address_cells();
            let child_address_cells = node.self_address_cells()?.checked_sub(1)?;
            let child_size_cells = node.self_size_cells()?;

            // Extract interupt map
            let interrupt_map = make_interrupt_map(node)?;

            let cell_sizes = (1, child_address_cells, parent_address_cells, child_size_cells);

            // Extract PCI ranges
            let mut bus_ranges = Vec::new();
            for (flags, child_base, parent_base, size) in ranges.iter_cells(cell_sizes) {
                let options = RangeProperties::from(flags as u32);

                if let Some(range) = make_range(context, options, child_base, parent_base, size) {
                    log::info!(
                        "PCIe: {:#x?} -> {}:{:#x?}, {:?}",
                        range.host_base.into_u64()..range.host_base.into_u64() + range.size as u64,
                        range.bus_number,
                        range.pci_base..range.pci_base + range.size as u64,
                        range.ty
                    );
                    bus_ranges.push(range);
                }
            }

            let msi_map = match make_msi_map(node) {
                Some(map) => map,
                None => PciMsiMap::Legacy
            };

            PciBusManager::add_segment_from_device_tree(
                ecam_base,
                // TODO bus-range
                0..255,
                bus_ranges,
                interrupt_map,
                msi_map,
            ).ok();

            None
        }
    }
}
