// Copyright 2022 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
//
//     https://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 crate::bit_vector::*;
use crate::block_device::*;
use crate::errors::*;
use crate::utils::*;

pub struct TestDevice<'a> {
    pub geom: BlockDeviceGeometry,
    pub flash_base: *mut u8,
    pub dirty: BitVector<'a>,
}

impl<'a> TestDevice<'a> {
    pub fn new(geom: BlockDeviceGeometry, flash_base: *mut u8, dirty_bits: &'a mut [u32]) -> Self {
        assert_eq!(geom.block_size * geom.block_count, dirty_bits.len() * 32);
        let result = TestDevice {
            geom: geom,
            flash_base: flash_base,
            dirty: BitVector::new(dirty_bits),
        };
        return result;
    }

    pub fn dirty_count(&self, addr: usize, size: usize) -> Result<usize, BFSErr> {
        let mut result: usize = 0;
        for i in addr..addr + size {
            result += self.dirty.get_bit(i)? as usize;
        }
        return Ok(result);
    }

    pub fn mark_dirty(&mut self, addr: usize, size: usize) -> Result<(), BFSErr> {
        self.dirty.set_range(addr, addr + size)?;
        return Ok(());
    }

    pub fn mark_clean(&mut self, addr: usize, size: usize) -> Result<(), BFSErr> {
        self.dirty.clear_range(addr, addr + size)?;
        return Ok(());
    }
}

/// Preconditions for TestDevice.
/// FIXME - I loosened these up because they were annoying, not sure how strict
/// we should be...

impl<'a> TestDevice<'a> {
    pub fn check_dirty(&self, addr: usize, size: usize) -> Result<(), BFSErr> {
        for i in addr..addr + size {
            dcheck!(self.dirty.get_bit(i)? == 1, BFSErr::CleanRead);
        }
        return Ok(());
    }

    pub fn check_clean(&self, addr: usize, size: usize) -> Result<(), BFSErr> {
        for i in addr..addr + size {
            dcheck!(self.dirty.get_bit(i)? == 0, BFSErr::DirtyWrite);
        }
        return Ok(());
    }

    pub fn check_is_block_sized(&self, data: &[u8]) -> Result<(), BFSErr> {
        dcheck!(data.len() == self.geom.block_size, BFSErr::OutOfBounds);
        return Ok(());
    }

    pub fn check_block_index(&self, iblock: usize) -> Result<(), BFSErr> {
        dcheck!(iblock < self.geom.block_count, BFSErr::OutOfBounds);
        return Ok(());
    }

    pub fn check_fits_in_block(&self, addr: usize, data: &[u8]) -> Result<(), BFSErr> {
        let block_a = addr / self.geom.block_size;
        let block_b = (addr + data.len() - 1) / self.geom.block_size;
        dcheck!(block_a == block_b, BFSErr::OutOfBounds);
        return Ok(());
    }

    pub fn check_read_block(&self, iblock: usize, block: &[u8]) -> Result<(), BFSErr> {
        self.check_block_index(iblock)?;
        self.check_fits_in_block(iblock * self.geom.block_size, block)?;
        //self.check_dirty(iblock * self.geom.block_size, block.len())?;
        return Ok(());
    }

    pub fn check_write_block(&self, iblock: usize, block: &[u8]) -> Result<(), BFSErr> {
        self.check_block_index(iblock)?;
        self.check_fits_in_block(iblock * self.geom.block_size, block)?;
        self.check_clean(iblock * self.geom.block_size, block.len())?;
        return Ok(());
    }

    pub fn check_erase_block(&self, iblock: usize) -> Result<(), BFSErr> {
        self.check_block_index(iblock)?;
        return Ok(());
    }

    pub fn check_erase_dirty_block(&self, iblock: usize) -> Result<(), BFSErr> {
        self.check_erase_block(iblock)?;
        let dirty_count = self.dirty_count(iblock * self.geom.block_size, self.geom.block_size)?;
        dcheck!(dirty_count > 0, BFSErr::BadErase);
        return Ok(());
    }

    pub fn check_range(&self, addr: usize, size: usize) -> Result<(), BFSErr> {
        let bs = self.geom.block_size;
        let bc = self.geom.block_count;
        dcheck!(addr + size < (bc * bs), BFSErr::OutOfBounds);
        return Ok(());
    }

    pub fn check_read_range(&self, addr: usize, data: &mut [u8]) -> Result<(), BFSErr> {
        self.check_range(addr, data.len())?;
        //self.check_dirty(addr, data.len())?;
        return Ok(());
    }

