// 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::dprintf;

extern crate alloc;
use alloc::string::{String, ToString};
use alloc::vec::Vec;
use core::cell::Cell;
use core::mem::{size_of, transmute};
use core::ptr;
use libtock::futures;
use libtock::syscalls;
use matcha_config::*;
use serde::{Deserialize, Serialize};

#[derive(Debug, Serialize, Deserialize)]
pub enum SECRequest<'a> {
    FindFile(&'a str),     // Find file by name -> (/*fid*/ u32, /*size*/ u32)
    GetFilePage(u32, u32), // Read data from fid at offset -> <attached page>

    Test(/*count*/ u32), // Scribble on count words of supplied page

    // NB: must be last to match up with stripped down rootserver driver
    GetBuiltins,           // Get package names -> Vec(String)
}

#[derive(Debug, Serialize, Deserialize)]
pub struct GetBuiltinsResponse {
    pub names: Vec<String>,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct FindFileResponse {
    pub fid: u32,        // Unique file identifier
    pub size_bytes: u32, // File size
}

#[derive(Debug, Serialize, Deserialize)]
pub struct GetFilePageResponse {
    pub count_bytes: u32, // Read count (bytes)
}

#[derive(Debug, Eq, PartialEq)]
pub enum SECRequestError {
    Success = 0,
    DeserializeFailed,
    SerializeFailed,
    UnknownError,
    PageInvalid,
    FileNotFound,
    FileOffsetInvalid,
    // Generic errors.
    SendFailed,
    RecvFailed,
}
impl From<SECRequestError> for Result<(), SECRequestError> {
    fn from(err: SECRequestError) -> Result<(), SECRequestError> {
        if err == SECRequestError::Success {
            Ok(())
        } else {
            Err(err)
        }
    }
}

//------------------------------------------------------------------------------

pub struct MailboxClient {}

pub type WaitResult = (Option<bool>, usize, Option<usize>);

impl MailboxClient {
    pub fn init() {
        let _ = syscalls::command(CAPSULE_MAILBOX, CMD_MAILBOX_INIT, 0, 0);
        let _ = syscalls::command(CAPSULE_STORAGE, CMD_STORAGE_INIT, 0, 0);
    }

    // Support for receiving inbound mailbox messages. This is split into
    // multiple methods to support a race-free usage pattern:
    //
    // use MailboxClient::WaitResult;
    // MailboxClient::init();
    // let wait_result = Cell::new(WaitResult::default());
    // MailboxClient::wait_message_init(&wait_result, unsafe { &mut MESSAGE_BUF });
    // loop {
    //    MailboxClient::wait_message_async(&wait_result).await;
    //    // Status, message size (bytes), Option<attached page>
    //    let (_, request_size, opt_page) = wait_result.take();
    //
    //     .. process received message ..
    //
    //     MailboxClient::send_message_sync(reply_size, reply_slice);
    // }
    // MailboxClient::wait_message_cleanup();
    //
    // TODO(sleffler): maybe wrap in a struct w/ RAII cleanup

    // Sets up the caller to receive a mailbox message in |buffer| using
    // |wait_result| to pass data about the received message from the irq
    // handler. The operation is completed asynchronously. The caller must
    // poll the implied future until |wait_result| holds the received
    // message data; typically using |wait_message_async|, Note that
    // |wait_message_setup| registers a subscriber for callbacks from the
    // irq handler; this must be cleared before |wait_result| is
    // deallocated--usually by calling |wait_message_cleanup|.
    pub fn wait_message_setup(wait_result: &Cell<WaitResult>, buffer: &mut [u8]) {
        // Callback on recv complete.
        unsafe extern "C" fn wait_callback(ok: usize, len: usize, page: usize, data: usize) {
            let wait_result = &*(data as *const Cell<WaitResult>);
            wait_result.set((
                Some(ok != 0),
                len,
                if page == 0 { None } else { Some(page) },
            ));
        }

        // Give Tock references to our callback, result, +buffer to write;
        // then wait until the callback sets the message status. Finally,
        // clear references to our local state.
        unsafe {
            syscalls::raw::allow(
                CAPSULE_MAILBOX,
                CMD_MAILBOX_RECV,
                buffer.as_mut_ptr(),
                buffer.len(),
            );
            syscalls::raw::subscribe(
                CAPSULE_MAILBOX,
                CMD_MAILBOX_RECV,
                wait_callback as *const _,
                wait_result as *const _ as usize,
            );
        }
    }
    // Cleans up state setup by |wait_message_setup|.
    pub fn wait_message_cleanup() {
        unsafe {
            syscalls::raw::subscribe(CAPSULE_MAILBOX, 0, ptr::null(), 0);
            syscalls::raw::allow(CAPSULE_MAILBOX, 0, ptr::null_mut(), 0);
        }
    }
    // Waits for a mailbox message to be received.
    pub async fn wait_message_async(wait_result: &Cell<WaitResult>) {
        futures::wait_until(|| (wait_result.get().0.is_some())).await;
    }

