// Copyright (C) 2024  Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use core::cell::RefCell;
use crc32fast::Hasher;
pub use gbl_storage::{
    alignment_scratch_size, is_aligned, is_buffer_aligned, required_scratch_size, AsBlockDevice,
    AsMultiBlockDevices, BlockDeviceEx, BlockInfo, BlockIo, BlockIoError, GptEntry, GptHeader,
    IoStatus, NonBlockingBlockIo, GPT_MAGIC, GPT_NAME_LEN_U16,
};
use safemath::SafeNum;
use std::collections::BTreeMap;
use zerocopy::AsBytes;

// Declares a per-thread global instance of timestamp. The timestamp is used to control the
// execution of non-blocking IO.
thread_local! {
    static TIMESTAMP: RefCell<u64> = RefCell::new(Default::default());
    /// Number of `TimestampPauser` in effect.
    static NUM_TIMESTAMP_PAUSER: RefCell<u64> = RefCell::new(Default::default());
}

/// Increases the value of timestamp.
pub fn advance_timestamp() {
    TIMESTAMP.with(|ts| (*ts.borrow_mut()) += 1);
}

/// Queries the current value of timestamp.
pub fn query_timestamp() -> u64 {
    NUM_TIMESTAMP_PAUSER.with(|v| (*v.borrow() == 0).then(|| advance_timestamp()));
    TIMESTAMP.with(|ts| *ts.borrow())
}

/// When a `TimestampPauser` is in scope, timestamp will not be increased by query.
pub struct TimestampPauser {}

impl Drop for TimestampPauser {
    fn drop(&mut self) {
        NUM_TIMESTAMP_PAUSER.with(|v| *v.borrow_mut() -= 1);
    }
}

impl TimestampPauser {
    /// Creates a new instance to pause the timestamp.
    pub fn new() -> Self {
        NUM_TIMESTAMP_PAUSER.with(|v| *v.borrow_mut() += 1);
        Self {}
    }

    /// Consumes the pauser, causing it to go out of scope. When all pausers go out of scope,
    /// timestamp resumes.
    pub fn resume(self) {}
}

/// `NonBlockingIoState` tracks the non-blocking IO state.
enum NonBlockingIoState {
    // (timestamp when initiated, blk offset, buffer, is read)
    Pending(u64, u64, &'static mut [u8], bool),
    Ready(IoStatus),
}

/// Helper `gbl_storage::BlockIo` struct for TestBlockDevice.
pub struct TestBlockIo {
    /// The storage block size in bytes.
    pub block_size: u64,
    /// The storage access alignment in bytes.
    pub alignment: u64,
    /// The backing storage data.
    pub storage: Vec<u8>,
    /// The number of successful write calls.
    pub num_writes: usize,
    /// The number of successful read calls.
    pub num_reads: usize,
    /// Pending non-blocking IO
    io: Option<NonBlockingIoState>,
}

impl TestBlockIo {
    pub fn new(block_size: u64, alignment: u64, data: Vec<u8>) -> Self {
        Self { block_size, alignment, storage: data, num_writes: 0, num_reads: 0, io: None }
    }

    fn check_alignment(&mut self, buffer: &[u8]) -> bool {
        matches!(is_buffer_aligned(buffer, self.alignment()), Ok(true))
            && matches!(is_aligned(buffer.len().into(), self.block_size().into()), Ok(true))
    }
}

impl BlockIo for TestBlockIo {
    fn info(&mut self) -> BlockInfo {
        NonBlockingBlockIo::info(self)
    }

    fn read_blocks(&mut self, blk_offset: u64, out: &mut [u8]) -> Result<(), BlockIoError> {
        // `BlockIo` is implemented for `&mut dyn NonBlockingBlockIo`
        BlockIo::read_blocks(&mut (self as &mut dyn NonBlockingBlockIo), blk_offset, out)
    }

