use crate::Error;
use core::marker::PhantomData;

pub const DAIF_MASK: u64 = 0x3c0;
pub const IRQ_MASK: u64 = 0x80;

/// Privileged, CPU-local register access. No method may migrate the caller.
///
/// # Safety
/// The implementation must report the actual logical CPU and DAIF, and write
/// only that CPU's DAIF. It must exclusively own register changes made through
/// this object. An implementation must not re-enter the guard while writing.
pub unsafe trait Backend {
    fn cpu(&self) -> u16;
    fn read_daif(&self) -> u64;
    fn write_daif(&mut self, value: u64);
}

pub struct IrqState<B: Backend> {
    backend: B,
    cpu: u16,
    depth: u32,
    faulted: bool,
    _local: PhantomData<*mut ()>,
}

impl<B: Backend> IrqState<B> {
    pub fn new(backend: B) -> Result<Self, (Error, B)> {
        if backend.cpu() == u16::MAX || backend.read_daif() & !DAIF_MASK != 0 {
            return Err((Error::Invalid, backend));
        }
        let cpu = backend.cpu();
        Ok(Self {
            backend,
            cpu,
            depth: 0,
            faulted: false,
            _local: PhantomData,
        })
    }

    pub fn enter(&mut self) -> Result<Guard<'_, B>, Error> {
        if self.faulted {
            return Err(Error::Quarantined);
        }
        if self.backend.cpu() != self.cpu {
            return Err(Error::Foreign);
        }
        let next = self.depth.wrapping_add(1);
        let saved = self.backend.read_daif();
        if saved & !DAIF_MASK != 0 {
            return Err(Error::Invalid);
        }
        self.backend.write_daif(saved | IRQ_MASK);
        if self.backend.cpu() != self.cpu || self.backend.read_daif() != saved | IRQ_MASK {
            self.faulted = true;
            return Err(Error::Backend);
        }
        self.depth = next;
        Ok(Guard {
            state: self,
            saved,
            active: true,
            _local: PhantomData,
        })
    }

    pub fn observation(&self) -> (u16, u64, u32, bool) {
        (
            self.backend.cpu(),
            self.backend.read_daif(),
            self.depth,
            self.faulted,
        )
    }

    pub fn into_backend(self) -> Result<B, (Error, Self)> {
        if self.faulted || self.depth != 0 {
            Err((Error::Quarantined, self))
        } else {
            Ok(self.backend)
        }
    }
}

pub struct Guard<'a, B: Backend> {
    state: &'a mut IrqState<B>,
    saved: u64,
    active: bool,
    _local: PhantomData<*mut ()>,
}

impl<B: Backend> Guard<'_, B> {
    pub fn enter(&mut self) -> Result<Guard<'_, B>, Error> {
        self.check_masked()?;
        self.state.enter()
    }

    pub fn cpu(&self) -> u16 {
        self.state.cpu
    }
    pub fn saved_daif(&self) -> u64 {
        self.saved
    }
    pub fn observation(&self) -> (u16, u64, u32, bool) {
        self.state.observation()
    }

    pub fn check_masked(&self) -> Result<(), Error> {
        if self.state.faulted {
            return Err(Error::Quarantined);
        }
        if self.state.backend.cpu() != self.state.cpu {
            return Err(Error::Foreign);
        }
        let daif = self.state.backend.read_daif();
        if daif & !DAIF_MASK != 0 || daif & IRQ_MASK == 0 {
            return Err(Error::Context);
        }
        Ok(())
    }

    pub fn finish(mut self) -> Result<(), Error> {
        self.restore()
    }

    fn restore(&mut self) -> Result<(), Error> {
        if !self.active {
            return Ok(());
        }
        self.active = false;
        if self.check_masked().is_err() || self.state.depth == 0 {
            self.state.faulted = true;
            return Err(Error::Quarantined);
        }
        self.state.backend.write_daif(self.saved);
        if self.state.backend.cpu() != self.state.cpu
            || self.state.backend.read_daif() != self.saved
        {
            self.state.faulted = true;
            return Err(Error::Backend);
        }
        self.state.depth -= 1;
        Ok(())
    }
}

impl<B: Backend> Drop for Guard<'_, B> {
    fn drop(&mut self) {
        let _ = self.restore();
    }
}

#[cfg(test)]
mod exhaustion_tests {
    use super::*;
    struct Registers(u64);
    unsafe impl Backend for Registers {
        fn cpu(&self) -> u16 {
            0
        }
        fn read_daif(&self) -> u64 {
            self.0
        }
        fn write_daif(&mut self, value: u64) {
            self.0 = value
        }
    }
    #[test]
    fn depth_exhaustion_never_writes_registers() {
        let mut state = IrqState::new(Registers(0x140)).ok().unwrap();
        state.depth = u32::MAX;
        assert!(matches!(state.enter(), Err(Error::Overflow)));
        assert_eq!(state.observation(), (0, 0x140, u32::MAX, false));
    }
}