// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

//! VMGS storage implementation on top of [`SimpleDisk`].

use disk_backend::DiskError;
use disk_backend::SimpleDisk;
use guestmem::GuestMemory;
use scsi_buffers::OwnedRequestBuffers;
use std::sync::Arc;
use thiserror::Error;

#[cfg_attr(feature = "inspect", derive(inspect::Inspect))]
pub(crate) struct VmgsStorage {
    disk: Arc<dyn SimpleDisk>,
    mem: GuestMemory,
    mem_size: usize,
    sector_size: u32,
    sector_shift: u32,
}

/// Error due to underlying VMGS storage.
#[derive(Debug, Error)]
pub enum StorageError {
    /// A disk IO error.
    #[error(transparent)]
    Disk(DiskError),
    /// The access was not aligned to a sector boundary.
    #[error("access is not aligned to a sector boundary")]
    Unaligned,
}

impl VmgsStorage {
    pub fn new(disk: Arc<dyn SimpleDisk>) -> Self {
        let sector_size = disk.sector_size();
        // Max IO size. Balance between performance and memory usage.
        let mem_size = 64 * 1024;
        Self {
            mem: GuestMemory::allocate(mem_size),
            mem_size,
            sector_size,
            sector_shift: sector_size.trailing_zeros(),
            disk,
        }
    }

    /// Read from the block device.
    ///
    /// The beginning of the read must be sector aligned, but the end need not
    /// be.
    pub async fn read_block(
        &mut self,
        mut byte_offset: u64,
        buf: &mut [u8],
    ) -> Result<(), StorageError> {
        if byte_offset & (self.sector_size as u64 - 1) != 0 {
            return Err(StorageError::Unaligned);
        }
        for buf in buf.chunks_mut(self.mem_size) {
            let sector = byte_offset >> self.sector_shift;
            let sector_end = (byte_offset + buf.len() as u64 + (self.sector_size as u64 - 1))
                >> self.sector_shift;

            let len = (sector_end - sector) << self.sector_shift;
            let buffers = OwnedRequestBuffers::linear(0, len as usize, true);
            self.disk
                .read_vectored(&buffers.buffer(&self.mem), sector)
                .await
                .map_err(StorageError::Disk)?;

            self.mem.read_at(0, buf).unwrap();
            byte_offset += len;
        }
        Ok(())
    }

    /// Write a block to the block device.
    ///
    /// The beginning of the write must be sector aligned, but the end need not
    /// be.
    ///
    /// CAUTION: if the end of the write is not sector aligned, the trailing sector
    /// data will be zeroed out.
    pub async fn write_block(
        &mut self,
        mut byte_offset: u64,
        buf: &[u8],
    ) -> Result<(), StorageError> {
        if byte_offset & (self.sector_size as u64 - 1) != 0 {
            return Err(StorageError::Unaligned);
        }
        for buf in buf.chunks(self.mem_size) {
            let sector = byte_offset >> self.sector_shift;
            let sector_end = (byte_offset + buf.len() as u64 + (self.sector_size as u64 - 1))
                >> self.sector_shift;
            let len = (sector_end - sector) << self.sector_shift;

            self.mem.write_at(0, buf).unwrap();
            // Zero the trailing sector data.
            self.mem
                .fill_at(buf.len() as u64, 0, len as usize - buf.len())
                .unwrap();

            let buffers = OwnedRequestBuffers::linear(0, len as usize, false);
            self.disk
                .write_vectored(&buffers.buffer(&self.mem), sector, false)
                .await
                .map_err(StorageError::Disk)?;
            byte_offset += len;
        }
        Ok(())
    }

    /// Flush any buffered data.
    pub async fn flush(&mut self) -> Result<(), StorageError> {
        self.disk.sync_cache().await.map_err(StorageError::Disk)
    }

    pub fn sector_size(&self) -> u32 {
        self.sector_size
    }

    pub fn sector_count(&self) -> u64 {
        self.disk.sector_count()
    }

    /// Capacity in VMGS blocks.
    pub fn block_capacity(&self) -> u32 {
        ((self.sector_count() * self.sector_size() as u64)
            .min(vmgs_format::VMGS_MAX_CAPACITY_BYTES)
            / vmgs_format::VMGS_BYTES_PER_BLOCK as u64) as u32
    }

    pub fn aligned_header_size(&self) -> u64 {
        assert!(self.sector_size() >= size_of::<vmgs_format::VmgsHeader>() as u32);
        self.sector_size().into()
    }
}