    fn write_blocks(&mut self, blk_offset: u64, data: &mut [u8]) -> Result<(), BlockIoError> {
        BlockIo::write_blocks(&mut (self as &mut dyn NonBlockingBlockIo), blk_offset, data)
    }
}

// SAFETY:
// * When `TestBlockIo::io` is `Some(NonBlockingIoState(Pending(_, _, buffer, _)))`,
//   `check_status()` always returns `IoStatus::Pending`. `check_status()` returns other `IoStatus`
//   values if and only if `TestBlockIo::io` is not `Some(NonBlockingIoState(Pending())`, in which
//   case the buffer is not tracked anymore and thus will not be retained again.
// * `Self::check_status()` does not dereference the input pointer.
// * `TestBlockIo::io` is set to `Some(NonBlockingIoState(Pending(_, _, buffer, _)))` and retains
//   the buffer only on success (returning Ok(())).
unsafe impl NonBlockingBlockIo for TestBlockIo {
    /// Returns a `BlockInfo` for the block device.
    fn info(&mut self) -> BlockInfo {
        BlockInfo {
            block_size: self.block_size,
            num_blocks: u64::try_from(self.storage.len()).unwrap() / self.block_size,
            alignment: self.alignment,
        }
    }

    unsafe fn write_blocks(
        &mut self,
        blk_offset: u64,
        buffer: *mut [u8],
    ) -> core::result::Result<(), BlockIoError> {
        match self.io {
            Some(_) => Err(BlockIoError::MediaBusy),
            _ => {
                self.num_writes += 1;
                // SAFETY: By safety requirement, trait implementation can retain the buffer until
                // it no longer returns `IoStatus::Pending` in `Self::check_status()`.
                let buffer = unsafe { &mut *buffer };
                assert!(self.check_alignment(buffer));
                self.io =
                    Some(NonBlockingIoState::Pending(query_timestamp(), blk_offset, buffer, false));
                Ok(())
            }
        }
    }

    unsafe fn read_blocks(
        &mut self,
        blk_offset: u64,
        buffer: *mut [u8],
    ) -> core::result::Result<(), BlockIoError> {
        match self.io {
            Some(_) => Err(BlockIoError::MediaBusy),
            _ => {
                self.num_reads += 1;
                // SAFETY: By safety requirement, trait implementation can retain the buffer until
                // it no longer returns `IoStatus::Pending` in `Self::check_status()`.
                let buffer = unsafe { &mut *buffer };
                assert!(self.check_alignment(buffer));
                self.io =
                    Some(NonBlockingIoState::Pending(query_timestamp(), blk_offset, buffer, true));
                Ok(())
            }
        }
    }

    fn check_status(&mut self, buf: *mut [u8]) -> IoStatus {
        match self.io.as_mut() {
            Some(NonBlockingIoState::Pending(ts, blk_offset, ref mut buffer, is_read))
                if std::ptr::eq(*buffer as *const [u8], buf as _) =>
            {
                // Executes the IO if current timestamp is newer.
                if query_timestamp() > *ts {
                    let offset = (SafeNum::from(*blk_offset) * self.block_size).try_into().unwrap();
                    match is_read {
                        true => buffer.clone_from_slice(&self.storage[offset..][..buffer.len()]),
                        _ => self.storage[offset..][..buffer.len()].clone_from_slice(buffer),
                    }
                    self.io = Some(NonBlockingIoState::Ready(IoStatus::Completed));
                }
                IoStatus::Pending
            }
            Some(NonBlockingIoState::Ready(v)) => {
                let res = *v;
                self.io.take();
                res
            }
            _ => IoStatus::NotFound,
        }
    }

