use core::{arch::global_asm, cell::UnsafeCell, marker::PhantomData};

use kernel_arch_interface::{
    mem::{KernelTableManager, PhysicalMemoryAllocator},
    task::{ForkFrame, StackBuilder, TaskContext, TaskFrame, UserContextInfo},
};
use kernel_arch_x86::registers::{FpuContext, CR3, MSR_IA32_FS_BASE};
use libk_mm_interface::address::PhysicalAddress;
use tock_registers::interfaces::Writeable;
use yggdrasil_abi::{arch::SavedFrame, error::Error};

use crate::{
    mem::{auto_lower_address, fixed},
    ArchitectureImpl,
};

/// Frame saved onto the stack when taking an IRQ
#[derive(Debug)]
#[repr(C)]
pub struct IrqFrame {
    pub rax: u64,
    pub rcx: u64,
    pub rdx: u64,
    pub rbx: u64,
    pub rsi: u64,
    pub rdi: u64,
    pub rbp: u64,
    pub r8: u64,
    pub r9: u64,
    pub r10: u64,
    pub r11: u64,
    pub r12: u64,
    pub r13: u64,
    pub r14: u64,
    pub r15: u64,

    pub rip: u64,
    pub cs: u64,
    pub rflags: u64,
    pub rsp: u64,
    pub ss: u64,
}

/// Set of registers saved when taking an exception/interrupt
#[derive(Debug)]
#[repr(C)]
pub struct ExceptionFrame {
    pub rax: u64,
    pub rcx: u64,
    pub rdx: u64,
    pub rbx: u64,
    pub rsi: u64,
    pub rdi: u64,
    pub rbp: u64,
    pub r8: u64,
    pub r9: u64,
    pub r10: u64,
    pub r11: u64,
    pub r12: u64,
    pub r13: u64,
    pub r14: u64,
    pub r15: u64,

    pub exc_number: u64,
    pub exc_code: u64,

    pub rip: u64,
    pub cs: u64,
    pub rflags: u64,
    pub rsp: u64,
    pub ss: u64,
}

/// Set of registers saved when taking a syscall instruction
#[derive(Debug)]
#[repr(C)]
pub struct SyscallFrame {
    pub rax: usize,
    pub args: [usize; 6],

    pub rcx: u64,
    pub r11: u64,

    pub user_ip: u64,
    pub user_sp: u64,
    pub user_flags: u64,

    pub rbx: u64,
    pub rbp: u64,
    pub r12: u64,
    pub r13: u64,
    pub r14: u64,
    pub r15: u64,
}

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

    fs_base: usize,
}

/// x86-64 implementation of a task context
#[allow(dead_code)]
pub struct TaskContextImpl<
    K: KernelTableManager,
    PA: PhysicalMemoryAllocator<Address = PhysicalAddress>,
