// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you 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 crate::{Error, ErrorKind, Result, Uuid};
#[cfg(not(feature = "std"))]
use alloc::{borrow::ToOwned, vec::Vec};
use optee_utee_sys as raw;

pub struct LoadablePlugin {
    uuid: Uuid,
}

pub struct LoadablePluginCommand<'a> {
    plugin: &'a LoadablePlugin,
    cmd_id: u32,
    sub_cmd_id: u32,
    buffer: Vec<u8>,
}

impl LoadablePlugin {
    pub fn new(uuid: &Uuid) -> Self {
        Self {
            uuid: uuid.to_owned(),
        }
    }
    /// Invoke plugin with given request data, use when you want to post something into REE.
    /// ``` rust,no_run
    /// # use optee_utee::{LoadablePlugin, Uuid};
    /// # fn main() -> optee_utee::Result<()> {
    /// # let uuid = Uuid::parse_str("").unwrap();
    /// # let command_id = 0;
    /// # let subcommand_id = 0;
    /// # let request_data = [0_u8; 0];
    /// let plugin = LoadablePlugin::new(&uuid);
    /// let result = plugin.invoke(command_id, subcommand_id, &request_data)?;
    /// # Ok(())
    /// # }
    /// ```
    /// Caution: the size of the shared buffer is set to the len of data, you could get a
    ///          ShortBuffer error if Plugin return more data than shared buffer, in that case,
    ///          use invoke_with_capacity and set the capacity manually.
    pub fn invoke(&self, command_id: u32, subcommand_id: u32, data: &[u8]) -> Result<Vec<u8>> {
        self.invoke_with_capacity(command_id, subcommand_id, data.len())
            .chain_write_body(data)
            .call()
    }
    /// Construct a command with shared buffer up to capacity size, write the buffer and call it
    /// manually, use when you need to control details of the invoking process.
    /// ```no_run
    /// # use optee_utee::{Uuid, LoadablePlugin};
    /// # fn main() -> optee_utee::Result<()> {
    /// # let plugin = LoadablePlugin::new(&Uuid::parse_str("").unwrap());
    /// # let request_data = [0_u8; 0];
    /// # let command_id = 0;
    /// # let sub_command_id = 0;
    /// # let capacity = 0;
    /// let mut cmd = plugin.invoke_with_capacity(command_id, sub_command_id, capacity);
    /// cmd.write_body(&request_data);
    /// let result = cmd.call()?;
    /// # Ok(())
    /// # }
    /// ```
    /// You can also imply a wrapper for performance, for example, imply a std::io::Write so
    /// serde_json can write to the buffer directly.
    /// ```no_run
    /// # use optee_utee::{LoadablePluginCommand, Uuid, LoadablePlugin, trace_println};
    /// # use optee_utee::ErrorKind;
    /// # fn main() -> optee_utee::Result<()> {
    /// # let command_id = 0;
    /// # let subcommand_id = 0;
    /// # let capacity = 0;
    /// # let plugin = LoadablePlugin::new(&Uuid::parse_str("").unwrap());
    /// struct Wrapper<'a, 'b>(&'b mut LoadablePluginCommand<'a>);
    /// impl<'a, 'b> std::io::Write for Wrapper<'a, 'b> {
    ///     fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
    ///         self.0.write_body(buf);
    ///         Ok(buf.len())
    ///     }
    ///     fn flush(&mut self) -> std::io::Result<()> {
    ///         Ok(())
    ///     }
    /// }
    /// // serialize data into command directly
    /// let request_data = serde_json::json!({
    ///     "age": 100,
    ///     "name": "name"
    /// });
    /// let mut cmd = plugin.invoke_with_capacity(command_id, subcommand_id, capacity);
    /// serde_json::to_writer(Wrapper(&mut cmd), &request_data).map_err(|err| {
    ///     trace_println!("serde error: {:?}", err);
    ///     ErrorKind::Unknown
    /// })?;
    /// let result = cmd.call()?;
    ///
    /// # Ok(())
    /// # }
    /// ```
    /// Notice: the shared buffer could grow to fit the request data automatically.
    pub fn invoke_with_capacity(
        &self,
        command_id: u32,
        subcommand_id: u32,
        capacity: usize,
    ) -> LoadablePluginCommand<'_> {
        LoadablePluginCommand::new_with_capacity(self, command_id, subcommand_id, capacity)
    }
}

