// 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.

//! StorageManager capsule

use core::cell::Cell;
use core::marker::PhantomData;
use core::mem::size_of;
use kernel::common::cells::TakeCell;
use kernel::hil;
use kernel::hil::flash::Flash;
use kernel::{AppId, AppSlice, Callback, Driver, Grant, ReturnCode, Shared};
use matcha_config::*;
use matcha_hal::dprintf;
use matcha_utils::SMC_PAGE_SIZE;
use matcha_utils::smc_is_page;
use matcha_utils::smc_ram_memcpy;
use matcha_utils::tar_loader::TarHeader;

const FLASH_END: usize = 16 * 1024 * 1024; // 16MiB SPI flash

// Returns true if the specified SMC page address fits in SPI flash
fn is_spi_page(addr: usize) -> bool {
    addr + SMC_PAGE_SIZE < FLASH_END
}

#[inline]
fn howmany(a: u32, b: u32) -> u32 { (a + b - 1) / b }
#[inline]
fn roundup(a: u32, b: u32) -> u32 { howmany(a, b) * b }

#[derive(Default)]
pub struct AppData {
    // XXX only need 1 callback since we're single-threaded
    pub get_name_callback: Option<Callback>,
    pub find_file_callback: Option<Callback>,
    pub read_page_callback: Option<Callback>,
    pub buffer: Option<AppSlice<Shared, u8>>,
}

#[derive(Clone, Copy, Debug, PartialEq)]
enum StorageState {
    Idle,
    FindFile(u32 /* cursor */), // NB: filename passed with allow
    ReadPage(u32 /* fid */, u32 /* offset */, u32 /* dest */),
    GetName(u32 /* cursor */),  // NB: filename returned with allow
}
impl StorageState {
    fn is_idle(&self) -> bool {
        match self {
            StorageState::Idle => true,
            _ => false,
        }
    }
}

pub struct StorageCapsule<'a, F: hil::flash::Flash + 'static> {
    app_data_grant: Grant<AppData>,
    pub current_app: Cell<Option<AppId>>,
    flash: Option<&'static capsules::virtual_flash::FlashUser<'static, F>>,
    flash_busy: Cell<bool>,
    read_page: TakeCell<'static, F::Page>,
    page_len: u32,
    state: Cell<StorageState>,
    phantom: PhantomData<&'a ()>,
}

impl<'a, F: hil::flash::Flash> StorageCapsule<'a, F> {
    pub fn new(app_data_grant: Grant<AppData>) -> Self {
        Self {
            app_data_grant: app_data_grant,
            current_app: Cell::new(None),
            flash: None,
            flash_busy: Cell::new(false),
            read_page: TakeCell::empty(),
            page_len: 0,
            state: Cell::new(StorageState::Idle),
            phantom: PhantomData,
        }
    }

    pub fn set_flash(
        &mut self,
        flash: &'static capsules::virtual_flash::FlashUser<'static, F>,
        read_page: &'static mut F::Page,
    ) {
        self.flash = Some(flash);
        self.read_page.replace(read_page);
        let mut page_len = 0;
        self.read_page.map(|page| {
            let mut_page = page.as_mut();
            page_len = mut_page.len() as u32;
        });
        self.page_len = page_len;
    }

    pub fn handle_command(
        &self,
        app_id: AppId,
        app_data: &mut AppData,
        minor_num: usize,
        arg2: usize,
        arg3: usize,
    ) -> ReturnCode {
        match minor_num {
            CMD_STORAGE_INIT => {
                self.current_app.set(Some(app_id));
                ReturnCode::SUCCESS
            }
            CMD_STORAGE_GET_NAME => {
                let state = self.state.get();
                if !state.is_idle() {
                    dprintf!("storage: GET_NAME: state {:?}\r\n", state);
                    return ReturnCode::EBUSY;
                }
                // NB: buffer is optional
                let fid = arg2 as u32;
                self.state.set(StorageState::GetName(fid));
                self.read_page(fid); // First read, kicks off state machine
                return ReturnCode::SUCCESS
            }
            CMD_STORAGE_FIND_FILE => {
                let state = self.state.get();
                if !state.is_idle() {
                    dprintf!("storage: FIND_FILE: state {:?}\r\n", state);
                    return ReturnCode::EBUSY;
                }
                // Verify the filename has been passed to allow and check
                // the string is valid so work inside the state machine
                // can just unwrap the result of str::from_utf8.
                if let Some(app_slice) = app_data.buffer.as_ref() {
                    if core::str::from_utf8(app_slice.as_ref()).is_ok() {
                        self.state.set(StorageState::FindFile(0));
                        self.read_page(0); // First read, kicks off state machine
                        return ReturnCode::SUCCESS
                    }
                }
                ReturnCode::EINVAL
            }
            CMD_STORAGE_READ_PAGE => {
                let state = self.state.get();
                if !state.is_idle() {
                    dprintf!("storage: READ: state {:?}\r\n", state);
                    return ReturnCode::EBUSY;
                }
                // NB: this checks the entire 4K page can be read from SPI;
                // we don't really handle reading less than a full page.
                let fid = arg2 as u32;
                if !is_spi_page(fid as usize) {
                    dprintf!("storage: READ: extends past the end-of-flash (fid  {})\r\n", fid);
                    return ReturnCode::EINVAL;
                }
                // The best we can do is check |dest| looks like an SMC
                // address; otherwise we just spray data where ever the
                // SecurityCoordinator tells us.
                let dest = arg3 as u32;
                if !smc_is_page(dest as usize) {
                    dprintf!("storage: READ: destination not in SMC (dest {:#x})\r\n", dest);
                    return ReturnCode::EINVAL;
                }
                self.state.set(StorageState::ReadPage(fid, 0, dest));
                self.read_page(fid); // First read, kicks off state machine
                ReturnCode::SUCCESS
            }
            _ => ReturnCode::EINVAL,
        }
    }

