use core::{
    mem::MaybeUninit,
    ops::{Index, IndexMut},
};

use gimli::{LittleEndian, X86_64};

pub type HostEndian = LittleEndian;

pub fn host_endian() -> HostEndian {
    LittleEndian
}

#[derive(Default)]
pub struct RegisterSet {
    pub rip: Option<u64>,
    pub rsp: Option<u64>,
    pub rbp: Option<u64>,
    pub ret: Option<u64>,
}

#[derive(Debug)]
#[repr(C)]
struct UnwindEntryInfo {
    rip: usize,
    rbp: usize,
    rsp: usize,
}

impl RegisterSet {
    pub fn current() -> Self {
        let mut cx = MaybeUninit::uninit();
        let cx = unsafe { get_context(&mut cx) };

        Self {
            rip: Some(cx.rip as _),
            rsp: Some(cx.rsp as _),
            rbp: Some(cx.rbp as _),
            ret: None,
        }
    }

    pub fn iter() -> impl Iterator<Item = gimli::Register> {
        [X86_64::RA, X86_64::RBP, X86_64::RSP].into_iter()
    }
}

impl Index<gimli::Register> for RegisterSet {
    type Output = Option<u64>;

    fn index(&self, index: gimli::Register) -> &Self::Output {
        match index {
            X86_64::RA => &self.ret,
            X86_64::RSP => &self.rsp,
            X86_64::RBP => &self.rbp,
            _ => panic!("Unhandled register"),
        }
    }
}

impl IndexMut<gimli::Register> for RegisterSet {
    fn index_mut(&mut self, index: gimli::Register) -> &mut Self::Output {
        match index {
            X86_64::RA => &mut self.ret,
            X86_64::RSP => &mut self.rsp,
            X86_64::RBP => &mut self.rbp,
            _ => panic!("Unhandled register"),
        }
    }
}

#[naked]
unsafe extern "C" fn get_context(unw: &mut MaybeUninit<UnwindEntryInfo>) -> &UnwindEntryInfo {
    // %rsp: return address (%rip)
    // %rdi: *mut UnwindInfo unw
    // %rax clobbered
    core::arch::naked_asm!(
        r#"
        // unw.rip = [%rsp]
        mov (%rsp), %rax
        mov %rax, 0(%rdi)
        // unw.rsp = %rsp + 8
        mov %rsp, %rax
        add $8, %rax
        mov %rax, 16(%rdi)
        // unw.rbp = %rbp
        mov %rbp, 8(%rdi)
        mov %rdi, %rax
        ret
    "#,
        options(att_syntax)
    )
}