/*
 * Copyright (c) 2025 Google Inc. All rights reserved
 *
 * Permission is hereby granted, free of charge, to any person obtaining
 * a copy of this software and associated documentation files
 * (the "Software"), to deal in the Software without restriction,
 * including without limitation the rights to use, copy, modify, merge,
 * publish, distribute, sublicense, and/or sell copies of the Software,
 * and to permit persons to whom the Software is furnished to do so,
 * subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be
 * included in all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
 * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
 * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
 * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
 * CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
 * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
 * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */

#![no_std]

use crate::sys::sm_vm_notifier_register;
use crate::sys::sm_vm_notifier_unregister;
use alloc::boxed::Box;
use core::cell::UnsafeCell;
use core::fmt;
use core::fmt::Debug;
use core::fmt::Formatter;
use core::mem::zeroed;
use core::mem::ManuallyDrop;
use core::ptr::null_mut;
use extmem::ext_mem_obj_id_t;
use rust_support::container_of;
use rust_support::list_node;
use rust_support::obj_ref;
use rust_support::status_t;
use rust_support::Error;

pub use crate::sys::sm_vm;
pub use crate::sys::sm_vm_del_ref;
pub use crate::sys::sm_vm_get;
pub use crate::sys::sm_vm_notifier;

mod sys {
    #![allow(unused)]
    #![allow(non_camel_case_types)]
    use extmem::ext_mem_obj_id_t;
    use rust_support::list_node;
    use rust_support::obj_ref;
    use rust_support::status_t;
    include!(env!("BINDGEN_INC_FILE"));
}

pub fn intc_raise_doorbell_irq() {
    // SAFETY: This function has no safety requirements and may be called from any context.
    unsafe {
        sys::sm_intc_raise_doorbell_irq();
    }
}

// Use repr(C) to ensure its layout is consistent between compilations to allow using
// container_of to cast from sm_vm_notifier pointers to VmNotifierInner pointers.
#[repr(C)]
struct VmNotifierInner {
    notif: sm_vm_notifier,
    // A rust-specific callback for VM destruction notification. The callback field in
    // sm_vm_notifier just reads the client ID and calls this function with it.
    destroy: VmDestroyCallback,
}

// Implement Debug manually since sm_vm_notifier contains a list_node which does not implement it.
// This impl omits the callback field in sm_vm_notifier since the rust bindings ensure it will
// always be the invoke_rust_callback function.
impl Debug for VmNotifierInner {
    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), fmt::Error> {
        f.debug_struct("VmNotifierInner")
            .field("client_id", &self.notif.client_id)
            .field("destroy", &self.destroy)
            .finish()
    }
}

// A notifier for invoking a callback when a VM is destroyed. This is boxed since sm_vm_notifier
// has a list_node field which must not move.
#[derive(Debug)]
pub struct VmNotifier(Box<VmNotifierInner>);

// SAFETY: VmNotifier provides no way to mutate it through a shared reference
unsafe impl Sync for VmNotifier {}

// SAFETY: VmNotifier only has a heap allocated field so it may be freely sent across threads
// without invalidating it.
unsafe impl Send for VmNotifier {}

// TODO: add Unpin as a negative trait bound once the rustc feature is stabilized.
// impl !Unpin for VmNotifier {}

pub type VmDestroyCallback = fn(ext_mem_obj_id_t) -> Result<(), Error>;

impl VmNotifier {
    /// Creates a destruction notifier for a VM with the given client ID.
    ///
    /// Unlike the C sm_vm_notifier interface, the destroy callback takes the client ID as its only
    /// argument and must be a safe function. Callers must also register and avoid dropping the
    /// VmNotifier to ensure Trusty will invoke it when the VM is destroyed.
    pub fn new(client_id: ext_mem_obj_id_t, destroy: VmDestroyCallback) -> Result<Self, Error> {
        // SAFETY: The c_notif argument must point to a sm_vm_notifier embedded in a VmNotifierInner
        // struct. The caller must ensure that the `notif.client_id` and `destroy` fields will not change
        // for the duration of this function.
        unsafe extern "C" fn invoke_rust_callback(c_notif: *mut sm_vm_notifier) -> status_t {
            // SAFETY: Safety delegated to caller.
            let client_id = unsafe { (*c_notif).client_id };
            // This cast is valid due to the layout of VmNotifierInner.
            // SAFETY: invoke_rust_callback's safety requirements ensure `c_notif` points to the
            // `notif` field in VmNotifierInner.
            let rust_notif = unsafe { container_of!(c_notif, VmNotifierInner, notif) };
            // SAFETY: Safety delegated to caller.
            let rust_callback = unsafe { (*rust_notif).destroy };
            let rc = rust_callback(client_id);
            match rc {
                Ok(()) => 0,
                Err(e) => status_t::from(e),
            }
        }
        let notif = sm_vm_notifier {
            node: list_node::default(),
            client_id,
            destroy: Some(invoke_rust_callback),
        };
        let rust_notif = Box::new(VmNotifierInner { notif, destroy });
        // TODO: Call sm_vm_notifier_init and assert it returns 0. That function currently only sets
        // fields and validates arguments like VmNotifier::new but we should use it instead of
        // relying on setting fields explicitly to avoid having to keep these two code paths in
        // sync. Using it currently triggers a linker error since it relies on bindgen's
        // --wrap-static-fns and we need to ensure the .o produced from that source file gets linked
        // into lk-crates.a
        // let rc = sm_vm_notifier_init(rust_notif.notif.get(), client_id, Some(invoke_rust_callback));
        // assert!(rc == 0);
        let mut res = Self(rust_notif);
        res.register()?;
        Ok(res)
    }

