//! Implements the core stack unwinding logic.
//!
//! # TODOs
//!
//! * Support evaluation of DWARF expressions, this might not be required
//!   however, because Rust doesn't tend to use this DWARF feature.

use fallible_iterator::FallibleIterator;
use gimli::{
    CfaRule, EndianSlice, LittleEndian, Register, RegisterRule, StoreOnHeap,
    UnwindContext, UnwindSection, X86_64,
};

use super::eh_info::EhInfo;

/// Implements the core stack unwinding logic by parsing the call frame
/// information. This also stores current (DWARF) register values.
///
/// # Sources
///
/// Taken from [lesenechal.fr](https://lesenechal.fr/en/linux/unwinding-the-stack-the-hard-way)
/// with some additional features to be added soon.
pub struct Unwinder {
    /// The call frame information.
    eh_info: EhInfo,
    /// An [UnwindContext] used by Gimli for optimisations.
    unwind_ctx: UnwindContext<usize, StoreOnHeap>,
    /// The current values of ABI/architecture independent registers. There are
    /// used by DWARF.
    pub regs: RegisterSet,
    /// The current CFA address.
    cfa: u64,
    /// Is this the first iteration?
    is_first: bool,
}

// TODO: Use map_err et al.
impl FallibleIterator for Unwinder {
    type Item = CallFrame;
    type Error = UnwinderError;

    /// Returns call frames of calling functions. This may be called to produce
    /// a stack trace or otherwise support physical unwinding.
    fn next(&mut self) -> Result<Option<Self::Item>, Self::Error> {
        // Gets the current program counter from the DWARF register set.
        let Some(pc) = self.regs.get_pc() else {
            return Err(UnwinderError::NoPcRegister);
        };
        // 0xffffffff8000014b b - 7 = 11 - 7 = 4

        if self.is_first {
            self.is_first = false;

            return Ok(Some(CallFrame { pc, symbol: 0 }));
        }

        // This is a row in the virtual unwind table AKA the CFI which will help
        // us find the CFA (canonical frame address) for a given program
        // counter.
        let Ok(row) = self.eh_info.hdr_table.unwind_info_for_address(
            &self.eh_info.eh_frame,
            &self.eh_info.base_addrs,
            &mut self.unwind_ctx,
            pc,
            |section, bases, offset|
                                // Finds a DWARF CIE using an offset as given.
                                section.cie_from_offset(bases, offset),
        ) else {
            return Err(UnwinderError::NoUnwindInfo);
        };

        // We compute the CFA (canonical frame address from its rule).
        // TODO: Support other rules such as DWARF expressions.
        match row.cfa() {
            CfaRule::RegisterAndOffset { register, offset } => {
                let Some(reg_val) = self.regs.get(*register) else {
                    return Err(UnwinderError::CfaRuleUnknownRegister(
                        *register,
                    ));
                };

                self.cfa = (reg_val as i64 + offset) as u64;
            }

            // TODO: Support other rules for computing the CFA.
            _ => return Err(UnwinderError::UnsupportedCfaRule),
        }

        for reg in RegisterSet::iter() {
            match row.register(reg) {
                RegisterRule::Undefined => self.regs.undef(reg),
                RegisterRule::SameValue => (),
                RegisterRule::Offset(offset) => {
                    // Adds the given offset to the register contents and
                    // retrieve the value from the stack at address CFA +
                    // offset.
                    let ptr = (self.cfa as i64 + offset) as u64 as *const usize;

                    self.regs.set(reg, unsafe { ptr.read() } as u64)?
                }

                _ => {
                    return Err(UnwinderError::UnimplementedRegisterRule);
                }
            }
        }

        // Get the new value for %rip from the function return value and
        // subtract one from it, because the address actually points the the
        // next instruction after the call.
        let Some(pc) = self.regs.get_ret() else {
            return Err(UnwinderError::NoReturnAddr);
        };

        // REVIEWME: Must be a nicer way of doing this.
        let Some(pc) = pc.checked_sub(1) else {
            // REVIEWME: This should handle underflow now.
            return Ok(None);
        };

        self.regs.set_pc(pc);

        // Set %rsp to the CFA. This simulates returning from the function,
        // destroying the call frame, so we were able to virtually unwind (to
        // the caller function).
        self.regs.set_stack_ptr(self.cfa);

        Ok(Some(CallFrame { pc, symbol: 0 }))
    }
}

impl Unwinder {
    pub fn new(eh_info: EhInfo, regset: RegisterSet) -> Self {
        Self {
            eh_info,
            regs: regset,
            unwind_ctx: UnwindContext::new(),
            cfa: 0,
            is_first: true,
        }
    }
}

/// The set of registers used by DWARF. This struct allows future portability if
/// we want to target other architectures than x86_64.
#[derive(Debug, Default)]
pub struct RegisterSet {
    rip: Option<u64>,
    rsp: Option<u64>,
    rbp: Option<u64>,
    /// The return address register.
    ret: Option<u64>,
}

impl RegisterSet {
    pub const fn get(&self, reg: Register) -> Option<u64> {
        match reg {
            X86_64::RSP => self.rsp,
            X86_64::RBP => self.rbp,
            X86_64::RA => self.ret,
            _ => None,
        }
    }

    pub const fn set(
        &mut self,
        reg: Register,
        val: u64,
    ) -> Result<(), UnwinderError> {
        *match reg {
            X86_64::RSP => &mut self.rsp,
            X86_64::RBP => &mut self.rbp,
            X86_64::RA => &mut self.ret,
            _ => return Err(UnwinderError::UnexpectedRegister(reg)),
        } = Some(val);

        Ok(())
    }

    const fn undef(&mut self, reg: Register) {
        *match reg {
            X86_64::RSP => &mut self.rsp,
            X86_64::RBP => &mut self.rbp,
            X86_64::RA => &mut self.ret,
            _ => return,
        } = None;
    }

    const fn get_pc(&self) -> Option<u64> {
        self.rip
    }

    pub const fn set_pc(&mut self, val: u64) {
        self.rip = Some(val);
    }

    const fn get_ret(&self) -> Option<u64> {
        self.ret
    }

    pub const fn set_stack_ptr(&mut self, val: u64) {
        self.rsp = Some(val);
    }

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

/// The current instruction pointer.
#[derive(Debug)]
pub struct CallFrame {
    /// The current instruction pointer.
    pub pc: u64,
    /// The symbol of the function.
    pub symbol: usize,
}

#[derive(Debug)]
/// A list of errors that could occur whilst unwinding the stack.
pub enum UnwinderError {
    UnexpectedRegister(Register),
    UnsupportedCfaRule,
    UnimplementedRegisterRule,
    CfaRuleUnknownRegister(Register),
    NoUnwindInfo,
    NoPcRegister,
    NoReturnAddr,
}