    // Waits for an inbound  message that it writes to |buffer|. The msg
    // size (in bytes) is returned along with any associated page frame
    #[allow(dead_code)]
    pub async fn wait_message_sync(buffer: &mut [u8]) -> (usize, Option<usize>) {
        let wait_result = Cell::new(WaitResult::default());
        Self::wait_message_setup(&wait_result, buffer);

        Self::wait_message_async(&wait_result).await;
        Self::wait_message_cleanup();

        // Status, message size (bytes), attached page (or 0)
        let (_, msg_size, attached_page) = wait_result.into_inner();
        (msg_size, attached_page)
    }

    // Sends the message of |size| bytes that resides in |buffer|
    // XXX why isn't buffer.len() used instead of |size|?
    pub fn send_message_sync(size: usize, buffer: &[u8]) {
        unsafe {
            syscalls::raw::allow(
                CAPSULE_MAILBOX,
                CMD_MAILBOX_SEND,
                buffer.as_ptr() as *mut u8,
                buffer.len(),
            );
            let _ = syscalls::command(CAPSULE_MAILBOX, CMD_MAILBOX_SEND, size, 0);
            syscalls::raw::allow(CAPSULE_MAILBOX, CMD_MAILBOX_SEND, ptr::null_mut(), 0);
        }
    }

    // Parse & dispatch the request in |request_slice| and construct the result
    // in |reply_slice|. |opt_page| holds any page attached to the request.
    pub async fn dispatch(
        request_slice: &[u8],
        opt_page: Option<usize>,
        reply_slice: &mut [u8],
    ) -> Result<usize, SECRequestError> {
        match postcard::from_bytes::<SECRequest>(request_slice) {
            Ok(request) => {
//                dprintf!("dispatch {:?} opt_page {:?}\r\n", &request, opt_page);
                match request {
                    SECRequest::GetBuiltins => {
                        MailboxClient::get_builtins_async(reply_slice).await
                    }
                    SECRequest::FindFile(name) => {
                        MailboxClient::find_file_async(name, reply_slice).await
                    }
                    SECRequest::GetFilePage(fid, offset) => {
                        MailboxClient::get_file_page_async(fid, offset, opt_page, reply_slice).await
                    }
                    SECRequest::Test(count) => {
                        MailboxClient::test_sync(count as usize, opt_page, reply_slice)
                    }
                }
            }
            Err(e) => {
                dprintf!("Invalid request: {:?}\r\n", e);
                Err(SECRequestError::DeserializeFailed)
            }
        }
    }

    // Returns the names of objects that can be looked up using FindFile.
    async fn get_builtins_async(reply_buffer: &mut [u8]) -> Result<usize, SECRequestError> {
        const NAME_SIZE: usize = 100;   // Arbitrary but matches TarHeader.name
        let mut name: [u8; NAME_SIZE] = [0u8; NAME_SIZE];
        let mut names = Vec::new();

        let mut fid: u32 = 0; // First file in SPI
        loop {
            // Status, name length, next fid
            type GetNameResult = (Option<bool>, usize, usize);
            let get_name_result: Cell<GetNameResult> = Cell::new((None, 0, 0));

            unsafe extern "C" fn get_name_callback(ok: usize, name_len: usize, next_fid: usize, data: usize) {
                let get_name_result = &*(data as *const Cell<GetNameResult>);
                get_name_result.set((Some(ok != 0), name_len, next_fid));
            }
            unsafe {
                syscalls::raw::allow(
                    CAPSULE_STORAGE,
                    CMD_STORAGE_GET_NAME,
                    name.as_mut_ptr() as _,
                    NAME_SIZE,
                );
                syscalls::raw::subscribe(
                    CAPSULE_STORAGE,
                    CMD_STORAGE_GET_NAME,
                    get_name_callback as *const _,
                    &get_name_result as *const _ as usize,
                );
                let _ = syscalls::command(
                    CAPSULE_STORAGE,
                    CMD_STORAGE_GET_NAME,
                    fid as usize,
                    0,
                );
            }
            futures::wait_until(|| (get_name_result.get().0.is_some())).await;
            unsafe {
                syscalls::raw::subscribe(CAPSULE_STORAGE, 0, ptr::null(), 0);
                syscalls::raw::allow(CAPSULE_STORAGE, 0, ptr::null_mut(), 0);
            }

            let (status, name_len, next_fid) = get_name_result.into_inner();
            if !status.unwrap() { break; }
            if let Ok(str) = core::str::from_utf8(&name[..name_len]) {
                // XXX maybe explicitly include instead of filtering out?
                match str {
                    // Filter bootstrap artifacts.
                    "kernel" | "capdl-loader" | "matcha-tock-bundle.bin" => {}
                    _ => {
                        // Filter CAmkES component executables.
                        if !str.ends_with("_group_bin") {
                            names.push(str.to_string());
                        }
                    }
                }
            }
            fid = next_fid as u32;
        }
        postcard::to_slice(&GetBuiltinsResponse { names }, reply_buffer)
            .or(Err(SECRequestError::SerializeFailed))
            .map(|slice| slice.len())
    }

