//! AArch64-specific task context implementation
use core::{arch::global_asm, cell::UnsafeCell, ffi::c_void, fmt, marker::PhantomData};

use kernel_arch_interface::{
    mem::{KernelTableManager, PhysicalMemoryAllocator},
    task::{StackBuilder, TaskContext, TaskFrame, UserContextInfo},
};
use libk_mm_interface::address::PhysicalAddress;
use yggdrasil_abi::{arch::SavedFrame, error::Error};

/// Struct for register values saved when taking an exception
#[repr(C)]
pub struct ExceptionFrame {
    /// General-purpose registers
    pub r: [usize; 32],
    /// SPSR_EL1, userspace flags register
    pub spsr_el1: u64,
    /// ELR_EL1, userspace program counter
    pub elr_el1: u64,
    /// SP_EL0, userspace stack pointer
    pub sp_el0: u64,
    // MDSCR_EL1, debug control
    pub mdscr_el1: u64,
}

#[repr(C, align(0x10))]
struct TaskContextInner {
    // 0x00
    sp: usize,
}

#[repr(align(0x10))]
pub struct FpContext {
    // q0..q31 x 128bit + fpcr/fpsr
    #[allow(unused)]
    inner: [u8; 16 * 32 + 16],
}

/// AArch64 implementation of a task context
#[allow(unused)]
pub struct TaskContextImpl<
    K: KernelTableManager,
    PA: PhysicalMemoryAllocator<Address = PhysicalAddress>,
> {
    inner: UnsafeCell<TaskContextInner>,
    fp_context: UnsafeCell<FpContext>,
    stack_base_phys: PhysicalAddress,
    stack_size: usize,

    _alloc: PhantomData<PA>,
    _table_manager: PhantomData<K>,
}

const COMMON_CONTEXT_SIZE: usize = 8 * 14;

impl FpContext {
    pub const fn new() -> Self {
        Self {
            inner: [0; 16 * 32 + 16],
        }
    }

    /// Stores the FPU context into the `this` pointer.
    ///
    /// # Safety
    ///
    /// It is up to the caller to ensure `this` is a valid pointer to store the FPU context in.
    pub unsafe fn store(this: *mut Self) {
        __aarch64_fp_store_context(this as _)
    }

    /// Loads the FPU with the context stored in `this` pointer.
    ///
    /// # Safety
    ///
    /// It is up to the caller to ensure `this` is a valid pointer to load the FPU context from.
    pub unsafe fn restore(this: *const Self) {
        __aarch64_fp_restore_context(this as _)
    }
}

impl TaskFrame for ExceptionFrame {
    fn store(&self) -> SavedFrame {
        SavedFrame {
            gp_regs: self.r,
            spsr_el1: self.spsr_el1,
            elr_el1: self.elr_el1,
            sp_el0: self.sp_el0,
            mdscr_el1: self.mdscr_el1,
        }
    }

    fn restore(&mut self, saved: &SavedFrame) {
        self.r = saved.gp_regs;
        self.spsr_el1 = saved.spsr_el1;
        self.elr_el1 = saved.elr_el1;
        self.sp_el0 = saved.sp_el0;
        self.mdscr_el1 = saved.mdscr_el1;
    }

    fn argument(&self) -> u64 {
        self.r[0] as _
    }

    fn user_ip(&self) -> usize {
        self.elr_el1 as _
    }

    fn user_sp(&self) -> usize {
        self.sp_el0 as _
    }

    fn set_argument(&mut self, value: u64) {
        self.r[0] = value as _;
    }

    fn set_return_value(&mut self, value: u64) {
        self.r[0] = value as _;
    }

    fn set_user_ip(&mut self, value: usize) {
        self.elr_el1 = value as _;
    }

    fn set_user_sp(&mut self, value: usize) {
        self.sp_el0 = value as _;
    }

    fn set_single_step(&mut self, step: bool) {
        if step {
            self.mdscr_el1 |= 1 << 0;
        } else {
            self.mdscr_el1 &= !(1 << 0);
        }
    }
}

impl fmt::Debug for ExceptionFrame {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        for i in (0..32).step_by(2) {
            write!(
                f,
                "x{:<2} = {:#020x}\tx{:<2} = {:#020x}",
                i,
                self.r[i],
                i + 1,
                self.r[i + 1]
            )?;
            if i != 30 {
                f.write_str("\n")?;
            }
        }

        Ok(())
    }
}

unsafe impl<K: KernelTableManager, PA: PhysicalMemoryAllocator<Address = PhysicalAddress>> Sync
    for TaskContextImpl<K, PA>
{
}

