/* buffered.rs
 *
 * Developed by Tim Walls <tim.walls@snowgoons.com>
 * Copyright (c) All Rights Reserved, Tim Walls
 */
//! A buffered storage driver layer implementation.

// Imports ===================================================================
use avr_oxide::alloc::boxed::Box;
use core::cell::UnsafeCell;
use avr_oxide::io::IoError;
use avrox_storage::{RandomRead, RandomWrite, SAddr, SSize};
use avrox_storage::random::{AccessHint, Storage};
use avr_oxide::OxideResult::Ok;

// Declarations ==============================================================
pub struct PageBuffer<const PAGE_SIZE: usize, SD> {
  data: UnsafeCell<PageBufferData<PAGE_SIZE,SD>>
}

struct PageBufferData<const PAGE_SIZE: usize, SD> {
  storage_driver: SD,
  page_buffer: Box<[u8; PAGE_SIZE]>,
  page_state: PageState
}

enum PageState {
  Invalid,
  Clean(SAddr),
  Dirty(SAddr)
}

// Code ======================================================================
impl<const PAGE_SIZE: usize, SD> PageBuffer<PAGE_SIZE,SD>
where
  SD: RandomRead + RandomWrite
{
  pub fn with_driver(driver: SD) -> Self {
    PageBuffer {
      data: UnsafeCell::new(PageBufferData {
        storage_driver: driver,
        page_buffer: Box::new([0x00; PAGE_SIZE]),
        page_state: PageState::Invalid
      })
    }
  }

  /// Calculate the page boundares (start address, offset and number of bytes
  /// remaining in page) for the given address
  fn page_boundaries_for_addr(&self, addr: SAddr) -> (SAddr,usize,usize) {
    let start_addr = ( addr / PAGE_SIZE as SAddr) * PAGE_SIZE as SAddr;
    let offset      = (addr - start_addr) as usize;
    let remaining   = PAGE_SIZE - offset;

    (start_addr, offset, remaining)
  }

  /// Is the given address within the current page
  fn is_addr_in_current_page(&self, addr: SAddr) -> bool {
    let data = unsafe { &mut *self.data.get() };

    match data.page_state {
      PageState::Invalid => false,
      PageState::Clean(my_pageaddr) |
      PageState::Dirty(my_pageaddr) => {
        let ( addr_pageaddr, _, _ ) = self.page_boundaries_for_addr(addr);
        addr_pageaddr == my_pageaddr
      }
    }
  }

  /// Load the given page.  If the current page is dirty, will write it out
  /// first if we need to change page.
  fn load_page_containing_addr(&self, addr: SAddr) -> avr_oxide::OxideResult<(),IoError>
  where
    SD: RandomRead + RandomWrite {
    let ( page_start, _offset, _remaining) = self.page_boundaries_for_addr(addr);
    let data = unsafe { &mut *self.data.get() };


    match data.page_state {
      PageState::Invalid => {
        data.storage_driver.read_full_at(page_start, &mut *data.page_buffer)?;
        data.page_state = PageState::Clean(page_start);
        Ok(())
      },
      PageState::Clean(current_page_addr) => {
        if current_page_addr != page_start {
          data.storage_driver.read_full_at(page_start, &mut *data.page_buffer)?;
          data.page_state = PageState::Clean(page_start);
        }
        Ok(())
      },
      PageState::Dirty(current_page_addr) => {
        if current_page_addr != page_start {
          data.storage_driver.write_all_at(current_page_addr, & *data.page_buffer)?;
          data.storage_driver.read_full_at(page_start, &mut *data.page_buffer)?;
          data.page_state = PageState::Clean(page_start);
        }
        Ok(())
      }
    }
  }
}

impl<const PAGE_SIZE: usize, SD> RandomRead for PageBuffer<PAGE_SIZE,SD>
where
  SD: RandomRead + RandomWrite
{
  fn read_at_hinted(&self, addr: SAddr, buf: &mut [u8], hint: AccessHint) -> avr_oxide::io::Result<usize> {
    if self.is_addr_in_current_page(addr) || !hint.is_nonsequential() {
      self.load_page_containing_addr(addr)?;
      let data = unsafe { &mut *self.data.get() };

      let (_,page_offset,remaining_in_page) = self.page_boundaries_for_addr(addr);

      if buf.len() >= remaining_in_page {
        buf[..remaining_in_page].copy_from_slice(&data.page_buffer[page_offset..(page_offset+remaining_in_page)]);
        Ok(remaining_in_page)
      } else {
        buf.copy_from_slice(&data.page_buffer[page_offset..(page_offset+buf.len())]);
        Ok(buf.len())
      }
    } else {
      // Non-sequential reads won't benefit from our page cache, so let's
      // just skip it (and thereby avoid invalidating whatever IS cached)
      let data = unsafe { &mut *self.data.get() };

      data.storage_driver.read_at_hinted(addr, buf, hint)
    }
  }
}

