//! ARM PSCI driver implementation

use abi::error::Error;
use alloc::sync::Arc;
use device_api::{
    device::{Device, DeviceInitContext},
    CpuBringupDevice, ResetDevice,
};
use device_tree::{
    driver::{device_tree_driver, Node, ProbeContext},
    DeviceTreePropertyRead,
};
use kernel_arch::{Architecture, ArchitectureImpl};

use crate::arch::PLATFORM;

enum CallMethod {
    Hvc,
    Smc,
}

/// ARM Power State Coordination Interface driver
pub struct Psci {
    method: CallMethod,
    cpu_on: u32,
    #[allow(dead_code)]
    cpu_off: u32,
    #[allow(dead_code)]
    cpu_suspend: u32,
}

impl Device for Psci {
    fn display_name(&self) -> &str {
        "ARM PSCI"
    }

    unsafe fn init(self: Arc<Self>, _cx: DeviceInitContext) -> Result<(), Error> {
        PLATFORM.psci.init(self.clone());
        Ok(())
    }
}

impl CpuBringupDevice for Psci {
    unsafe fn start_cpu(&self, id: usize, ip: usize, arg0: usize) -> Result<(), Error> {
        self.call(self.cpu_on as _, id as _, ip as _, arg0 as _);
        Ok(())
    }
}

impl ResetDevice for Psci {
    unsafe fn reset(&self) -> ! {
        ArchitectureImpl::set_interrupt_mask(true);

        self.call(Self::SYSTEM_RESET as _, 0, 0, 0);

        loop {
            ArchitectureImpl::wait_for_interrupt();
        }
    }
}

impl Psci {
    const SYSTEM_RESET: u32 = 0x84000009;
    const SYSTEM_OFF: u32 = 0x84000008;

    #[inline]
    unsafe fn call(&self, mut x0: u64, x1: u64, x2: u64, x3: u64) -> u64 {
        match self.method {
            CallMethod::Hvc => {
                core::arch::asm!("hvc #0", inlateout("x0") x0, in("x1") x1, in("x2") x2, in("x3") x3)
            }
            CallMethod::Smc => {
                core::arch::asm!("smc #0", inlateout("x0") x0, in("x1") x1, in("x2") x2, in("x3") x3)
            }
        }
        x0
    }

    /// Shut down the system
    pub unsafe fn power_off(&self) -> Result<!, Error> {
        self.call(Self::SYSTEM_OFF as _, 0, 0, 0);
        unreachable!()
    }
}

device_tree_driver! {
    compatible: ["arm,psci-1.0", "arm,psci"],
    driver: {
        fn probe(&self, node: &Arc<Node>, _context: &mut ProbeContext) -> Option<Arc<dyn Device>> {
            let method = node.property("method")?;
            let method = match method.as_str()? {
                "hvc" => CallMethod::Hvc,
                "smc" => CallMethod::Smc,
                method => {
                    log::warn!("Ignoring PSCI with unknown method: {method:?}");
                    return None;
                }
            };
            let cpu_on = node.property("cpu_on")?.read_cell(0, 1)? as u32;
            let cpu_off = node.property("cpu_off")?.read_cell(0, 1)? as u32;
            let cpu_suspend = node.property("cpu_suspend")?.read_cell(0, 1)? as u32;

            Some(Arc::new(Psci {
                method,
                cpu_on,
                cpu_off,
                cpu_suspend
            }))
        }
    }
}