impl<K: KernelTableManager, PA: PhysicalMemoryAllocator<Address = PhysicalAddress>>
    TaskContext<K, PA> for TaskContextImpl<K, PA>
{
    const USER_STACK_EXTRA_ALIGN: usize = 0;
    const SIGNAL_STACK_EXTRA_ALIGN: usize = 0;

    fn kernel(entry: extern "C" fn(usize) -> !, arg: usize) -> Result<Self, Error> {
        const KERNEL_TASK_PAGES: usize = 8;
        let stack_base_phys = PA::allocate_contiguous_pages(KERNEL_TASK_PAGES)?;
        let stack_base = stack_base_phys.raw_virtualize::<K>();

        let mut stack = StackBuilder::new(stack_base, KERNEL_TASK_PAGES * 0x1000);

        // Entry and argument
        stack.push(entry as _);
        stack.push(arg);

        setup_common_context(&mut stack, __aarch64_task_enter_kernel as _, 0, 0);

        let sp = stack.build();

        // TODO stack is leaked

        Ok(Self {
            inner: UnsafeCell::new(TaskContextInner { sp }),
            fp_context: UnsafeCell::new(FpContext::new()),
            stack_base_phys,
            stack_size: KERNEL_TASK_PAGES * 0x1000,

            _alloc: PhantomData,
            _table_manager: PhantomData,
        })
    }

    fn user(context: UserContextInfo) -> Result<Self, Error> {
        const USER_TASK_PAGES: usize = 16;
        let stack_base_phys = PA::allocate_contiguous_pages(USER_TASK_PAGES)?;
        let stack_base = stack_base_phys.raw_virtualize::<K>();

        let mut stack = StackBuilder::new(stack_base, USER_TASK_PAGES * 0x1000);

        let mdscr_el1 = if context.single_step { 1 << 0 } else { 0 };

        stack.push(context.entry);
        stack.push(context.argument);
        stack.push(mdscr_el1);
        stack.push(context.stack_pointer);

        let ttbr0 = context.address_space | (context.asid << 48) | 1;

        setup_common_context(
            &mut stack,
            __aarch64_task_enter_user as _,
            ttbr0,
            context.thread_pointer as _,
        );

        let sp = stack.build();

        Ok(Self {
            inner: UnsafeCell::new(TaskContextInner { sp }),
            fp_context: UnsafeCell::new(FpContext::new()),
            stack_base_phys,
            stack_size: USER_TASK_PAGES * 0x1000,

            _alloc: PhantomData,
            _table_manager: PhantomData,
        })
    }

    unsafe fn enter(&self) -> ! {
        FpContext::restore(self.fp_context.get());

        __aarch64_enter_task(self.inner.get())
    }

    unsafe fn switch(&self, from: &Self) {
        let dst = self.inner.get();
        let src = from.inner.get();

        if dst != src {
            // Save the old context
            FpContext::store(from.fp_context.get());
            // Load next context
            FpContext::restore(self.fp_context.get());

            __aarch64_switch_task(self.inner.get(), from.inner.get())
        }
    }

    unsafe fn switch_and_drop(&self, thread: *const ()) {
        FpContext::restore(self.fp_context.get());

        __aarch64_switch_task_and_drop(self.inner.get(), thread);
    }

    fn set_thread_pointer(&self, _tp: usize) {
        // Do nothing: tp can be set from EL0 by writing to TPIDR_EL0 directly
    }
}

impl<K: KernelTableManager, PA: PhysicalMemoryAllocator<Address = PhysicalAddress>> Drop
    for TaskContextImpl<K, PA>
{
    fn drop(&mut self) {
        assert_eq!(self.stack_size % 0x1000, 0);

        for offset in (0..self.stack_size).step_by(0x1000) {
            unsafe {
                PA::free_page(self.stack_base_phys.add(offset));
            }
        }
    }
}

fn setup_common_context(builder: &mut StackBuilder, entry: usize, ttbr0: u64, tpidr_el0: u64) {
    builder.push(ttbr0 as _); // ttbr0_el1
    builder.push(tpidr_el0 as _); // tpidr_el0

    builder.push(entry); // x30/lr
    builder.push(0); // x29
    builder.push(0); // x28
    builder.push(0); // x27
    builder.push(0); // x26
    builder.push(0); // x25
    builder.push(0); // x24
    builder.push(0); // x23
    builder.push(0); // x22
    builder.push(0); // x21
    builder.push(0); // x20
    builder.push(0); // x19
}

extern "C" {
    fn __aarch64_enter_task(to: *mut TaskContextInner) -> !;
    fn __aarch64_switch_task(to: *mut TaskContextInner, from: *mut TaskContextInner);
    fn __aarch64_switch_task_and_drop(to: *mut TaskContextInner, thread: *const ()) -> !;
    fn __aarch64_task_enter_kernel();
    fn __aarch64_task_enter_user();
    fn __aarch64_fp_store_context(to: *mut c_void);
    fn __aarch64_fp_restore_context(from: *const c_void);
}

global_asm!(include_str!("context.S"), context_size = const COMMON_CONTEXT_SIZE);
