use riscv::register::sstatus;

/// This struct should be stored in your per-Hart data 
/// (e.g., accessed via the `tp` register or an array indexed by `mhartid`).
#[derive(Clone, Copy)]
pub struct HartInterruptState {
    /// Depth of nested spinlocks.
    pub noff: usize,
    /// Were interrupts enabled before the outermost lock was acquired?
    pub interrupt_enabled_before: bool,
}

impl HartInterruptState {
    pub const fn new() -> Self {
        Self {
            noff: 0,
            interrupt_enabled_before: false,
        }
    }
}

/// Disables interrupts and increments the nesting count.
pub fn push_off() {
    let old_sie = sstatus::read().sie();
    unsafe { sstatus::clear_sie() };
    
    // TODO: Fetch your current hart's state. 
    // Example: let hart = my_cpu();
    let hart = get_current_hart_state_mut(); 
    
    // If this is the very first lock being acquired, record the original SIE state.
    if hart.noff == 0 {
        hart.interrupt_enabled_before = old_sie;
    }
    
    hart.noff += 1;
}

/// Decrements the nesting count and restores interrupts if count reaches 0.
pub fn pop_off() {
    // TODO: Fetch your current hart's state.
    let hart = get_current_hart_state_mut();
    
    // Panic if we are popping more than we pushed (lock bug).
    assert!(hart.noff > 0, "pop_off: nesting count already 0");
    
    hart.noff -= 1;
    
    // If we just released the outermost lock, AND interrupts were originally on, restore them.
    if hart.noff == 0 && hart.interrupt_enabled_before {
        unsafe { sstatus::set_sie() };
    }
}

static mut HART_STATE: [HartInterruptState; 4] = [HartInterruptState::new(); 4];

// Dummy function to represent getting your per-CPU state.
// In a real kernel, this often reads the `tp` register and casts it to a pointer.
#[cfg(target_arch = "riscv64")]
fn get_current_hart_state_mut() -> &'static mut HartInterruptState {
    let id: usize;
    unsafe { core::arch::asm!("mv {}, tp", out(reg) id) };
    unsafe { &mut HART_STATE[id] }
}

// Dummy function to represent getting your per-CPU state.
// In a real kernel, this often reads the `tp` register and casts it to a pointer.
#[cfg(not(target_arch = "riscv64"))]
fn get_current_hart_state_mut() -> &'static mut HartInterruptState {
    let _ = &raw mut HART_STATE;
    unimplemented!("Only supported on riscv64")
}