    // Lookup a file by name. A file identifier (the offset of the file's
    // data in SPI flash) and the data size (in bytes) is returned.
    // XXX could use GET_NAME
    async fn find_file_async(
        name: &str,
        reply_buffer: &mut [u8],
    ) -> Result<usize, SECRequestError> {
        // Status, file identifier, file size (bytes)
        type FindResult = (Option<bool>, usize, usize);
        let find_result: Cell<FindResult> = Cell::new((None, 0, 0));

        unsafe extern "C" fn find_callback(ok: usize, fid: usize, size_bytes: usize, data: usize) {
            let find_result = &*(data as *const Cell<FindResult>);
            find_result.set((Some(ok != 0), fid, size_bytes));
        }
        unsafe {
            syscalls::raw::allow(
                CAPSULE_STORAGE,
                CMD_STORAGE_FIND_FILE,
                name.as_ptr() as _, // XXX const
                name.len(),
            );
            syscalls::raw::subscribe(
                CAPSULE_STORAGE,
                CMD_STORAGE_FIND_FILE,
                find_callback as *const _,
                &find_result as *const _ as usize,
            );
            let _ = syscalls::command(CAPSULE_STORAGE, CMD_STORAGE_FIND_FILE, 0, 0);
        }
        futures::wait_until(|| (find_result.get().0.is_some())).await;
        unsafe {
            syscalls::raw::subscribe(CAPSULE_STORAGE, 0, ptr::null(), 0);
            syscalls::raw::allow(CAPSULE_STORAGE, 0, ptr::null_mut(), 0);
        }

        let (ok, fid, size_bytes) = find_result.into_inner();
        if ok.unwrap() {
            postcard::to_slice(
                &FindFileResponse {
                    fid: fid as u32,
                    size_bytes: size_bytes as u32,
                },
                reply_buffer,
            )
            .or(Err(SECRequestError::SerializeFailed))
            .map(|slice| slice.len())
        } else {
            Err(SECRequestError::FileNotFound)
        }
    }

    // Get/read 4KiB (or less at EOF) from |fid| starting at |offset| from
    // the start of the file. Data is returned in an associated page that
    // must reside in SMC memory
    async fn get_file_page_async(
        fid: u32,
        offset: u32,
        opt_page: Option<usize>,
        reply_buffer: &mut [u8],
    ) -> Result<usize, SECRequestError> {
        // Status, read count (bytes)
        type ReadResult = (Option<bool>, usize);
        let read_result: Cell<ReadResult> = Cell::new((None, 0));

        unsafe extern "C" fn read_callback(
            ok: usize,
            _fid: usize,
            count_bytes: usize,
            data: usize,
        ) {
            let read_result = &*(data as *const Cell<ReadResult>);
            read_result.set((Some(ok != 0), count_bytes));
        }
        unsafe {
// XXX check returns
            syscalls::raw::subscribe(
                CAPSULE_STORAGE,
                CMD_STORAGE_READ_PAGE,
                read_callback as *const _,
                &read_result as *const _ as usize,
            );
            // NB: cannot use allow to share opt_page because the address
            // is in the SMC memory and not the app/process. Instead we
            // pass an opaque value through command.
            let _ = syscalls::command(
                CAPSULE_STORAGE,
                CMD_STORAGE_READ_PAGE,
                (fid + offset) as usize,
                opt_page.unwrap() as _,
            );
        }
        futures::wait_until(|| (read_result.get().0.is_some())).await;
        unsafe { syscalls::raw::subscribe(CAPSULE_STORAGE, 0, ptr::null(), 0); }

        let (ok, count_bytes) = read_result.into_inner();
        if ok.unwrap() {
            postcard::to_slice(
                &GetFilePageResponse {
                    count_bytes: count_bytes as u32,
                },
                reply_buffer,
            )
            .or(Err(SECRequestError::SerializeFailed))
            .map(|slice| slice.len())
        } else {
            Err(SECRequestError::FileOffsetInvalid) // XXX guess
        }
    }

    fn test_sync(
        count: usize,
        opt_page: Option<usize>,
        _reply_buffer: &mut [u8],
    ) -> Result<usize, SECRequestError> {
        if let Some(page) = opt_page {
            if count >= (4096 / size_of::<u32>()) {
                return Err(SECRequestError::PageInvalid); // XXX
            }
            // XXX cannot address page; must push to capsule
            unsafe {
                dprintf!("page {:#x}", page);
                let dst: *mut u32 = transmute(page);
                // Tweak the first and last dword of the message.
                dprintf!("request[0] = {:#x}\r\n", dst.add(0).read());
                dst.add(0).write(0x12345678);
                if count > 1 {
                    let last = count - 1;
                    dprintf!("request[{}] = {:#x}\r\n", last, dst.add(last).read());
                    dst.add(last).write(0x87654321);
                }
            }
            Ok(0) // NB: modified page has result
        } else {
            Err(SECRequestError::PageInvalid)
        }
    }
}