use {
    crate::utils::{
        copyhashmap::CopyHashMap,
        mmap::{Mmapped, mmap},
        oserror::OsError,
        page_size::page_size,
        ptr_ext::{MutPtrExt, PtrExt},
    },
    std::{marker::PhantomData, ops::Range, rc::Rc},
    thiserror::Error,
    uapi::{OwnedFd, Pod, c},
};

#[derive(Default)]
pub struct PwMemPool {
    pub mems: CopyHashMap<u32, Rc<PwMem>>,
}

#[derive(Debug, Copy, Clone, Eq, PartialEq)]
pub enum PwMemType {
    MemFd,
    DmaBuf,
}

pub struct PwMem {
    pub _ty: PwMemType,
    pub read: bool,
    pub write: bool,
    pub fd: Rc<OwnedFd>,
}

pub struct PwMemMap {
    pub _mem: Rc<PwMem>,
    pub range: Range<usize>,
    pub map: Mmapped,
}

pub struct PwMemTyped<T> {
    mem: Rc<PwMemMap>,
    offset: usize,
    _phantom: PhantomData<T>,
}

#[expect(dead_code)]
pub struct PwMemSlice {
    mem: Rc<PwMemMap>,
    range: Range<usize>,
}

impl PwMemPool {
    pub fn map(&self, memid: u32, offset: u32, size: u32) -> Result<Rc<PwMemMap>, PwMemError> {
        match self.mems.get(&memid) {
            Some(m) => m.map(offset, size),
            _ => Err(PwMemError::MemidDoesNotExist(memid)),
        }
    }
}

impl PwMem {
    pub fn map(self: &Rc<Self>, offset: u32, size: u32) -> Result<Rc<PwMemMap>, PwMemError> {
        let mask = page_size() - 1;
        let offset = offset as usize;
        let size = size as usize;
        let start = offset & !mask;
        let dist = offset - start;
        let len = (size + dist + mask) & !mask;
        let range = dist..dist + size;
        let mut prot = 0;
        if self.read {
            prot |= c::PROT_READ;
        }
        if self.write {
            prot |= c::PROT_WRITE;
        }
        let map = match mmap(len as _, prot, c::MAP_SHARED, self.fd.raw(), start as _) {
            Ok(m) => m,
            Err(e) => return Err(PwMemError::MmapFailed(e)),
        };
        Ok(Rc::new(PwMemMap {
            _mem: self.clone(),
            range,
            map,
        }))
    }
}

impl PwMemMap {
    #[expect(dead_code)]
    pub unsafe fn read<T: Pod>(&self) -> &T {
        self.check::<T>(0);
        unsafe { (self.map.ptr.cast::<u8>().add(self.range.start) as *const T).deref() }
    }

    #[expect(dead_code)]
    pub unsafe fn write<T: Pod>(&self) -> &mut T {
        self.check::<T>(0);
        unsafe { (self.map.ptr.cast::<u8>().add(self.range.start) as *mut T).deref_mut() }
    }

    #[expect(dead_code)]
    pub unsafe fn bytes_mut(&self) -> &mut [u8] {
        unsafe {
            std::slice::from_raw_parts_mut(
                self.map.ptr.cast::<u8>().add(self.range.start) as _,
                self.range.len(),
            )
        }
    }

    fn check<T>(&self, offset: usize) {
        assert!(offset <= self.range.len());
        assert!(size_of::<T>() <= self.range.len() - offset);
        assert_eq!((align_of::<T>() - 1) & (self.range.start + offset), 0);
    }

    pub fn typed<T: Pod>(self: &Rc<Self>) -> Rc<PwMemTyped<T>> {
        self.typed_at(0)
    }

    pub fn typed_at<T: Pod>(self: &Rc<Self>, offset: usize) -> Rc<PwMemTyped<T>> {
        self.check::<T>(offset);
        Rc::new(PwMemTyped {
            mem: self.clone(),
            offset: self.range.start + offset,
            _phantom: Default::default(),
        })
    }

    pub fn slice(self: &Rc<Self>, range: Range<usize>) -> Rc<PwMemSlice> {
        assert!(range.start <= self.range.len());
        assert!(range.len() <= self.range.len() - range.start);
        Rc::new(PwMemSlice {
            mem: self.clone(),
            range: self.range.start + range.start..self.range.start + range.end,
        })
    }
}

impl<T: Pod> PwMemTyped<T> {
    pub unsafe fn read(&self) -> &T {
        unsafe { (self.mem.map.ptr.cast::<u8>().add(self.offset) as *const T).deref() }
    }

    pub unsafe fn write(&self) -> &mut T {
        unsafe { (self.mem.map.ptr.cast::<u8>().add(self.offset) as *mut T).deref_mut() }
    }
}

#[derive(Debug, Error)]
pub enum PwMemError {
    #[error("mmap failed")]
    MmapFailed(#[source] OsError),
    #[error("memid {0} does not exist")]
    MemidDoesNotExist(u32),
}