//! Block devices: things that store fixed-size *sectors* (512 bytes) you can
//! read and write by number (LBA = logical block address). Disks, partitions
//! of disks, and a RAM disk for tests all implement [`BlockDevice`], so file
//! systems don't care which one they sit on.
//!
//! Layers, bottom to top:
//!   virtio_blk (a QEMU disk)  ->  partition (a slice of a disk)  ->  fs::fat (FAT32)
//!
//! Devices are registered by name ("vda", "vda1", like Linux) in a global
//! list, see [`register`] and [`devices`].

pub mod partition;
pub mod virtio_blk;

use crate::errno::Errno;
use alloc::{string::String, sync::Arc, vec, vec::Vec};
use spin::Mutex;

pub const SECTOR_SIZE: usize = 512;

pub trait BlockDevice: Send + Sync {
    /// Number of sectors.
    fn sector_count(&self) -> u64;
    /// Read `buf.len() / 512` sectors starting at `lba` (buf.len() must be a
    /// multiple of 512).
    fn read_sectors(&self, lba: u64, buf: &mut [u8]) -> Result<(), Errno>;
    /// Write `buf.len() / 512` sectors starting at `lba`.
    fn write_sectors(&self, lba: u64, buf: &[u8]) -> Result<(), Errno>;
}

/// Bounds check shared by all devices: whole sectors, inside the device.
pub fn check_range(dev: &dyn BlockDevice, lba: u64, len: usize) -> Result<u64, Errno> {
    if len % SECTOR_SIZE != 0 {
        return Err(Errno::EINVAL);
    }
    let count = (len / SECTOR_SIZE) as u64;
    match lba.checked_add(count) {
        Some(end) if end <= dev.sector_count() => Ok(count),
        _ => Err(Errno::EIO), // like reading past the end of a real disk
    }
}

/// A disk in RAM (for tests, and handy for experiments).
pub struct RamDisk {
    data: Mutex<Vec<u8>>,
}

impl RamDisk {
    pub fn new(sectors: usize) -> Self {
        RamDisk { data: Mutex::new(vec![0; sectors * SECTOR_SIZE]) }
    }
}

impl BlockDevice for RamDisk {
    fn sector_count(&self) -> u64 {
        (self.data.lock().len() / SECTOR_SIZE) as u64
    }
    fn read_sectors(&self, lba: u64, buf: &mut [u8]) -> Result<(), Errno> {
        check_range(self, lba, buf.len())?;
        let start = lba as usize * SECTOR_SIZE;
        buf.copy_from_slice(&self.data.lock()[start..start + buf.len()]);
        Ok(())
    }
    fn write_sectors(&self, lba: u64, buf: &[u8]) -> Result<(), Errno> {
        check_range(self, lba, buf.len())?;
        let start = lba as usize * SECTOR_SIZE;
        self.data.lock()[start..start + buf.len()].copy_from_slice(buf);
        Ok(())
    }
}

static DEVICES: Mutex<Vec<(String, Arc<dyn BlockDevice>)>> = Mutex::new(Vec::new());

pub fn register(name: &str, dev: Arc<dyn BlockDevice>) {
    DEVICES.lock().push((String::from(name), dev));
}

/// All registered devices, (name, device), in registration order.
pub fn devices() -> Vec<(String, Arc<dyn BlockDevice>)> {
    DEVICES.lock().clone()
}

pub fn get(name: &str) -> Option<Arc<dyn BlockDevice>> {
    DEVICES.lock().iter().find(|(n, _)| n == name).map(|(_, d)| d.clone())
}

/// Boot-time storage setup: find virtio disks on the PCI bus ("vda", "vdb"...),
/// register them, and register their partitions ("vda1", ...). Returns
/// lines describing what was found, for printing.
pub fn init(pci_devices: &[crate::pci::PciDevice]) -> Vec<String> {
    use alloc::format;
    let mut log = Vec::new();
    let mut letter = b'a';
    for pci in pci_devices.iter().filter(|d| virtio_blk::is_virtio_blk(d)) {
        let name = format!("vd{}", letter as char);
        match virtio_blk::VirtioBlk::new(pci) {
            Ok(disk) => {
                let disk: Arc<dyn BlockDevice> = Arc::new(disk);
                log.push(format!(
                    "{}: virtio-blk, {} sectors ({} MiB)",
                    name,
                    disk.sector_count(),
                    disk.sector_count() * SECTOR_SIZE as u64 / (1024 * 1024)
                ));
                register(&name, disk.clone());
                match partition::scan(&disk) {
                    Ok((scheme, parts)) => {
                        log.push(format!("{}: {} partition table, {} partition(s)", name, scheme, parts.len()));
                        for (i, p) in parts.into_iter().enumerate() {
                            let pname = format!("{}{}", name, i + 1);
                            log.push(format!(
                                "{}: start {} sectors {} type {}{}",
                                pname, p.start, p.sectors, p.type_name,
                                if p.name.is_empty() { String::new() } else { format!(" name \"{}\"", p.name) }
                            ));
                            register(&pname, Arc::new(partition::Partition::new(disk.clone(), p.start, p.sectors)));
                        }
                    }
                    Err(e) => log.push(format!("{}: no partition table ({:?})", name, e)),
                }
            }
            Err(e) => log.push(format!("{}: virtio-blk init failed: {:?}", name, e)),
        }
        letter += 1;
    }
    log
}