    fn register(&mut self) -> Result<(), Error> {
        // SAFETY: This makes Trusty call the C callback `invoke_rust_callback` when a VM is
        // destroyed. The safety requirements on it are that the argument is an sm_vm_notifier
        // embedded within a VmNotifierInner and that the `client_id` and `destroy` (rust
        // callback) fields do not change for the duration of that function. The former holds
        // because of the definition of the VmNotifier type. The latter two hold because VmNotifier
        // provides no way to change these fields after its initialized. The constructor also
        // ensures that `invoke_rust_callback` is always the unsafe function that will get called
        // when the VM is destroyed.
        let rc = unsafe { sm_vm_notifier_register(&raw mut self.0.notif) };
        Error::from_lk(rc)
    }

    pub fn unregister(self) -> Result<(), (Self, Error)> {
        // Skip dropping self to avoid calling sm_vm_notifier_unregister multiple times
        let mut rust_notif = ManuallyDrop::new(self);
        // SAFETY: This function removes the notifier its argument points to from the bst of VMs
        // preventing the C callback from being called. `unregister` takes `self` by value so only a
        // single thread may attempt to remove the notifier from the bst and there cannot be a race
        // condition. If the callback is already in progress this function will just wait for it to
        // finish and return NO_ERROR. If the callback has already been called it will just return
        // ERR_NOT_FOUND.
        let rc = unsafe { sm_vm_notifier_unregister(&raw mut rust_notif.0.notif) };
        if rc != 0 {
            return Error::from_lk(rc).map_err(|e| (ManuallyDrop::into_inner(rust_notif), e));
        }
        Ok(())
    }
}

impl Drop for VmNotifier {
    fn drop(&mut self) {
        // SAFETY: This function removes the notifier its argument points to from the bst of VMs
        // preventing the C callback from being called. Only a single thread may drop the notifier
        // so only a single thread may attempt to remove the notifier from the bst and there cannot
        // be a race condition. If the callback is already in progress it will just wait for it to
        // finish and return NO_ERROR. If the callback has already been called it will return
        // ERR_NOT_FOUND. Any other error code is unexpected so we panic since Drop doesn't allow
        // bubbling up an error to the caller.
        let rc = unsafe { sm_vm_notifier_unregister(&raw mut self.0.notif) };
        if rc != Error::NO_ERROR.0 && rc != Error::ERR_NOT_FOUND.0 {
            panic!("Failed to unregister VmNotifier");
        }
    }
}

#[derive(Debug)]
pub struct VmRef {
    vm_ptr: *mut sm_vm,
    // TODO: Remove the UnsafeCell when the box_as_ptr feature is stabilized
    // https://github.com/rust-lang/rust/issues/129090
    obj_ref_ptr: Box<UnsafeCell<obj_ref>>,
}

impl VmRef {
    pub fn new(client_id: ext_mem_obj_id_t) -> Result<Self, Error> {
        // sm_vm_get expects the caller to allocate memory for the struct obj_ref so we place it on
        // the heap to ensure it remains valid after VmRef::new returns
        // SAFETY: obj_ref is a C type with two pointers which can be zeroed.
        let obj_ref_init: obj_ref = unsafe { zeroed() };
        let boxed_obj_ref = Box::new(UnsafeCell::new(obj_ref_init));
        let mut vm_ptr: *mut sm_vm = null_mut();
        // Grab a reference to the VM by calling sm_vm_get
        // SAFETY: This function takes a non-null pointer to a struct obj_ref which may be
        // uninitialized and a pointer for the VM object. The struct obj_ref pointee has a fixed
        // address so the value it's initialized to remains valid after this function returns. The
        // pointer to the obj_ref that sm_vm_get stores also remains valid since it's boxed and
        // owned by the returned Self
        let rc = unsafe { sm_vm_get(client_id, boxed_obj_ref.get(), &raw mut vm_ptr) };
        Error::from_lk(rc)?;
        // If it succeeded store the out-parameters for when we want to release the VM reference
        Ok(Self { vm_ptr, obj_ref_ptr: boxed_obj_ref })
    }
}

impl Drop for VmRef {
    fn drop(&mut self) {
        // Release the reference
        // SAFETY: These pointers were initialized to valid values in VmRef::new and their pointees
        // remain valid for the lifetime of the VmRef since obj_ref is owned by the VmRef and libsm
        // will not deallocate the backing sm_vm struct while there are references to the VM.
        unsafe { sm_vm_del_ref(self.vm_ptr, self.obj_ref_ptr.get()) }
    }
}

// SAFETY: Both pointers in VmRef are heap allocated so it may freely be sent between threads. The
// thread which releases a refcount does not necessarily need to be same that grabbed it so its Drop
// impl is allowed to run in a different thread than the one that created it.
unsafe impl Send for VmRef {}