//! Probe and restore two real QEMU VirtIO PCI memory BARs.

use super::*;
use os_pci_discovery::{ClaimTable, EcamWindow};
use os_pci_resources::{Error as BarError, RangeTable, Width,
    decode_memory_bar_sizing};

pub fn run(boot: &BootData) -> Result<(), &'static str> {
    let blob = unsafe { core::slice::from_raw_parts(boot.dtb_address as *const u8, boot.dtb_size) };
    let host = dtb::discover_pci_ecam(blob).map_err(|_| "k1445-dtb-host")?;
    let window = EcamWindow::new(host.base, host.size, host.first_bus, host.last_bus)
        .map_err(|_| "k1445-ecam-window")?;
    let mut configs = ClaimTable::<1>::new();
    let leaf = configs.claim(window, 2, 1, 0, 0)
        .map_err(|_| "k1445-leaf-claim")?;
    if leaf.physical != host.base + 0x208000 {
        return Err("k1445-leaf-address");
    }
    let base = leaf.physical;
    // SAFETY: mode 80 maps this DTB-bounded ECAM page as Device-nGnRnE
    // before TTBR0. The isolated guest owns this disposable VirtIO PCI leaf.
    let identity = unsafe { read_volatile(base as *const u32) };
    if identity & 0xffff != 0x1af4 { return Err("k1445-leaf-identity"); }
    let command = unsafe { read_volatile((base + 0x04) as *const u16) };
    let bar1 = unsafe { read_volatile((base + 0x14) as *const u32) };
    let bar4 = unsafe { read_volatile((base + 0x20) as *const u32) };
    let bar5 = unsafe { read_volatile((base + 0x24) as *const u32) };
    let (mask1, mask4, mask5);
    unsafe {
        write_volatile((base + 0x04) as *mut u16, command & !0x3);
        asm!("dsb sy", options(nostack, preserves_flags));
        write_volatile((base + 0x14) as *mut u32, u32::MAX);
        asm!("dsb sy", options(nostack, preserves_flags));
        mask1 = read_volatile((base + 0x14) as *const u32);
        write_volatile((base + 0x14) as *mut u32, bar1);
        write_volatile((base + 0x20) as *mut u32, u32::MAX);
        write_volatile((base + 0x24) as *mut u32, u32::MAX);
        asm!("dsb sy", options(nostack, preserves_flags));
        mask4 = read_volatile((base + 0x20) as *const u32);
        mask5 = read_volatile((base + 0x24) as *const u32);
        write_volatile((base + 0x24) as *mut u32, bar5);
        write_volatile((base + 0x20) as *mut u32, bar4);
        write_volatile((base + 0x04) as *mut u16, command);
        asm!("dsb sy", options(nostack, preserves_flags));
    }
    if unsafe { read_volatile((base + 0x04) as *const u16) } != command
        || unsafe { read_volatile((base + 0x14) as *const u32) } != bar1
        || unsafe { read_volatile((base + 0x20) as *const u32) } != bar4
        || unsafe { read_volatile((base + 0x24) as *const u32) } != bar5 {
        return Err("k1445-bar-restore");
    }
    let first = decode_memory_bar_sizing(bar1, mask1, 0)
        .map_err(|_| "k1445-bar1-size")?;
    let second = decode_memory_bar_sizing(bar4, mask4, mask5)
        .map_err(|_| "k1445-bar4-size")?;
    if first.width != Width::Bits32 || second.width != Width::Bits64
        || first.size < 16 || second.size < 16 {
        return Err("k1445-bar-types");
    }
    let mut ranges = RangeTable::<2>::new();
    for (low, high, low_mask, high_mask) in
        [(bar1, 0, mask1, 0), (bar4, bar5, mask4, mask5)] {
        match ranges.claim(low, high, low_mask, high_mask) {
            Ok(range) => ranges.release(range).map_err(|_| "k1445-range-release")?,
            Err(BarError::Unassigned) => {},
            Err(_) => return Err("k1445-range-decode"),
        }
    }
    configs.release(leaf).map_err(|_| "k1445-leaf-release")?;
    if ranges.count() != 0 || configs.count() != 0 {
        return Err("k1445-claim-leak");
    }
    write_str("TURKOS/OS/K1445-BAR: BAR1-SIZE=0x");
    write_hex(first.size as usize);
    write_str(" BAR1-WIDTH=32 BAR4-SIZE=0x");
    write_hex(second.size as usize);
    write_line(" BAR4-WIDTH=64 RESTORED=1 CLAIMS=0");
    Ok(())
}