    pub fn check_write_range(&self, addr: usize, data: &[u8]) -> Result<(), BFSErr> {
        self.check_range(addr, data.len())?;
        self.check_clean(addr, data.len())?;
        return Ok(());
    }

    /// Check that this overwrite only changes bits from 1->0
    pub fn check_overwrite_range(&self, addr: usize, data: &[u8]) -> Result<(), BFSErr> {
        self.check_range(addr, data.len())?;
        for i in 0..data.len() {
            let src = data[i];
            unsafe {
                let dst = self.flash_base.add(addr + i).read();
                dcheck!((src & !dst) == 0, BFSErr::BadOverwrite);
            }
        }
        return Ok(());
    }
}

impl<'a> BlockDevice for TestDevice<'a> {
    fn geom(&self) -> BlockDeviceGeometry {
        self.geom
    }

    /// Read a chunk of flash contained in a single block.
    fn read_block(&self, iblock: usize, block: &mut [u8]) -> Result<(), BFSErr> {
        self.check_read_block(iblock, block)?;

        let bs = self.geom.block_size;
        unsafe {
            self.flash_base
                .add(iblock * bs)
                .copy_to(block.as_mut_ptr(), block.len());
        }

        return Ok(());
    }

    /// Write a a chunk of flash contained in a single block.
    fn write_block(&mut self, iblock: usize, block: &[u8]) -> Result<(), BFSErr> {
        self.check_write_block(iblock, block)?;

        let bs = self.geom.block_size;
        unsafe {
            self.flash_base
                .add(iblock * bs)
                .copy_from(block.as_ptr(), bs);
        }

        self.mark_dirty(iblock * bs, bs)?;
        return Ok(());
    }

    /// Erase an entire block of flash.
    fn erase_block(&mut self, iblock: usize) -> Result<(), BFSErr> {
        self.check_erase_block(iblock)?;

        let bs = self.geom.block_size;
        self.mark_clean(iblock * bs, bs)?;
        unsafe {
            self.flash_base.add(iblock * bs).write_bytes(0xFF, bs);
        }

        return Ok(());
    }

    /// Read a range of bytes in flash, checking first that those bytes have
    /// been written since they were erased.
    fn read_range(&self, addr: usize, data: &mut [u8]) -> Result<(), BFSErr> {
        self.check_read_range(addr, data)?;

        unsafe {
            self.flash_base
                .add(addr)
                .copy_to(data.as_mut_ptr(), data.len());
        }

        return Ok(());
    }

    /// Write to a range of bytes in flash, checking first that those bytes have
    /// not been written since they were erased.
    fn write_range(&mut self, addr: usize, data: &[u8]) -> Result<(), BFSErr> {
        self.check_write_range(addr, data)?;

        unsafe {
            self.flash_base
                .add(addr)
                .copy_from(data.as_ptr(), data.len());
        }

        self.mark_dirty(addr, data.len())?;
        return Ok(());
    }

    /// Write to a range of bytes in flash, but _do_ allow writing over dirty
    /// bytes - the result in flash will be the logical AND of the old and new
    /// data.
    fn overwrite_range(&mut self, addr: usize, data: &[u8]) -> Result<(), BFSErr> {
        self.check_overwrite_range(addr, data)?;

        unsafe {
            for i in 0..data.len() {
                let ptr = self.flash_base.add(addr + i);
                ptr.write(ptr.read() & data[i]);
            }
        }

        self.mark_dirty(addr, data.len())?;
        return Ok(());
    }
}

/// Overwriting existing data should write the logical AND of the old and new
/// data.

#[test]
fn test_overwrite() {
    const BLOCK_SIZE: usize = 16;
    const BLOCK_COUNT: usize = 4;
    let geom = BlockDeviceGeometry {
        block_size: BLOCK_SIZE,
        block_count: BLOCK_COUNT,
    };
    let mut buf: [u8; BLOCK_SIZE * BLOCK_COUNT] = [0; BLOCK_SIZE * BLOCK_COUNT];
    let mut dirty_bits = [0; BLOCK_SIZE * BLOCK_COUNT / 32];
    let bd: &mut dyn BlockDevice = &mut TestDevice::new(geom, buf.as_mut_ptr(), &mut dirty_bits);

    let write1: u32 = 0xFF0FF0FF;
    let write2: u32 = 0x0F00F000;

    assert_ok!(bd.write_range(0, &write1.to_le_bytes()));
    assert_ok!(bd.overwrite_range(0, &write2.to_le_bytes()));

    let mut result: u32 = 0;
    assert_ok!(bd.read_range(0, as_unsafe_blob_mut(&mut result)));
    assert_eq!(result, 0xFF0FF0FF & 0x0FF0FF00);
}