impl<const PAGE_SIZE: usize, SD> RandomWrite for PageBuffer<PAGE_SIZE,SD>
  where
    SD: RandomRead + RandomWrite
{
  fn write_at_hinted(&mut self, addr: SAddr, buf: &[u8], hint: AccessHint) -> avr_oxide::io::Result<usize> {
    if (hint.is_writeonly() || hint.is_nonsequential()) && !self.is_addr_in_current_page(addr) {
      // The start address is not in our current page, but we have also been
      // told it is write-only; in that case, we can write the data direct
      // to the underlying store without invalidating our cached page.
      //
      // We should still only write a page-worth max though (maybe the write is
      // big enough later bytes *do* overlap our current page)
      let data = unsafe { &mut *self.data.get() };
      let (_, _, remaining_in_page) = self.page_boundaries_for_addr(addr);

      if buf.len() >= remaining_in_page {
        data.storage_driver.write_at_hinted(addr, &buf[..remaining_in_page], hint)
      } else {
        data.storage_driver.write_at_hinted(addr, buf, hint)
      }

    } else {
      self.load_page_containing_addr(addr)?;
      let data = unsafe { &mut *self.data.get() };

      let (page_start,page_offset,remaining_in_page) = self.page_boundaries_for_addr(addr);

      if buf.len() >= remaining_in_page {
        data.page_buffer[page_offset..(page_offset+remaining_in_page)].copy_from_slice(&buf[..remaining_in_page]);
        data.page_state = PageState::Dirty(page_start);
        Ok(remaining_in_page)
      } else {
        data.page_buffer[page_offset..(page_offset+buf.len())].copy_from_slice(&buf);
        data.page_state = PageState::Dirty(page_start);
        Ok(buf.len())
      }
    }
  }

  fn flush(&mut self) -> avr_oxide::io::Result<()> {
    let data = unsafe { &mut *self.data.get() };

    match data.page_state {
      PageState::Invalid => {},
      PageState::Clean(_) => {},
      PageState::Dirty(current_page_addr) => {
        data.storage_driver.write_all_at(current_page_addr, & *data.page_buffer)?;
        data.page_state = PageState::Clean(current_page_addr);
      }
    }

    data.storage_driver.flush()
  }
}

impl<const PAGE_SIZE: usize, SD> Storage for PageBuffer<PAGE_SIZE,SD>
  where
    SD: RandomRead + RandomWrite + Storage
{
  const ADDRESSABLE_BYTES: SSize = SD::ADDRESSABLE_BYTES;
}


// Tests =====================================================================
#[cfg(test)]
mod tests {
  use avrox_storage::buffered::PageBuffer;
  use avrox_storage::serprom::generic::dummy::DummyPromBusClient;
  use avrox_storage::serprom::composite::tests::TestCompositeProm;
  use avr_oxide::devices::serialbus::UsesSerialBusClient;
  use avrox_storage::{RandomRead,RandomWrite};

  type TestBuffer = PageBuffer<32,TestCompositeProm>;

  #[test]
  pub fn test_page_boundary_maths() {
    let mut test_buffer = TestBuffer::with_driver(TestCompositeProm::using_client(DummyPromBusClient::new()));

    assert_eq!(test_buffer.page_boundaries_for_addr(0), (0, 0, 32));
    assert_eq!(test_buffer.page_boundaries_for_addr(1), (0, 1, 31));
    assert_eq!(test_buffer.page_boundaries_for_addr(31), (0, 31, 1));
    assert_eq!(test_buffer.page_boundaries_for_addr(32), (32, 0, 32));
  }

  #[test]
  pub fn test_buffered_prom_operations() {
    // We copy the same tests we used in the basic PROM; it should all give
    // the same results, just with extra cacheing under the bonnet.

    let mut test_prom = TestBuffer::with_driver(TestCompositeProm::using_client(DummyPromBusClient::new()));

    let mut buffer = [ 0x00u8, 0x00u8, 0x00u8, 0x00u8 ];

    // Check basic write/read
    test_prom.write_all_at(0x0000u32, &[ 0x01u8, 0x02u8, 0x03u8, 0x04u8 ]).unwrap();
    test_prom.read_full_at(0x0000u32, &mut buffer).unwrap();
    assert_eq!(buffer, [ 0x01u8, 0x02u8, 0x03u8, 0x04u8 ]);

    // Check basic write/read at a different address
    test_prom.write_all_at(0x0123u32, &[ 0x01u8, 0x02u8, 0x03u8, 0x04u8 ]).unwrap();
    test_prom.read_full_at(0x0123u32, &mut buffer).unwrap();
    assert_eq!(buffer, [ 0x01u8, 0x02u8, 0x03u8, 0x04u8 ]);

    // Check basic write/read at the second chip
    test_prom.write_all_at(0x10123u32, &[ 0x08u8, 0x07u8, 0x06u8, 0x05u8 ]).unwrap();
    test_prom.read_full_at(0x10123u32, &mut buffer).unwrap();

    assert_eq!(buffer, [ 0x08u8, 0x07u8, 0x06u8, 0x05u8 ]);

    // Make sure that didn't actually write to the first chip
    test_prom.read_full_at(0x0123u32, &mut buffer).unwrap();
    assert_eq!(buffer, [ 0x01u8, 0x02u8, 0x03u8, 0x04u8 ]);

    test_prom.flush().unwrap();
  }

  #[test]
  fn test_buffered_prom_chip_boundaries() {
    // The handy thing about this test is that a chip boundary will also
    // necessarily be a page boundary.
    let mut test_prom = TestBuffer::with_driver(TestCompositeProm::using_client(DummyPromBusClient::new()));

    let mut buffer = [ 0x00u8, 0x00u8, 0x00u8, 0x00u8 ];

    // The following write should wrap over the end of a chip boundary,
    // so if the PROM driver implementation is wrong we'll feck it up
    test_prom.write_all_at(0x00fffeu32, &[ 0xde, 0xad, 0xbe, 0xef ]).unwrap();

    // This read should 'appear' to work, because the read will also wrap
    // at the boundary
    test_prom.read_full_at(0x00fffeu32, &mut buffer).unwrap();
    assert_eq!(buffer, [ 0xde, 0xad, 0xbe, 0xef ]);

    // But we must check it's actually right...
    test_prom.read_full_at(0x10000u32, &mut buffer).unwrap();
    assert_eq!(buffer, [ 0xbe, 0xef, 0xff, 0xff ]);
  }

}