use {
    ahash::AHashMap,
    byteorder::{NativeEndian, WriteBytesExt},
    jay_video_types::Modifier,
    std::{io::Write, rc::Rc},
    thiserror::Error,
    uapi::{OwnedFd, c},
};

#[derive(Debug)]
pub struct DrmFeedbackIds {
    next: std::cell::Cell<u64>,
}

impl Default for DrmFeedbackIds {
    fn default() -> Self {
        Self {
            next: std::cell::Cell::new(1),
        }
    }
}

impl DrmFeedbackIds {
    pub fn next(&self) -> DrmFeedbackId {
        let id = self.next.get();
        self.next.set(id + 1);
        DrmFeedbackId(id)
    }
}

#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash)]
pub struct DrmFeedbackId(u64);

#[derive(Debug)]
pub struct DrmFeedbackShared {
    pub fd: Rc<OwnedFd>,
    pub size: usize,
    pub main_device: c::dev_t,
    pub indices: AHashMap<(u32, Modifier), u16>,
}

#[derive(Debug)]
pub struct DrmFeedback {
    pub id: DrmFeedbackId,
    pub shared: Rc<DrmFeedbackShared>,
    pub tranches: Vec<DrmFeedbackTranche>,
}

#[derive(Clone, Debug)]
pub struct DrmFeedbackTranche {
    pub device: c::dev_t,
    pub indices: Vec<u16>,
    pub scanout: bool,
}

impl DrmFeedback {
    pub fn new<C: DrmFeedbackContext + ?Sized>(
        ids: &DrmFeedbackIds,
        render_ctx: &C,
    ) -> Result<Self, DrmFeedbackError> {
        let main_device = match render_ctx.main_device() {
            Some(dev) => dev,
            _ => return Err(DrmFeedbackError::NoDrmDevice),
        };
        let (data, index_map) = create_fd_data(render_ctx);
        let mut memfd =
            uapi::memfd_create("drm_feedback", c::MFD_CLOEXEC | c::MFD_ALLOW_SEALING).unwrap();
        memfd.write_all(&data).unwrap();
        uapi::lseek(memfd.raw(), 0, c::SEEK_SET).unwrap();
        uapi::fcntl_add_seals(
            memfd.raw(),
            c::F_SEAL_SEAL | c::F_SEAL_GROW | c::F_SEAL_SHRINK | c::F_SEAL_WRITE,
        )
        .unwrap();
        Ok(Self {
            id: ids.next(),
            tranches: vec![DrmFeedbackTranche {
                device: main_device,
                indices: (0..index_map.len()).map(|v| v as u16).collect(),
                scanout: false,
            }],
            shared: Rc::new(DrmFeedbackShared {
                fd: Rc::new(memfd),
                size: data.len(),
                main_device,
                indices: index_map,
            }),
        })
    }

    pub fn for_scanout(
        &self,
        ids: &DrmFeedbackIds,
        devnum: c::dev_t,
        formats: &[(u32, Modifier)],
    ) -> Result<Option<Self>, DrmFeedbackError> {
        let mut tranches = vec![];
        {
            let mut indices = vec![];
            for (format, modifier) in formats {
                if let Some(idx) = self.shared.indices.get(&(*format, *modifier)) {
                    indices.push(*idx);
                }
            }
            if indices.len() > 0 {
                tranches.push(DrmFeedbackTranche {
                    device: devnum,
                    indices,
                    scanout: true,
                });
            } else {
                return Ok(None);
            }
        }
        tranches.extend(self.tranches.iter().cloned());
        Ok(Some(Self {
            id: ids.next(),
            shared: self.shared.clone(),
            tranches,
        }))
    }
}

pub trait DrmFeedbackContext {
    fn main_device(&self) -> Option<c::dev_t>;
    fn for_each_read_format(&self, f: &mut dyn FnMut(u32, Modifier));
}

fn create_fd_data<C: DrmFeedbackContext + ?Sized>(
    ctx: &C,
) -> (Vec<u8>, AHashMap<(u32, Modifier), u16>) {
    let mut vec = vec![];
    let mut map = AHashMap::new();
    let mut pos = 0;
    ctx.for_each_read_format(&mut |format, modifier| {
        vec.write_u32::<NativeEndian>(format).unwrap();
        vec.write_u32::<NativeEndian>(0).unwrap();
        vec.write_u64::<NativeEndian>(modifier).unwrap();
        map.insert((format, modifier), pos);
        pos += 1;
    });
    (vec, map)
}

#[derive(Debug, Error)]
pub enum DrmFeedbackError {
    #[error("Graphics API does not have a DRM device")]
    NoDrmDevice,
}