//
// Copyright (C) 2024 The Android Open Source Project
//
// 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
//
//      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::Error;
use alloc::boxed::Box;
use core::marker::PhantomData;
use core::ptr;
use optee_utee::{ParamType, Uuid};
use optee_utee_sys::{
    Memref, TEE_CloseTASession, TEE_InvokeTACommand, TEE_OpenTASession, TEE_Param,
    TEE_TASessionHandle, Value, TEE_SUCCESS, TEE_TIMEOUT_INFINITE,
};

// Matches optee_os/lib/libutee/include/pta_system.h
pub const PTA_SYSTEM_ADD_RNG_ENTROPY: u32 = 0;
#[allow(dead_code)]
pub const PTA_SYSTEM_DERIVE_TA_UNIQUE_KEY: u32 = 1;
pub const PTA_SYSTEM_UUID: &'static str = "3a2f8978-5dc0-11e8-9c2d-fa7ae01bbebc";

#[allow(dead_code)]
pub enum Direction {
    In,
    Out,
    Both,
}

pub struct Param<'a> {
    raw: TEE_Param,
    param_type: ParamType,
    phantom: PhantomData<&'a mut i32>,
}

impl<'a> Param<'a> {
    pub fn empty() -> Params<'a> {
        Params(Param::none(), Param::none(), Param::none(), Param::none())
    }

    pub fn none() -> Param<'a> {
        Param {
            raw: TEE_Param { value: Value { a: 0, b: 0 } },
            param_type: ParamType::None,
            phantom: PhantomData,
        }
    }

    pub fn from_mut_slice(data: &'a mut [u8], direction: Direction) -> Param<'a> {
        Param {
            raw: TEE_Param { memref: Memref { buffer: data.as_mut_ptr() as _, size: data.len() } },
            param_type: match direction {
                Direction::In => ParamType::MemrefInput,
                Direction::Out => ParamType::MemrefOutput,
                Direction::Both => ParamType::MemrefInout,
            },
            phantom: PhantomData,
        }
    }

    pub fn from_slice(data: &'a [u8]) -> Param<'a> {
        Param {
            raw: TEE_Param { memref: Memref { buffer: data.as_ptr() as _, size: data.len() } },
            param_type: ParamType::MemrefInput,
            phantom: PhantomData,
        }
    }

    #[allow(dead_code)]
    pub fn from_value(a: u32, b: u32, direction: Direction) -> Param<'a> {
        Param {
            raw: TEE_Param { value: Value { a: a, b: b } },
            param_type: match direction {
                Direction::In => ParamType::ValueInput,
                Direction::Out => ParamType::ValueOutput,
                Direction::Both => ParamType::ValueInout,
            },
            phantom: PhantomData,
        }
    }
}

pub struct Params<'a>(pub Param<'a>, pub Param<'a>, pub Param<'a>, pub Param<'a>);

pub struct Ta;

impl Ta {
    pub fn open(uuid: &Uuid, params: &Params) -> Result<Session, kmr_common::Error> {
        let raw_handle: *mut TEE_TASessionHandle = Box::into_raw(Box::new(ptr::null_mut()));
        let types = Ta::to_types(&params);
        let parambuf: [TEE_Param; 4] = [params.0.raw, params.1.raw, params.2.raw, params.3.raw];
        let mut origin: u32 = 0;

        // SAFETY: pointers are valid during the function call.
        match unsafe {
            TEE_OpenTASession(
                uuid.as_raw_ptr(),
                TEE_TIMEOUT_INFINITE,
                types,
                parambuf.as_ptr() as _,
                raw_handle,
                &mut origin,
            )
        } {
            TEE_SUCCESS => Ok(Session { handle: raw_handle, origin: origin }),
            code => Err(optee_utee::Error::from_raw_error(code)).map_err(Error::kmerr),
        }
    }

    pub fn open_system_pta(params: &Params) -> Result<Session, kmr_common::Error> {
        let uuid = Uuid::parse_str(PTA_SYSTEM_UUID).map_err(Error::kmerr)?;
        Ta::open(&uuid, params)
    }

    fn to_types(params: &Params) -> u32 {
        (params.0.param_type as u32 & 0xf)
            | ((params.1.param_type as u32 & 0xf) << 4)
            | ((params.2.param_type as u32 & 0xf) << 8)
            | ((params.3.param_type as u32 & 0xf) << 12)
    }
}

pub struct Session {
    handle: *mut TEE_TASessionHandle,
    #[allow(dead_code)]
    pub origin: u32,
}

impl Session {
    pub fn invoke(&self, command: u32, params: &Params) -> Result<u32, kmr_common::Error> {
        let types = Ta::to_types(params);
        let parambuf: [TEE_Param; 4] = [params.0.raw, params.1.raw, params.2.raw, params.3.raw];
        let mut origin: u32 = 0;

        // SAFETY: pointers are valid during the function call.
        match unsafe {
            TEE_InvokeTACommand(
                *self.handle,
                TEE_TIMEOUT_INFINITE,
                command,
                types,
                parambuf.as_ptr() as _,
                &mut origin,
            )
        } {
            TEE_SUCCESS => Ok(origin),
            code => Err(optee_utee::Error::from_raw_error(code)).map_err(Error::kmerr),
        }
    }
}

impl Drop for Session {
    fn drop(&mut self) {
        unsafe {
            if self.handle != ptr::null_mut() {
                TEE_CloseTASession(*self.handle);
            }
            drop(Box::from_raw(self.handle));
        }
    }
}