    pub fn handle_subscribe(
        &self,
        _app_id: AppId,
        app_data: &mut AppData,
        minor_num: usize,
        callback: Option<Callback>,
    ) -> ReturnCode {
        match minor_num {
            CMD_STORAGE_GET_NAME => {
                app_data.get_name_callback = callback;
                ReturnCode::SUCCESS
            }
            CMD_STORAGE_FIND_FILE => {
                app_data.find_file_callback = callback;
                ReturnCode::SUCCESS
            }
            CMD_STORAGE_READ_PAGE => {
                app_data.read_page_callback = callback;
                ReturnCode::SUCCESS
            }
            _ => ReturnCode::EINVAL,
        }
    }

    pub fn handle_allow(
        &self,
        _app_id: AppId,
        app_data: &mut AppData,
        minor_num: usize,
        slice: Option<AppSlice<Shared, u8>>,
    ) -> ReturnCode {
        match minor_num {
            CMD_STORAGE_GET_NAME
            | CMD_STORAGE_FIND_FILE => {
                app_data.buffer = slice;
                ReturnCode::SUCCESS
            }
            _ => ReturnCode::EINVAL,
        }
    }

    // GetName state machine callback on flash read complete.
    fn get_name_callback(&self) {
        self.current_app.get().map(|app_id| {
            let _ = self.app_data_grant.enter(app_id, |app_data, _| {
                self.read_page.map(|page| {
                    let mut_page = page.as_mut(); // XXX mut?
                    match self.state.get() {
                        StorageState::GetName(cursor) => {
                            let tar_header = TarHeader::from_bytes(mut_page);
                            if tar_header.has_magic() {
                                // Copy back filename if a buffer is setup.
                                let name_len = if let Some(app_slice) = app_data.buffer.as_mut() {
                                    let slice = app_slice.as_mut();
                                    let tar_bytes = tar_header.name().as_bytes();
                                    if slice.len() > tar_bytes.len() {
                                        &mut slice[..tar_bytes.len()].copy_from_slice(tar_bytes);
                                        tar_bytes.len()
                                    } else {
                                        slice.copy_from_slice(&tar_bytes[..slice.len()]);
                                        slice.len()
                                    }
                                } else { 0 };
                                let tar_size = tar_header.size();
                                let next_cursor = cursor
                                    + self.page_len
                                    + roundup(tar_size, size_of::<TarHeader>() as u32);
                                self.state.set(StorageState::Idle);
                                app_data.get_name_callback.map(|mut callback| {
                                    callback.schedule(1, name_len, next_cursor as usize);
                                });
                            } else {
                                // Tar header looks invalid, this is typically how a failed lookup ends.
                                self.state.set(StorageState::Idle);
                                app_data.get_name_callback.map(|mut callback| {
                                    callback.schedule(0, 0, 0);
                                });
                            }
                        }
                        _ => panic!("get_name: bad state"),
                    }
                });
            });
        });
    }

    // FindFile state machine callback on flash read complete.
    fn find_file_callback(&self) {
        self.current_app.get().map(|app_id| {
            let _ = self.app_data_grant.enter(app_id, |app_data, _| {
                self.read_page.map(|page| {
                    let mut_page = page.as_mut(); // XXX mut?
                    match self.state.get() {
                        StorageState::FindFile(cursor) => {
                            let tar_header = TarHeader::from_bytes(mut_page);
                            if tar_header.has_magic() {
                                // NB: app_slice is checked to be valid in command; we assume
                                //   it does not change while we're running
                                let app_slice = app_data.buffer.as_ref().unwrap();
                                let name = core::str::from_utf8(app_slice.as_ref()).unwrap();
                                let tar_name = tar_header.name();
                                let tar_size = tar_header.size();
                                if tar_name == name {
                                    // Found file, return fid = the offset to the first byte of data
                                    self.state.set(StorageState::Idle);
                                    app_data.find_file_callback.map(|mut callback| {
                                        callback.schedule(
                                            1,                                          // success
                                            (cursor as usize) + size_of::<TarHeader>(), // fid
                                            tar_size as usize,                          // file size (bytes)
                                        );
                                    });
                                } else {
                                    // Advance to the next file if possible. We depend on
                                    // each header having a magic marker; otherwise we use
                                    // a fixed bound on the flash size to (try to) avoid
                                    // reading past the end of memory.
                                    let next_cursor = cursor
                                        + self.page_len
                                        + roundup(tar_size, size_of::<TarHeader>() as u32);
                                    if ((next_cursor as usize) + size_of::<TarHeader>()) >= FLASH_END {
                                        dprintf!("storage: {} not found, end of flash\r\n", name);
                                        self.state.set(StorageState::Idle);
                                        app_data.find_file_callback.map(|mut callback| {
                                            callback.schedule(0, 0, 0);
                                        });
                                    } else {
                                        self.state.set(StorageState::FindFile(next_cursor));
                                        // NB: defer read to below so we release self.read_page
                                    }
                                }
                            } else {
                                // Tar header looks invalid, this is typically how a failed lookup ends.
                                self.state.set(StorageState::Idle);
                                app_data.find_file_callback.map(|mut callback| {
                                    callback.schedule(0, 0, 0);
                                });
                            }
                        }
                        _ => panic!("find_file: bad state"),
                    }
                });
                match self.state.get() {
                    StorageState::FindFile(cursor) => self.read_page(cursor),
                    _ => {}
                }
            });
        });
    }