> {
    inner: UnsafeCell<Inner>,
    fpu_context: UnsafeCell<FpuContext>,
    stack_base_phys: PhysicalAddress,
    stack_size: usize,
    tss_rsp0: usize,

    cr3: usize,

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

impl TaskFrame for IrqFrame {
    fn store(&self) -> SavedFrame {
        SavedFrame {
            rax: self.rax,
            rcx: self.rcx,
            rdx: self.rdx,
            rbx: self.rbx,
            rsi: self.rsi,
            rdi: self.rdi,
            rbp: self.rbp,
            r8: self.r8,
            r9: self.r9,
            r10: self.r10,
            r11: self.r11,
            r12: self.r12,
            r13: self.r13,
            r14: self.r14,
            r15: self.r15,
            user_ip: self.rip,
            user_sp: self.rsp,
            rflags: self.rflags,
        }
    }

    fn restore(&mut self, _saved: &SavedFrame) {
        todo!()
    }

    fn argument(&self) -> u64 {
        self.rdi as _
    }

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

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

    fn set_argument(&mut self, value: u64) {
        self.rdi = value;
    }

    fn set_return_value(&mut self, value: u64) {
        self.rax = value;
    }

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

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

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

impl TaskFrame for ExceptionFrame {
    fn store(&self) -> SavedFrame {
        SavedFrame {
            rax: self.rax,
            rcx: self.rcx,
            rdx: self.rdx,
            rbx: self.rbx,
            rsi: self.rsi,
            rdi: self.rdi,
            rbp: self.rbp,
            r8: self.r8,
            r9: self.r9,
            r10: self.r10,
            r11: self.r11,
            r12: self.r12,
            r13: self.r13,
            r14: self.r14,
            r15: self.r15,
            user_ip: self.rip,
            user_sp: self.rsp,
            rflags: self.rflags,
        }
    }

    fn restore(&mut self, _saved: &SavedFrame) {
        todo!()
    }

    fn argument(&self) -> u64 {
        0
    }

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

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

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

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

    fn set_return_value(&mut self, _value: u64) {
        // Not in syscall, do not overwrite
    }

    fn set_argument(&mut self, value: u64) {
        self.rdi = value;
    }

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

impl<K: KernelTableManager, PA: PhysicalMemoryAllocator<Address = PhysicalAddress>> ForkFrame<K, PA>
    for SyscallFrame
{
    type Context = TaskContextImpl<K, PA>;

    unsafe fn fork(&self, address_space: u64) -> Result<TaskContextImpl<K, PA>, Error> {
        TaskContextImpl::from_syscall_frame(self, address_space)
    }

    fn set_return_value(&mut self, value: u64) {
        self.rax = value as _;
    }
}

impl TaskFrame for SyscallFrame {
    fn store(&self) -> SavedFrame {
        SavedFrame {
            rax: self.rax as _,
            rcx: self.rcx,
            rdx: self.args[2] as _,
            rbx: self.rbx,
            rsi: self.args[1] as _,
            rdi: self.args[0] as _,
            rbp: self.rbp,
            r8: self.args[4] as _,
            r9: self.args[5] as _,
            r10: self.args[3] as _,
            r11: self.r11,
            r12: self.r12,
            r13: self.r13,
            r14: self.r14,
            r15: self.r15,
            user_ip: self.user_ip,
            user_sp: self.user_sp,
            rflags: self.user_flags,
        }
    }

    fn restore(&mut self, saved: &SavedFrame) {
        self.rax = saved.rax as _;
        self.args[0] = saved.rdi as _;
        self.args[1] = saved.rsi as _;
        self.args[2] = saved.rdx as _;
        self.args[3] = saved.r10 as _;
        self.args[4] = saved.r8 as _;
        self.args[5] = saved.r9 as _;

        self.rcx = saved.rcx;
        self.r11 = saved.r11;

        self.user_ip = saved.user_ip;
        self.user_sp = saved.user_sp;
        self.user_flags = saved.rflags;

        self.rbx = saved.rbx;
        self.rbp = saved.rbp;
        self.r12 = saved.r12;
        self.r13 = saved.r13;
        self.r14 = saved.r14;
        self.r15 = saved.r15;
    }

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

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

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

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

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

    fn set_return_value(&mut self, value: u64) {
        self.rax = value as _;
    }

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

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

impl<K: KernelTableManager, PA: PhysicalMemoryAllocator<Address = PhysicalAddress>>
    TaskContextImpl<K, PA>
{
    /// Constructs a new task context from a "forked" syscall frame
    pub(super) unsafe fn from_syscall_frame(
        _frame: &SyscallFrame,
        _cr3: u64,
    ) -> Result<Self, Error> {
        todo!()
        // const USER_TASK_PAGES: usize = 8;

        // 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);

        // // iretq frame
        // stack.push(0x1B);
        // stack.push(frame.user_sp as _);
        // stack.push(0x200);
        // stack.push(0x23);
        // stack.push(frame.user_ip as _);

        // stack.push(frame.args[5] as _); // r9
        // stack.push(frame.args[4] as _); // r8
        // stack.push(frame.args[3] as _); // r10
        // stack.push(frame.args[2] as _); // rdx
        // stack.push(frame.args[1] as _); // rsi
        // stack.push(frame.args[0] as _); // rdi

        // // callee-saved registers
        // stack.push(__x86_64_task_enter_from_fork as _);

        // stack.push(cr3 as _);

        // stack.push(frame.rbp as _);
        // stack.push(0x12345678); // XXX TODO: fs_base from SyscallFrame
        // stack.push(frame.r15 as _);
        // stack.push(frame.r14 as _);
        // stack.push(frame.r13 as _);
        // stack.push(frame.r12 as _);
        // stack.push(frame.rbx as _);

        // let sp = stack.build();
        // let rsp0 = stack_base + USER_TASK_PAGES * 0x1000;

        // Ok(Self {
        //     inner: UnsafeCell::new(Inner { sp, tss_rsp0: rsp0 }),
        //     fpu_context: UnsafeCell::new(FpuContext::new(true)),
        //     stack_base_phys,
        //     stack_size: USER_TASK_PAGES * 0x1000,
        //     _alloc: PhantomData,
        //     _table_manager: PhantomData,
        // })
    }

    unsafe fn store_state(&self) {
        FpuContext::store(self.fpu_context.get());
        // No need to save TSS/%cr3/%fs base back into the TCB, only the kernel
        // can make changes to those
    }

    unsafe fn load_state(&self) {
        FpuContext::restore(self.fpu_context.get());
        // When the task is interrupted from Ring 3, make the CPU load
        // the top of its kernel stack
        ArchitectureImpl::set_local_tss_sp0(self.tss_rsp0);
        MSR_IA32_FS_BASE.set((*self.inner.get()).fs_base as u64);
        CR3.set_address(self.cr3);
    }
}

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 SIGNAL_STACK_EXTRA_ALIGN: usize = 8;
    const USER_STACK_EXTRA_ALIGN: usize = 8;

    fn kernel(entry: extern "C" fn(usize) -> !, arg: usize) -> Result<Self, Error> {
        const KERNEL_TASK_PAGES: usize = 32;

        let cr3: usize = auto_lower_address(&raw const fixed::KERNEL_PML4); // unsafe { KERNEL_TABLES.lock().as_physical_address() }.into();

        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, __x86_64_task_enter_kernel as _);

        let sp = stack.build();

        // TODO stack is leaked

        Ok(Self {
            inner: UnsafeCell::new(Inner { sp, fs_base: 0 }),
            fpu_context: UnsafeCell::new(FpuContext::new(false)),
            stack_base_phys,
            stack_size: KERNEL_TASK_PAGES * 0x1000,

            tss_rsp0: 0,
            cr3,

            _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 mut flags = 0x200;
        if context.single_step {
            flags |= 1 << 8;
        }

        stack.push(flags);
        stack.push(context.entry as _);
        stack.push(context.argument);
        stack.push(context.stack_pointer);

        setup_common_context(&mut stack, __x86_64_task_enter_user as _);

        let sp = stack.build();
        let rsp0 = stack_base + USER_TASK_PAGES * 0x1000;

        Ok(Self {
            inner: UnsafeCell::new(Inner {
                sp,
                fs_base: context.thread_pointer,
            }),
            fpu_context: UnsafeCell::new(FpuContext::new(true)),
            stack_base_phys,
            stack_size: USER_TASK_PAGES * 0x1000,

            tss_rsp0: rsp0,
            cr3: context.address_space as usize,

            _alloc: PhantomData,
            _table_manager: PhantomData,
        })
    }

    unsafe fn enter(&self) -> ! {
        self.load_state();
        __x86_64_enter_task(self.inner.get())
    }

    unsafe fn switch(&self, from: &Self) {
        if core::ptr::addr_eq(self, from) {
            return;
        }

        from.store_state();
        self.load_state();
        __x86_64_switch_task(self.inner.get(), from.inner.get())
    }

    unsafe fn switch_and_drop(&self, thread: *const ()) {
        self.load_state();
        __x86_64_switch_and_drop(self.inner.get(), thread)
    }

    fn set_thread_pointer(&self, tp: usize) {
        unsafe { (*self.inner.get()).fs_base = tp };
        MSR_IA32_FS_BASE.set(tp as _);
    }

    fn align_stack_for_entry(sp: usize) -> usize {
        (sp & !0xF) - 8
    }
}

impl<K: KernelTableManager, PA: PhysicalMemoryAllocator<Address = PhysicalAddress>> Drop
    for TaskContextImpl<K, PA>
{
    fn drop(&mut self) {
        log::trace!("Drop Context {:#p}", 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) {
    builder.push(entry);

    builder.push(0); // %rbp
    builder.push(0); // %r15
    builder.push(0); // %r14
    builder.push(0); // %r13
    builder.push(0); // %r12
    builder.push(0); // %rbx
}

extern "C" {
    fn __x86_64_task_enter_kernel();
    fn __x86_64_task_enter_user();
    fn __x86_64_task_enter_from_fork();
    fn __x86_64_enter_task(to: *mut Inner) -> !;
    fn __x86_64_switch_task(to: *mut Inner, from: *mut Inner);
    fn __x86_64_switch_and_drop(to: *mut Inner, from: *const ());
}

global_asm!(include_str!("context.S"), options(att_syntax));