    fn abort(&mut self) -> core::result::Result<(), BlockIoError> {
        match self.io {
            Some(NonBlockingIoState::Pending(_, _, _, _)) => {
                self.io = Some(NonBlockingIoState::Ready(IoStatus::Aborted));
            }
            _ => {}
        }
        Ok(())
    }
}

/// Simple RAM based block device used by unit tests.
pub struct TestBlockDevice {
    /// The BlockIo helper struct.
    pub io: TestBlockIo,
    /// In-memory backing store.
    pub scratch: Vec<u8>,
    max_gpt_entries: u64,
}

impl TestBlockDevice {
    pub fn as_block_device_ex(&mut self) -> BlockDeviceEx {
        BlockDeviceEx::new((&mut self.io as &mut dyn NonBlockingBlockIo).into())
    }
}

impl From<&[u8]> for TestBlockDevice {
    fn from(data: &[u8]) -> Self {
        TestBlockDeviceBuilder::new().set_data(data).build()
    }
}

impl AsBlockDevice for TestBlockDevice {
    fn with(&mut self, f: &mut dyn FnMut(&mut dyn BlockIo, &mut [u8], u64)) {
        f(&mut self.io, &mut self.scratch[..], self.max_gpt_entries)
    }
}

impl Default for TestBlockDevice {
    fn default() -> Self {
        TestBlockDeviceBuilder::new().build()
    }
}

/// A description of the backing data store for a block device or partition.
/// Can either describe explicit data the device or partition is initialized with
/// OR a size in bytes if the device or partition can be initialized in a blank state.
#[derive(Copy, Clone)]
pub enum BackingStore<'a> {
    Data(&'a [u8]),
    Size(usize),
}

impl<'a> BackingStore<'a> {
    fn size(&self) -> usize {
        match self {
            Self::Data(slice) => slice.len(),
            Self::Size(size) => *size,
        }
    }
}

enum DiskDescription<'a> {
    Disk(BackingStore<'a>),
    Partitions(BTreeMap<&'static str, BackingStore<'a>>),
}

/// Builder struct for TestBlockDevice.
/// Most tests will want either:
/// 1) A blank device of a reasonable size OR
/// 2) A device with specific initial data.
/// Other customizations include block size,
/// the maximum number of GPT entries,
/// the alignment requirements,
/// and the size of the scratch buffer.
///
/// Note: setting the storage size or storage data is generally safe,
///       as long as the backing store is large enough,
///       but customizing other attributes may generate a block device
///       that cannot successfully complete any operations.
///       This may be exactly the intention, but be warned that it can be tricky
///       to customize the device and generate something that works without errors.
pub struct TestBlockDeviceBuilder<'a> {
    block_size: u64,
    max_gpt_entries: u64,
    alignment: u64,
    disk_description: DiskDescription<'a>,
    scratch_size: Option<usize>,
}

impl<'a> TestBlockDeviceBuilder<'a> {
    /// The default access alignment in bytes.
    pub const DEFAULT_ALIGNMENT: u64 = 64;
    /// The default block size in bytes.
    pub const DEFAULT_BLOCK_SIZE: u64 = 512;
    /// The default maximum number of GPT entries.
    pub const MAX_GPT_ENTRIES: u64 = 128;

    /// Creates a new TestBlockDeviceBuilder with defaults for all attributes.
    pub fn new() -> Self {
        Self {
            block_size: Self::DEFAULT_BLOCK_SIZE,
            max_gpt_entries: Self::MAX_GPT_ENTRIES,
            alignment: Self::DEFAULT_ALIGNMENT,
            disk_description: DiskDescription::Disk(BackingStore::Size(
                (Self::DEFAULT_BLOCK_SIZE * 32) as usize,
            )),
            scratch_size: None,
        }
    }

    /// Set the block size of the block device in bytes.
    /// The default is `DEFAULT_BLOCK_SIZE`.
    pub fn set_block_size(mut self, block_size: u64) -> Self {
        self.block_size = block_size;
        self
    }

    /// Set the maximum number of GPT entries for the GPT header.
    /// The default is `MAX_GPT_ENTRIES`.
    /// Note: setting too large a number of entries will make a device
    ///       that fails to sync its GPT.
    pub fn set_max_gpt_entries(mut self, max_gpt_entries: u64) -> Self {
        self.max_gpt_entries = max_gpt_entries;
        self
    }

    /// Set the required alignment for the TestBlockDevice.
    /// An alignment of `0` means there are no alignment requirements.
    /// The default is `DEFAULT_ALIGNMENT`.
    pub fn set_alignment(mut self, alignment: u64) -> Self {
        self.alignment = alignment;
        self
    }

