//! Finds the board's devices (`types::bootinfo`): disk, GPU, and keyboard
//! (or the UART when there is none).

use hw_virtio::block::CONFIG_CAPACITY;
use hw_virtio::mmio::{
    DEVICE_ID_BLOCK, DEVICE_ID_GPU, DEVICE_ID_INPUT, DEVICE_ID_RNG, MAGIC, REG_DEVICE_ID, REG_MAGIC,
};
use libkos::device::{self, Device};
use libkos::syscall;
use types::badge::Badge;
use types::address::VirtualAddress;
use types::bootinfo::{BootInfo, KIND_UART_16550, KIND_VIRTIO, NO_SLOT};
use types::bounded::CapSlot;
use types::permission::{Permission, PermissionSet};

/// Where init briefly maps a device to ask what it is (top of the window).
const PROBE_VA: VirtualAddress = VirtualAddress::from_raw(0x01F0_0000);

pub struct Devices {
    pub disk: Device,
    pub gpu: Device,
    pub input: Device,
    /// The gatekeeper's: a small disk of its own for the key slots.
    pub keystore: Device,
    pub random: Device,
}

/// A block device this small (1 MiB) is the keystore, not the disk.
const KEYSTORE_MAX_SECTORS: u64 = 2048;

fn slot(raw: u64) -> Option<CapSlot> {
    if raw == NO_SLOT {
        return None;
    }
    CapSlot::new(u8::try_from(raw).ok()?)
}

/// The virtio device id behind `region` (0 for an empty slot) and, for a
/// block device, its size in sectors. Read through a copy of the
/// capability, mapped only while looking and then dropped, which unmaps it.
fn virtio_id(region: CapSlot) -> (u32, u64) {
    let perms = PermissionSet::empty().with(Permission::Read).with(Permission::Write);
    let Ok(copy) = syscall::mint(region, Badge::NONE, perms) else {
        return (0, 0);
    };
    let id = match device::map_mmio(copy, PROBE_VA) {
        // The registers stay mapped until the drop below.
        Ok(regs) => match regs.read::<u32>(REG_MAGIC) == MAGIC {
            true => {
                let id = regs.read(REG_DEVICE_ID);
                let size = match id {
                    DEVICE_ID_BLOCK => {
                        regs.read::<u32>(CONFIG_CAPACITY) as u64 | (regs.read::<u32>(CONFIG_CAPACITY + 4) as u64) << 32
                    }
                    _ => 0,
                };
                (id, size)
            }
            false => (0, 0),
        },
        Err(_) => (0, 0),
    };
    syscall::drop_cap(copy);
    id
}

/// What boot left for init; readers check its magic.
fn boot_info() -> BootInfo {
    libkos::boot_info()
}

/// The capability that halts the machine, if boot said where it is.
pub fn system_control() -> Option<CapSlot> {
    let info = boot_info();
    if info.magic != types::bootinfo::MAGIC {
        return None;
    }
    slot(info.system_control_slot)
}

/// The console capability, until init hands it to the log server.
pub fn console() -> Option<CapSlot> {
    let info = boot_info();
    if info.magic != types::bootinfo::MAGIC {
        return None;
    }
    slot(info.console_slot)
}

/// Sorts the board's devices out.
pub fn discover() -> Devices {
    let info = boot_info();
    let mut found = Devices {
        disk: Device::NONE,
        gpu: Device::NONE,
        input: Device::NONE,
        keystore: Device::NONE,
        random: Device::NONE,
    };
    let mut uart = Device::NONE;
    for d in info.devices() {
        let Some(region) = slot(d.region_slot) else {
            continue;
        };
        let dev = Device { mmio: Some(region), irq: slot(d.irq_slot), kind: d.kind as usize };
        let place = match d.kind {
            KIND_UART_16550 => &mut uart,
            KIND_VIRTIO => match virtio_id(region) {
                (DEVICE_ID_BLOCK, size) if size <= KEYSTORE_MAX_SECTORS => &mut found.keystore,
                (DEVICE_ID_BLOCK, _) => &mut found.disk,
                (DEVICE_ID_GPU, _) => &mut found.gpu,
                (DEVICE_ID_INPUT, _) => &mut found.input,
                (DEVICE_ID_RNG, _) => &mut found.random,
                _ => continue,
            },
            _ => continue,
        };
        // First device of each kind wins.
        if place.mmio.is_none() {
            *place = dev;
        }
    }
    if found.input.mmio.is_none() {
        found.input = uart;
    }
    found
}