/// Trying to set bits that have already been cleared by a previous write should
/// cause a panic in the test device.

#[test]
#[should_panic]
fn test_bad_overwrite() {
    const BLOCK_SIZE: usize = 16;
    const BLOCK_COUNT: usize = 4;
    let geom = BlockDeviceGeometry {
        block_size: BLOCK_SIZE,
        block_count: BLOCK_COUNT,
    };
    let mut buf: [u8; BLOCK_SIZE * BLOCK_COUNT] = [0; BLOCK_SIZE * BLOCK_COUNT];
    let mut dirty_bits = [0; BLOCK_SIZE * BLOCK_COUNT / 32];
    let bd: &mut dyn BlockDevice = &mut TestDevice::new(geom, buf.as_mut_ptr(), &mut dirty_bits);

    let all_1: u32 = 0xFFFFFFFF;
    let all_0: u32 = 0x00000000;

    assert_ok!(bd.write_range(0, &all_0.to_le_bytes()));
    assert_err!(bd.overwrite_range(0, &all_1.to_le_bytes()));
}

/// Trying to read blocks that have never been written should cause a panic in
/// the test device.

#[test]
#[should_panic]
fn test_read_unwritten_block() {
    const BLOCK_SIZE: usize = 16;
    const BLOCK_COUNT: usize = 4;
    let geom = BlockDeviceGeometry {
        block_size: BLOCK_SIZE,
        block_count: BLOCK_COUNT,
    };
    let mut buf: [u8; BLOCK_SIZE * BLOCK_COUNT] = [0; BLOCK_SIZE * BLOCK_COUNT];
    let mut dirty_bits = [0; BLOCK_SIZE * BLOCK_COUNT / 32];
    let bd: &mut dyn BlockDevice = &mut TestDevice::new(geom, buf.as_mut_ptr(), &mut dirty_bits);

    let mut dst_block = [0; BLOCK_SIZE];
    assert_err!(bd.read_block(0, &mut dst_block));
}

/// Trying to read blocks that have been written and then erased should cause a
/// panic in the test device.

#[test]
#[should_panic]
fn test_read_erased_block() {
    const BLOCK_SIZE: usize = 16;
    const BLOCK_COUNT: usize = 4;
    let geom = BlockDeviceGeometry {
        block_size: BLOCK_SIZE,
        block_count: BLOCK_COUNT,
    };
    let mut buf: [u8; BLOCK_SIZE * BLOCK_COUNT] = [0; BLOCK_SIZE * BLOCK_COUNT];
    let mut dirty_bits = [0; BLOCK_SIZE * BLOCK_COUNT / 32];
    let bd: &mut dyn BlockDevice = &mut TestDevice::new(geom, buf.as_mut_ptr(), &mut dirty_bits);

    let mut src_block = [0; BLOCK_SIZE];
    let mut dst_block = [0; BLOCK_SIZE];

    for i in 0..geom.block_size {
        src_block[i] = i as u8;
    }

    assert_ok!(bd.write_block(0, &mut src_block));
    assert_ok!(bd.erase_block(0));
    assert_err!(bd.read_block(0, &mut dst_block));
}

/// All blocks in a device should be readable and writable.

#[test]
fn test_read_write() {
    const BLOCK_SIZE: usize = 16;
    const BLOCK_COUNT: usize = 4;
    let geom = BlockDeviceGeometry {
        block_size: BLOCK_SIZE,
        block_count: BLOCK_COUNT,
    };
    let mut buf: [u8; BLOCK_SIZE * BLOCK_COUNT] = [0; BLOCK_SIZE * BLOCK_COUNT];
    let mut dirty_bits = [0; BLOCK_SIZE * BLOCK_COUNT / 32];
    let bd: &mut dyn BlockDevice = &mut TestDevice::new(geom, buf.as_mut_ptr(), &mut dirty_bits);

    let mut src_block = [0; BLOCK_SIZE];
    let mut dst_block = [0; BLOCK_SIZE];

    for i in 0..geom.block_size {
        src_block[i] = i as u8;
    }

    // Writing clean blocks should succeed.
    for i in 0..geom.block_count {
        assert_ok!(bd.write_block(i, &src_block));
    }

    // Reading those blocks back should succeed, and the contents should match
    // the original block.
    for i in 0..geom.block_count {
        assert_ok!(bd.read_block(i, &mut dst_block));
        for j in 0..geom.block_size {
            assert_eq!(src_block[j], dst_block[j]);
        }
    }

    // Erasing blocks should succeed.
    for i in 0..geom.block_count {
        assert_ok!(bd.erase_block(i));
    }
}