    /// Set the size of TestBlockDevice in bytes.
    /// When built, the TestBlockDevice will have a blank backing store of size `size`.
    /// The default is `DEFAULT_BLOCK_SIZE` * 32.
    ///
    /// Note: This option is mutually exclusive with `set_data` and `add_partition`.
    ///       If `set_data` or `add_partition` have been called, `set_size` overrides
    ///       those customizations.
    pub fn set_size(mut self, size: usize) -> Self {
        self.disk_description = DiskDescription::Disk(BackingStore::Size(size));
        self
    }

    /// Sets the block device's backing data to the provided slice.
    ///
    /// Note: This option is mutually exclusive with `set_size` and `add_partition`.
    ///       If `set_size` or `add_partition` have been called, `set_data` overrides
    ///       those customizations.
    pub fn set_data(mut self, data: &'a [u8]) -> Self {
        self.disk_description = DiskDescription::Disk(BackingStore::Data(data));
        self
    }

    /// Adds a partition description.
    /// Partitions can be defined either with a specific backing store
    /// from a slice OR from a specific size in bytes.
    /// Partition sizes are rounded up to full blocks.
    /// If the same partition name is added multiple times,
    /// the last definition is used.
    ///
    /// Note: explicitly added partitions are mutually exclusive with
    ///       `set_size` and `set_data`.
    ///       If either have been called, `add_partition` overrides that customization.
    pub fn add_partition(mut self, name: &'static str, backing: BackingStore<'a>) -> Self {
        match self.disk_description {
            DiskDescription::Disk(_) => {
                let mut map = BTreeMap::new();
                map.insert(name, backing);
                self.disk_description = DiskDescription::Partitions(map);
            }
            DiskDescription::Partitions(ref mut map) => {
                map.insert(name, backing);
            }
        };
        self
    }

    /// Customize the size of the block device's scratch buffer.
    /// The default size is a known safe minimum calculated when `build()` is called.
    ///
    /// Note: Too small a scratch buffer will generate errors.
    ///       Unless a test is specifically interested in a non-default
    ///       scratch size, it's better to rely on the default size.
    pub fn set_scratch_size(mut self, scratch_size: usize) -> Self {
        self.scratch_size = Some(scratch_size);
        self
    }

    /// Consumes the builder and generates a TestBlockDevice
    /// with the desired customizations.
    pub fn build(self) -> TestBlockDevice {
        let storage = match self.disk_description {
            DiskDescription::Disk(BackingStore::Data(slice)) => Vec::from(slice),
            DiskDescription::Disk(BackingStore::Size(size)) => vec![0u8; size],
            DiskDescription::Partitions(partitions) => {
                partitions_to_disk_data(&partitions, self.block_size as usize)
            }
        };
        assert!(storage.len() % (self.block_size as usize) == 0);
        let mut io = TestBlockIo::new(self.block_size, self.alignment, storage);
        let scratch_size = match self.scratch_size {
            Some(s) => s,
            None => required_scratch_size(&mut io, self.max_gpt_entries).unwrap(),
        };
        TestBlockDevice {
            io,
            scratch: vec![0u8; scratch_size],
            max_gpt_entries: self.max_gpt_entries,
        }
    }
}

fn str_to_utf16_entry_name(name: &str) -> [u16; GPT_NAME_LEN_U16] {
    assert!(name.len() < GPT_NAME_LEN_U16);
    let mut data = [0; GPT_NAME_LEN_U16];
    let tmp: Vec<u16> = name.encode_utf16().collect();
    for (d, t) in std::iter::zip(data.iter_mut(), tmp) {
        *d = t;
    }
    data
}

fn pad_to_block_size(store: &mut Vec<u8>, block_size: usize) {
    let delta = (block_size - store.len() % block_size) % block_size;
    for _ in 0..delta {
        store.push(0);
    }
}

fn add_blocks(store: &mut Vec<u8>, data: &[u8], block_size: usize) {
    store.extend(data.iter());
    pad_to_block_size(store, block_size);
}

fn pad_bytes(store: &mut Vec<u8>, size: usize, block_size: usize) {
    for _ in 0..size {
        store.push(0);
    }
    pad_to_block_size(store, block_size);
}