    // ReadPage state machine callback on flash read complete.
    fn read_page_callback(&self) {
        self.current_app.get().map(|app_id| {
            let _ = self.app_data_grant.enter(app_id, |app_data, _| {
                self.read_page.map(|page| {
                    let mut_page = page.as_mut();
                    match self.state.get() {
                        StorageState::ReadPage(fid, offset, dest) => {
//                            dprintf!("copy page {} to {:x}\r\n", fid + offset, dest);
                            smc_ram_memcpy(
                                mut_page,
                                dest,
                                self.page_len as usize, // XXX short read?
                            );
                            let next_offset = offset + self.page_len;
                            if next_offset >= (SMC_PAGE_SIZE as u32) {
                                self.state.set(StorageState::Idle);
                                app_data.read_page_callback.map(|mut callback| {
                                    callback.schedule(1, fid as usize, SMC_PAGE_SIZE);
                                });
                            } else {
                                self.state.set(
                                    StorageState::ReadPage(
                                        fid,
                                        next_offset,
                                        dest + self.page_len,
                                    )
                                );
                                // NB: defer read to below so we release self.read_page
                            }
                        }
                        _ => panic!("read_page: bad state"),
                    }
                });
                match self.state.get() {
                    StorageState::ReadPage(fid, offset, _) => self.read_page(fid + offset),
                    _ => {}
                }
            });
        });
    }

    // Kicks off a SPI flash read at the byte offset |page|.
    fn read_page(&self, page: u32) {
        debug_assert!((page % self.page_len) == 0);
        self.flash.map(|flash| {
            self.read_page.take().map(|read_page| {
                self.flash_busy.set(true);
                if let Err((_, buf)) = flash.read_page((page / self.page_len) as usize, read_page) {
                    self.read_page.replace(buf);
                }
            });
        });
    }
}

/// Driver impl just enters the app_data grant and delegates to StorageCapsule.

impl<'a, F: hil::flash::Flash> Driver for StorageCapsule<'a, F> {
    fn subscribe(&self, minor_num: usize, callback: Option<Callback>, app_id: AppId) -> ReturnCode {
        self.app_data_grant
            .enter(app_id, |app_data, _| {
                self.handle_subscribe(app_id, app_data, minor_num, callback)
            })
            .unwrap_or_else(|err| err.into())
    }

    fn command(&self, minor_num: usize, r2: usize, r3: usize, app_id: AppId) -> ReturnCode {
        self.app_data_grant
            .enter(app_id, |app_data, _| {
                self.handle_command(app_id, app_data, minor_num, r2, r3)
            })
            .unwrap_or_else(|err| err.into())
    }

    fn allow(
        &self,
        app_id: AppId,
        minor_num: usize,
        slice: Option<AppSlice<Shared, u8>>,
    ) -> ReturnCode {
        self.app_data_grant
            .enter(app_id, |app_data, _| {
                self.handle_allow(app_id, app_data, minor_num, slice)
            })
            .unwrap_or_else(|err| err.into())
    }
}

impl<'a, F: hil::flash::Flash> hil::flash::Client<capsules::virtual_flash::FlashUser<'a, F>>
    for StorageCapsule<'a, F>
{
    fn read_complete(&self, read_page: &'static mut F::Page, _err: kernel::hil::flash::Error) {
        self.read_page.replace(read_page);
        self.flash_busy.set(false);
        match self.state.get() {
            StorageState::GetName(..) => self.get_name_callback(),
            StorageState::FindFile(..) => self.find_file_callback(),
            StorageState::ReadPage(..) => self.read_page_callback(),
            _ => panic!("bad state"),
        }
    }
    fn write_complete(
        &self,
        _: &'static mut <F as kernel::hil::flash::Flash>::Page,
        _: kernel::hil::flash::Error,
    ) {
        todo!()
    }
    fn erase_complete(&self, _: kernel::hil::flash::Error) {
        todo!()
    }
}