//! psci 固件

use core::{arch::global_asm, mem::offset_of};

use moa_bsp_define::smp::{CpuBootMethod, PsciCompatible};

use crate::{
    bsp::__smp_define,
    println,
    processor::{PROCESSOR_CPU_ON, PROCESSOR_SYS_POWEROFF, PROCESSOR_SYS_RESTART},
};

const PSCI_0_2_FN_BASE: usize = 0x84000000;

const fn psci_0_2_fn(n: usize) -> usize {
    PSCI_0_2_FN_BASE + n
}

const PSCI_0_2_64BIT: usize = 0x40000000;
const PSCI_0_2_FN64_BASE: usize = PSCI_0_2_FN_BASE + PSCI_0_2_64BIT;

const fn psci_0_2_fn64(n: usize) -> usize {
    PSCI_0_2_FN64_BASE + n
}

const PSCI_0_2_FN_PSCI_VERSION: usize = psci_0_2_fn(0);
const PSCI_0_2_FN64_CPU_ON: usize = psci_0_2_fn64(3);
const PSCI_0_2_FN_SYSTEM_OFF: usize = psci_0_2_fn(8);
const PSCI_0_2_FN_SYSTEM_RESET: usize = psci_0_2_fn(9);

const PSCI_VERSION_MAJOR_SHIFT: usize = 16;
const PSCI_VERSION_MINOR_MASK: usize = (1 << PSCI_VERSION_MAJOR_SHIFT) - 1;
const PSCI_VERSION_MAJOR_MASK: usize = !PSCI_VERSION_MINOR_MASK;

const fn psci_version_major(ver: usize) -> usize {
    (ver & PSCI_VERSION_MAJOR_MASK) >> PSCI_VERSION_MAJOR_SHIFT
}

const fn psci_version_minor(ver: usize) -> usize {
    ver & PSCI_VERSION_MINOR_MASK
}

global_asm!(r"
ARM_SMCCC_RES_X0_OFFS = {}
ARM_SMCCC_RES_X2_OFFS = {}

    .macro SMCCC instr
    .cfi_startproc
    \instr	#0
    ldr	x4, [sp]
    stp	x0, x1, [x4, #ARM_SMCCC_RES_X0_OFFS]
    stp	x2, x3, [x4, #ARM_SMCCC_RES_X2_OFFS]
    ret
    .cfi_endproc
    .endm

.global __arm_smccc_smc; .align 2; __arm_smccc_smc:
    SMCCC   smc
.type __arm_smccc_smc, @function; .size __arm_smccc_smc, .-__arm_smccc_smc

.global __arm_smccc_hvc; .align 2; __arm_smccc_hvc:
    SMCCC   hvc
.type __arm_smccc_hvc, @function; .size __arm_smccc_hvc, .-__arm_smccc_hvc
", const ARM_SMCCC_RES_X0_OFFS, const ARM_SMCCC_RES_X2_OFFS);

unsafe extern "C" {
    fn __arm_smccc_smc(
        a0: usize,
        a1: usize,
        a2: usize,
        a3: usize,
        a4: usize,
        a5: usize,
        a6: usize,
        a7: usize,
        res: &mut ArmSmcccRes,
        a9: usize,
    );
    fn __arm_smccc_hvc(
        a0: usize,
        a1: usize,
        a2: usize,
        a3: usize,
        a4: usize,
        a5: usize,
        a6: usize,
        a7: usize,
        res: &mut ArmSmcccRes,
        a9: usize,
    );
}

#[repr(C)]
struct ArmSmcccRes {
    a0: usize,
    a1: usize,
    a2: usize,
    a3: usize,
}

const ARM_SMCCC_RES_X0_OFFS: usize = offset_of!(ArmSmcccRes, a0);
const ARM_SMCCC_RES_X2_OFFS: usize = offset_of!(ArmSmcccRes, a2);

static mut PSCI_FN_CPU_ON_ID: usize = 0;

fn psci_cpu_on(cpuid: usize, entry_point: usize) -> bool {
    let ret = unsafe { INVOKE_PSCI_FN(PSCI_FN_CPU_ON_ID, cpuid, entry_point, 0) };
    ret == 0
}

fn psci_sys_reset() {
    unsafe {
        INVOKE_PSCI_FN(PSCI_0_2_FN_SYSTEM_RESET, 0, 0, 0);
    }
}

fn psci_sys_poweroff() {
    unsafe {
        INVOKE_PSCI_FN(PSCI_0_2_FN_SYSTEM_OFF, 0, 0, 0);
    }
}

fn psci_0_2_set_functions() {
    println!("Using standard PSCI v0.2 function IDs");
    unsafe {
        PSCI_FN_CPU_ON_ID = PSCI_0_2_FN64_CPU_ON;
        PROCESSOR_CPU_ON = psci_cpu_on;

        PROCESSOR_SYS_RESTART = psci_sys_reset;
        PROCESSOR_SYS_POWEROFF = psci_sys_poweroff;
    }
}

fn psci_probe() {
    let ver = unsafe { INVOKE_PSCI_FN(PSCI_0_2_FN_PSCI_VERSION, 0, 0, 0) };
    println!("PSCIv{}.{} detected in firmware.", psci_version_major(ver), psci_version_minor(ver));
    assert!(!(psci_version_major(ver) == 0 && psci_version_minor(ver) < 2));

    psci_0_2_set_functions();
}

fn __invoke_psci_fn_hvc(function_id: usize, arg0: usize, arg1: usize, arg2: usize) -> usize {
    let mut res = ArmSmcccRes { a0: 0, a1: 0, a2: 0, a3: 0 };
    unsafe {
        __arm_smccc_hvc(function_id, arg0, arg1, arg2, 0, 0, 0, 0, &mut res, 0);
    }
    res.a0
}

fn __invoke_psci_fn_smc(function_id: usize, arg0: usize, arg1: usize, arg2: usize) -> usize {
    let mut res = ArmSmcccRes { a0: 0, a1: 0, a2: 0, a3: 0 };
    unsafe {
        __arm_smccc_smc(function_id, arg0, arg1, arg2, 0, 0, 0, 0, &mut res, 0);
    }
    res.a0
}

static mut INVOKE_PSCI_FN: fn(usize, usize, usize, usize) -> usize = __invoke_psci_fn_hvc;

fn set_conduit(method: moa_bsp_define::smp::PsciEnableMethod) {
    unsafe {
        if matches!(method, moa_bsp_define::smp::PsciEnableMethod::Smc) {
            INVOKE_PSCI_FN = __invoke_psci_fn_smc;
        }
    }
}

fn psci_0_2_init(enable_method: moa_bsp_define::smp::PsciEnableMethod) {
    set_conduit(enable_method);
    psci_probe();
}

fn psci_0_1_init(enable_method: moa_bsp_define::smp::PsciEnableMethod, cpu_on: usize) {
    set_conduit(enable_method);

    println!("Using PSCI v0.1 Function IDs");

    unsafe {
        PSCI_FN_CPU_ON_ID = cpu_on;
        PROCESSOR_CPU_ON = psci_cpu_on;
    }
}

pub(crate) fn psci_init() {
    let boot_method = unsafe { __smp_define().boot_method() };

    match boot_method {
        CpuBootMethod::Psci { compatible, enable_method, cpu_on } => match compatible {
            PsciCompatible::Psci0_1 => {
                psci_0_1_init(enable_method, cpu_on);
            },
            _ => {
                psci_0_2_init(enable_method);
            },
        },
        CpuBootMethod::Sbi => {},
    }
}