impl<'a> LoadablePluginCommand<'a> {
    // use this to write request body if needed
    pub fn write_body(&mut self, data: &[u8]) {
        self.buffer.extend_from_slice(data);
    }
    // same with write_body, but chainable
    pub fn chain_write_body(mut self, data: &[u8]) -> Self {
        self.write_body(data);
        self
    }
    // invoke the command, and get result from it
    pub fn call(self) -> Result<Vec<u8>> {
        let mut outlen: usize = 0;
        let mut buffer = self.buffer;
        buffer.resize(buffer.capacity(), 0); // resize to capacity first
        match unsafe {
            raw::tee_invoke_supp_plugin(
                self.plugin.uuid.as_raw_ptr(),
                self.cmd_id,
                self.sub_cmd_id,
                // convert the pointer manually, as in some platform c_char is i8
                buffer.as_mut_slice().as_mut_ptr() as *mut _,
                buffer.len(),
                &mut outlen as *mut usize,
            )
        } {
            raw::TEE_SUCCESS => {
                if outlen > buffer.len() {
                    return Err(ErrorKind::ShortBuffer.into());
                }
                buffer.resize(outlen, 0);
                Ok(buffer)
            }
            code => Err(Error::from_raw_error(code)),
        }
    }
}

impl<'a> LoadablePluginCommand<'a> {
    fn new_with_capacity(
        plugin: &'a LoadablePlugin,
        cmd_id: u32,
        sub_cmd_id: u32,
        capacity: usize,
    ) -> Self {
        Self {
            plugin,
            cmd_id,
            sub_cmd_id,
            buffer: Vec::with_capacity(capacity),
        }
    }
}

#[cfg(test)]
pub mod test_loadable_plugin {
    extern crate std;
    use super::*;
    use alloc::string::ToString;
    use optee_utee_sys::{mock_api, mock_utils::SERIAL_TEST_LOCK};
    use rand::distributions::Alphanumeric;
    use rand::Rng;

    fn generate_random_bytes(len: usize) -> Vec<u8> {
        rand::thread_rng()
            .sample_iter(&Alphanumeric)
            .take(len)
            .collect()
    }

    fn generate_test_pairs(
        request_size: usize,
        response_size: usize,
    ) -> (u32, u32, Vec<u8>, Vec<u8>) {
        let cmd: u32 = rand::random();
        let sub_cmd: u32 = rand::random();
        let random_request: Vec<u8> = generate_random_bytes(request_size);
        let random_response: Vec<u8> = generate_random_bytes(response_size);
        (cmd, sub_cmd, random_request, random_response)
    }

    fn random_uuid() -> Uuid {
        Uuid::new_raw(
            rand::random(),
            rand::random(),
            rand::random(),
            rand::random(),
        )
    }

    fn expect_success_request(
        ctx: &mock_api::extension::__tee_invoke_supp_plugin::Context,
        exp_uuid: &Uuid,
        exp_cmd: u32,
        exp_sub_cmd: u32,
        exp_request: &[u8],
        exp_response: &[u8],
    ) {
        let exp_request = exp_request.to_vec();
        let exp_response = exp_response.to_vec();
        let exp_uuid = exp_uuid.to_string();
        ctx.expect()
            .return_once_st(move |uuid, cmd, sub_cmd, buf, len, outlen| {
                let request_uuid = Uuid::from(unsafe { *uuid }).to_string();
                debug_assert_eq!(exp_uuid, request_uuid);
                debug_assert_eq!(cmd, exp_cmd);
                debug_assert_eq!(sub_cmd, exp_sub_cmd);
                debug_assert_eq!(
                    unsafe { core::slice::from_raw_parts(buf as *mut u8, exp_request.len()) },
                    exp_request.as_slice()
                );
                debug_assert!(len >= exp_response.len());
                let buffer: &mut [u8] =
                    unsafe { core::slice::from_raw_parts_mut(buf as *mut u8, len) };
                buffer[0..exp_response.len()].copy_from_slice(&exp_response);
                unsafe { *outlen = exp_response.len() };
                raw::TEE_SUCCESS
            });
    }

