//! Block layer: PCI enumeration, the virtio-blk driver (against the
//! target/disk.img QEMU attaches), partition table parsing, RAM disks.
#![no_std]
#![no_main]
#![feature(custom_test_frameworks)]
#![test_runner(rust_kernel::test_runner)]
#![reexport_test_harness_main = "test_main"]

extern crate alloc;

use alloc::{sync::Arc, vec};
use bootloader::{BootInfo, entry_point};
use core::panic::PanicInfo;
use rust_kernel::block::{self, BlockDevice, RamDisk, SECTOR_SIZE, partition};
use rust_kernel::errno::Errno;
use rust_kernel::pci;

entry_point!(main);

fn main(boot_info: &'static BootInfo) -> ! {
    rust_kernel::init();
    rust_kernel::init_memory(boot_info);
    block::init(&pci::scan());
    test_main();
    rust_kernel::hlt_loop();
}

#[test_case]
fn pci_scan_finds_the_virtio_disk() {
    let devices = pci::scan();
    // QEMU's i440FX machine always has the host bridge at 00:00.0.
    assert_eq!((devices[0].bus, devices[0].device, devices[0].vendor_id), (0, 0, 0x8086));
    assert!(devices.iter().any(|d| d.vendor_id == 0x1af4 && d.device_id == 0x1001));
}

#[test_case]
fn virtio_disk_reads_the_gpt_image() {
    let disk = block::get("vda").expect("no vda: is target/disk.img attached?");
    assert_eq!(disk.sector_count(), 64 * 1024 * 1024 / 512);
    let mut sector = [0u8; SECTOR_SIZE];
    disk.read_sectors(0, &mut sector).unwrap();
    assert_eq!(&sector[510..], &[0x55, 0xAA]); // protective MBR signature
    assert_eq!(sector[446 + 4], 0xEE); // partition type "GPT protective"
    disk.read_sectors(1, &mut sector).unwrap();
    assert_eq!(&sector[..8], b"EFI PART");
}

#[test_case]
fn multi_sector_reads_match_single_sector_reads() {
    // 20 sectors = more than one 4 KiB bounce-buffer page per request.
    let disk = block::get("vda").unwrap();
    let mut big = vec![0u8; 20 * SECTOR_SIZE];
    disk.read_sectors(2048, &mut big).unwrap();
    let mut one = [0u8; SECTOR_SIZE];
    for i in 0..20 {
        disk.read_sectors(2048 + i as u64, &mut one).unwrap();
        assert_eq!(&big[i * SECTOR_SIZE..(i + 1) * SECTOR_SIZE], &one[..]);
    }
}

#[test_case]
fn virtio_write_then_read_back() {
    // QEMU runs the disk with snapshot=on, so this doesn't touch the file.
    let disk = block::get("vda").unwrap();
    let data: alloc::vec::Vec<u8> = (0..2 * SECTOR_SIZE).map(|i| (i % 251) as u8).collect();
    disk.write_sectors(100, &data).unwrap();
    let mut back = vec![0u8; 2 * SECTOR_SIZE];
    disk.read_sectors(100, &mut back).unwrap();
    assert_eq!(back, data);
}

#[test_case]
fn out_of_range_and_partial_sectors_fail() {
    let disk = block::get("vda").unwrap();
    let mut sector = [0u8; SECTOR_SIZE];
    assert_eq!(disk.read_sectors(disk.sector_count(), &mut sector), Err(Errno::EIO));
    assert_eq!(disk.read_sectors(0, &mut sector[..100]), Err(Errno::EINVAL));
}

#[test_case]
fn gpt_partition_is_found_and_registered() {
    let disk = block::get("vda").unwrap();
    let (scheme, parts) = partition::scan(&disk).unwrap();
    assert_eq!(scheme, "GPT");
    assert_eq!(parts.len(), 1);
    assert_eq!((parts[0].start, parts[0].sectors), (2048, 126976));
    assert_eq!(parts[0].name, "fat32");
    // vda1 is a block device whose sector 0 is the disk's sector 2048:
    // the FAT boot sector (ends in 55 AA, says "FAT32   " at offset 82).
    let vda1 = block::get("vda1").unwrap();
    assert_eq!(vda1.sector_count(), 126976);
    let mut sector = [0u8; SECTOR_SIZE];
    vda1.read_sectors(0, &mut sector).unwrap();
    assert_eq!(&sector[82..90], b"FAT32   ");
    assert_eq!(vda1.read_sectors(126976, &mut sector), Err(Errno::EIO)); // past its end
}

#[test_case]
fn mbr_partitions_on_a_ram_disk() {
    // Build a classic MBR by hand: two primary partitions.
    let disk: Arc<dyn BlockDevice> = Arc::new(RamDisk::new(1000));
    let mut mbr = [0u8; SECTOR_SIZE];
    let entry = |mbr: &mut [u8; SECTOR_SIZE], i: usize, kind: u8, start: u32, len: u32| {
        let e = 446 + 16 * i;
        mbr[e + 4] = kind;
        mbr[e + 8..e + 12].copy_from_slice(&start.to_le_bytes());
        mbr[e + 12..e + 16].copy_from_slice(&len.to_le_bytes());
    };
    entry(&mut mbr, 0, 0x0c, 63, 400); // FAT32 (LBA)
    entry(&mut mbr, 1, 0x83, 500, 300); // Linux
    // No signature yet -> not a partition table.
    assert!(matches!(partition::scan(&disk), Err(partition::PartitionError::NoSignature)));
    mbr[510] = 0x55;
    mbr[511] = 0xAA;
    disk.write_sectors(0, &mbr).unwrap();
    let (scheme, parts) = partition::scan(&disk).unwrap();
    assert_eq!(scheme, "MBR");
    assert_eq!(parts.len(), 2);
    assert_eq!((parts[0].start, parts[0].sectors), (63, 400));
    assert_eq!((parts[1].start, parts[1].sectors), (500, 300));

    // A partition device maps its sector 0 to the partition's start.
    let part = partition::Partition::new(disk.clone(), 500, 300);
    part.write_sectors(0, &[7u8; SECTOR_SIZE]).unwrap();
    let mut back = [0u8; SECTOR_SIZE];
    disk.read_sectors(500, &mut back).unwrap();
    assert_eq!(back, [7u8; SECTOR_SIZE]);
}

#[panic_handler]
fn panic(info: &PanicInfo) -> ! {
    rust_kernel::test_panic_handler(info)
}