fn partitions_to_disk_data(
    partitions: &BTreeMap<&'static str, BackingStore>,
    block_size: usize,
) -> Vec<u8> {
    let gpt_max_entries = 128;
    assert!(partitions.len() <= gpt_max_entries);
    let entry_blocks: u64 = ((SafeNum::from(partitions.len()) * std::mem::size_of::<GptEntry>())
        .round_up(block_size)
        / block_size)
        .try_into()
        .unwrap();
    let mut block = entry_blocks
        + 1  // Protective MBR
        + 1 // Primary GPT header
        ;
    // Leading mbr
    let mut store = vec![0; block_size];
    let mut header = GptHeader {
        magic: GPT_MAGIC,
        current: 1,
        size: std::mem::size_of::<GptHeader>() as u32,
        first: block,
        entries: 2,
        entries_count: std::cmp::min(partitions.len(), gpt_max_entries) as u32,
        entries_size: std::mem::size_of::<GptEntry>() as u32,
        ..Default::default()
    };

    // Define gpt entry structures
    let entries: Vec<GptEntry> = partitions
        .iter()
        .take(gpt_max_entries)
        .map(|(k, v)| {
            let last = (SafeNum::from(v.size()).round_up(block_size) / block_size + block - 1)
                .try_into()
                .unwrap();
            let mut entry = GptEntry {
                part_type: Default::default(),
                guid: Default::default(),
                first: block,
                last,
                flags: 0,
                name: str_to_utf16_entry_name(k),
            };
            entry.guid[0] = block as u8;
            block = last + 1;
            entry
        })
        .collect();

    // Patch last fields of header
    header.last = block - 1;
    header.backup = block + entry_blocks;
    header.entries_crc = entries
        .iter()
        .fold(Hasher::new(), |mut h, e| {
            h.update(e.as_bytes());
            h
        })
        .finalize();
    header.update_crc();

    // Primary header
    add_blocks(&mut store, header.as_bytes(), block_size);

    // Primary entries
    for e in &entries {
        store.extend(e.as_bytes());
    }
    pad_to_block_size(&mut store, block_size);

    // Partition store
    for p in partitions.values() {
        match p {
            BackingStore::Data(d) => add_blocks(&mut store, d, block_size),
            BackingStore::Size(s) => pad_bytes(&mut store, *s, block_size),
        };
    }

    // Backup entries
    let backup_entries_block = store.len() / block_size;
    for e in entries {
        store.extend(e.as_bytes());
    }
    pad_to_block_size(&mut store, block_size);
    // Tweak header to make it the backup.
    header.current = header.backup;
    header.backup = 1;
    header.entries = backup_entries_block.try_into().unwrap();
    header.update_crc();
    add_blocks(&mut store, header.as_bytes(), block_size);

    store
}

/// Simple RAM based multi-block device used for unit tests.
pub struct TestMultiBlockDevices(pub Vec<TestBlockDevice>);

impl AsMultiBlockDevices for TestMultiBlockDevices {
    fn for_each(
        &mut self,
        f: &mut dyn FnMut(&mut dyn AsBlockDevice, u64),
    ) -> core::result::Result<(), Option<&'static str>> {
        let _ = self
            .0
            .iter_mut()
            .enumerate()
            .for_each(|(idx, ele)| f(ele, u64::try_from(idx).unwrap()));
        Ok(())
    }
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn test_builder_partitions() {
        let data: [u8; 8] = [1, 2, 3, 4, 5, 6, 7, 8];
        let mut actual: [u8; 8] = Default::default();
        let mut block_dev = TestBlockDeviceBuilder::new()
            .add_partition("squid", BackingStore::Data(&data))
            .add_partition("clam", BackingStore::Size(28))
            .build();

        assert!(block_dev.sync_gpt().is_ok());
        assert!(block_dev.read_gpt_partition("squid", 0, actual.as_mut_slice()).is_ok());
        assert_eq!(actual, data);

        assert!(block_dev.read_gpt_partition("clam", 0, actual.as_mut_slice()).is_ok());
        assert_eq!(actual, [0u8; 8]);
    }
}