    #[test]
    fn test_invoke() {
        let _lock = SERIAL_TEST_LOCK.lock().expect("should get the lock");

        let uuid: Uuid = random_uuid();
        let plugin = LoadablePlugin::new(&uuid);
        const REQUEST_LEN: usize = 32;
        let run_test = |request_size: usize, response_size: usize| {
            let (cmd, sub_cmd, request, exp_response) =
                generate_test_pairs(request_size, response_size);
            let fn1 = mock_api::extension::tee_invoke_supp_plugin_context();
            expect_success_request(&fn1, &uuid, cmd, sub_cmd, &request, &exp_response);
            let response = plugin.invoke(cmd, sub_cmd, &request).expect("should be ok");
            std::println!("*TA*: response is {:?}", response);
            debug_assert_eq!(response, exp_response);
        };

        // test calling with output size less than input
        run_test(REQUEST_LEN, REQUEST_LEN / 2);
        // test calling with output size equals to input
        run_test(REQUEST_LEN, REQUEST_LEN);
        // test calling with output size greater than input.
        // Mark: Without explicitly setting the response size, this function
        // must not be called with a response size larger than the request size.
        {
            let (cmd, sub_cmd, request, exp_response) =
                generate_test_pairs(REQUEST_LEN, 2 * REQUEST_LEN);
            let fn1 = mock_api::extension::tee_invoke_supp_plugin_context();
            fn1.expect().return_once_st(move |_, _, _, _, _, outlen| {
                unsafe { *outlen = exp_response.len() };
                raw::TEE_SUCCESS
            });
            let err = plugin
                .invoke(cmd, sub_cmd, &request)
                .expect_err("should be err");
            debug_assert_eq!(err.kind(), ErrorKind::ShortBuffer);
        }
    }

    // This test is equivalent to test_invoke, with the added verification that
    // capacity permits the response size to be larger than the request.
    #[test]
    fn test_invoke_with_capacity() {
        let _lock = SERIAL_TEST_LOCK.lock().expect("should get the lock");
        let uuid: Uuid = random_uuid();
        let plugin = LoadablePlugin::new(&uuid);
        const RESPONSE_LEN: usize = 32;

        let run_test = |request_size: usize, response_size: usize| {
            let (cmd, sub_cmd, request, exp_response) =
                generate_test_pairs(request_size, response_size);
            let fn1 = mock_api::extension::tee_invoke_supp_plugin_context();
            expect_success_request(&fn1, &uuid, cmd, sub_cmd, &request, &exp_response);

            let response = plugin
                .invoke_with_capacity(cmd, sub_cmd, exp_response.len())
                .chain_write_body(&request)
                .call()
                .unwrap();
            std::println!("*TA*: response is {:?}", response);
            debug_assert_eq!(response, exp_response);
        };

        // test calling with output size less than input
        run_test(2 * RESPONSE_LEN, RESPONSE_LEN);
        // test calling with output size equals to input
        run_test(RESPONSE_LEN, RESPONSE_LEN);
        // test calling with output size greater than input
        run_test(RESPONSE_LEN / 2, RESPONSE_LEN);
    }

    #[test]
    fn test_invoke_with_writer() {
        let _lock = SERIAL_TEST_LOCK.lock().expect("should get the lock");
        let uuid: Uuid = random_uuid();
        let plugin = LoadablePlugin::new(&uuid);
        let fn1 = mock_api::extension::tee_invoke_supp_plugin_context();
        // impl a writer for Command
        struct Wrapper<'a, 'b>(&'b mut LoadablePluginCommand<'a>);
        impl<'a, 'b> std::io::Write for Wrapper<'a, 'b> {
            fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
                self.0.write_body(buf);
                Ok(buf.len())
            }
            fn flush(&mut self) -> std::io::Result<()> {
                Ok(())
            }
        }
        // serialize data into command directly
        let test_data = serde_json::json!({
            "code": 100,
            "message": "error"
        });
        let mut exp_request = serde_json::to_vec(&test_data).unwrap();
        let buffer_len = exp_request.len() * 2;
        let (cmd, sub_cmd, _, exp_response) = generate_test_pairs(0, buffer_len);
        let mut plugin_cmd = plugin.invoke_with_capacity(cmd, sub_cmd, buffer_len);
        exp_request.resize(exp_response.len(), 0);

        expect_success_request(&fn1, &uuid, cmd, sub_cmd, &exp_request, &exp_response);
        serde_json::to_writer(Wrapper(&mut plugin_cmd), &test_data).unwrap();
        let response = plugin_cmd.call().unwrap();
        std::println!("*TA*: response is {:?}", response);
        debug_assert_eq!(response, exp_response);
    }

    #[test]
    fn test_invoke_with_no_data() {
        let _lock = SERIAL_TEST_LOCK.lock().expect("should get the lock");

        let uuid: Uuid = random_uuid();
        let plugin = LoadablePlugin::new(&uuid);
        let fn1 = mock_api::extension::tee_invoke_supp_plugin_context();
        const OUTPUT_LEN: usize = 50;
        let (cmd, sub_cmd, request, exp_response) = generate_test_pairs(0, OUTPUT_LEN);
        expect_success_request(&fn1, &uuid, cmd, sub_cmd, &request, &exp_response);

        let response = plugin
            .invoke_with_capacity(cmd, sub_cmd, OUTPUT_LEN)
            .call()
            .unwrap();
        std::println!("*TA*: response is {:?}", response);
        debug_assert_eq!(response, exp_response);
    }
}