onix 0.1.0

use std::fs::File;
use std::fs::OpenOptions;

use std::os::fd::{AsFd, AsRawFd, BorrowedFd};

use drm::buffer::{self, DrmFourcc, DrmModifier};
use drm::control::{self, atomic, connector, crtc, property, AtomicCommitFlags};
use drm::{control::Device as _, Device as _};

struct DmaBuf {
    width: u32,
    height: u32,
    format: DrmFourcc,
    modifier: DrmModifier,
    handles: [Option<buffer::Handle>; 4],
    pitches: [u32; 4],
    offsets: [u32; 4],
}

impl DmaBuf {
    pub fn new(handle: buffer::Handle, frame: &onix::Frame) -> DmaBuf {
        let mut handles = [None; 4];
        let mut pitches = [0; 4];
        let mut offsets = [0; 4];
        for (i, plane) in frame.planes.iter().enumerate() {
            if let Some(plane) = plane {
                handles[i] = Some(handle);
                pitches[i] = plane.pitch;
                offsets[i] = plane.offset;
            }
        }
        DmaBuf {
            width: frame.width,
            height: frame.height,
            format: DrmFourcc::try_from(frame.format).unwrap(),
            modifier: DrmModifier::try_from(frame.modifier).unwrap(),
            handles,
            pitches,
            offsets,
        }
    }
}

impl buffer::PlanarBuffer for DmaBuf {
    fn size(&self) -> (u32, u32) {
        (self.width, self.height)
    }

    fn format(&self) -> DrmFourcc {
        self.format
    }

    fn modifier(&self) -> Option<DrmModifier> {
        Some(self.modifier)
    }

    fn handles(&self) -> [Option<buffer::Handle>; 4] {
        self.handles
    }

    fn pitches(&self) -> [u32; 4] {
        self.pitches
    }

    fn offsets(&self) -> [u32; 4] {
        self.offsets
    }
}

#[derive(Debug)]
/// A simple wrapper for a device node.
struct Card(File);

/// Implementing [`AsFd`] is a prerequisite to implementing the traits found
/// in this crate. Here, we are just calling [`File::as_fd()`] on the inner
/// [`File`].
impl AsFd for Card {
    fn as_fd(&self) -> BorrowedFd<'_> {
        self.0.as_fd()
    }
}

/// With [`AsFd`] implemented, we can now implement [`drm::Device`].
impl drm::Device for Card {}
impl drm::control::Device for Card {}

impl Card {
    /// Simple helper method for opening a [`Card`].
    fn open(path: &str) -> std::io::Result<Self> {
        let mut options = OpenOptions::new();
        options.read(true);
        options.write(true);

        // The normal location of the primary device node on Linux
        Ok(Card(options.open(path)?))
    }
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let args: Vec<_> = std::env::args().collect();
    if args.len() < 2 || args.len() > 3 {
        eprintln!("Usage: {} <image.webp> [/dev/dri/card1]", args[0]);
        std::process::exit(1);
    }

    let card_filename = args
        .get(2)
        .cloned()
        .unwrap_or_else(|| String::from("/dev/dri/card1"));

    let webp = std::fs::read(&args[1])?;
    assert_eq!(&webp[..4], b"RIFF");
    assert_eq!(&webp[8..12], b"WEBP");
    assert_eq!(&webp[12..16], b"VP8 ");

    let size1 = u32::from_le_bytes([webp[4], webp[5], webp[6], webp[7]]);
    let size2 = u32::from_le_bytes([webp[16], webp[17], webp[18], webp[19]]);
    assert_eq!(size1 as usize, webp.len() - 8);
    assert_eq!(size2 as usize, webp.len() - 20);
    let vp8 = &webp[20..];

    let ctrl = {
        let mut parser = onix::vp8::Parser::new(vp8);
        parser.parse_vp8()?
    };

    // Decode our VP8 using V4L2.
    let decoder = onix::Decoder::find_devices()
        .expect("Unable to find a V4L2 M2M decoder corresponding to our criteria");
    let mut decoder = decoder
        .for_vp8(ctrl)
        .expect("No available decoder which supports VP8 on this system");

    println!(
        "Decoding from media {} and video {}",
        decoder.media_device_info()?.driver(),
        decoder.video_querycap()?.driver()
    );

    decoder
        .set_input_length(vp8.len() as u32)
        .expect("Unable to set length of the input VP8 data");
    let format = decoder
        .set_output_format(DrmFourcc::Nv12 as u32)
        .expect("Unable to set the output format to NV12");
    let fd = decoder
        .start_vp8_decode(vp8)
        .expect("Unable to start decoding our VP8 data");

    // Now display to DRM.
    let card = Card::open(&card_filename).expect("Couldn’t open {card_filename}");
    let driver = card.get_driver()?;
    let drm_name = driver.name.to_str().unwrap();
    let drm_desc = driver.desc.to_str().unwrap();
    println!("Outputting to {} ({})", drm_name, drm_desc);

    card.set_client_capability(drm::ClientCapability::UniversalPlanes, true)
        .expect("Unable to request UniversalPlanes capability");
    card.set_client_capability(drm::ClientCapability::Atomic, true)
        .expect("Unable to request Atomic capability");

    // Load the information.
    let res = card
        .resource_handles()
        .expect("Could not load normal resource ids.");
    let coninfo: Vec<connector::Info> = res
        .connectors()
        .iter()
        .flat_map(|con| card.get_connector(*con, true))
        .collect();
    let crtcinfo: Vec<crtc::Info> = res
        .crtcs()
        .iter()
        .flat_map(|crtc| card.get_crtc(*crtc))
        .collect();

    // Filter each connector until we find one that's connected.
    let con = coninfo
        .iter()
        .find(|&i| i.state() == connector::State::Connected)
        .expect("No connected connectors");

    // Get the first (usually best) mode
    let &mode = con.modes().get(0).expect("No modes found on connector");

    let (width, height) = mode.size();

    // Find a crtc and FB
    let crtc = crtcinfo.get(0).expect("No crtcs found");

    // Select the pixel format
    let fmt = DrmFourcc::Xrgb8888;

    // Create a DB
    // If buffer resolution is above display resolution, a ENOSPC (not enough GPU memory) error may
    // occur
    let db = card
        .create_dumb_buffer((width as u32, height as u32), fmt, 32)
        .expect("Could not create dumb buffer");

    // Create an FB:
    let fb = card
        .add_framebuffer(&db, 24, 32)
        .expect("Could not create FB");

    let planes = card.plane_handles().expect("Could not list planes");
    let (better_planes, compatible_planes): (
        Vec<control::plane::Handle>,
        Vec<control::plane::Handle>,
    ) = planes
        .iter()
        .filter(|&&plane| {
            card.get_plane(plane)
                .map(|plane_info| {
                    let compatible_crtcs = res.filter_crtcs(plane_info.possible_crtcs());
                    compatible_crtcs.contains(&crtc.handle())
                })
                .unwrap_or(false)
        })
        .partition(|&&plane| {
            if let Ok(props) = card.get_properties(plane) {
                for (&id, &val) in props.iter() {
                    if let Ok(info) = card.get_property(id) {
                        if info.name().to_str().map(|x| x == "type").unwrap_or(false) {
                            return val == (drm::control::PlaneType::Primary as u32).into();
                        }
                    }
                }
            }
            false
        });
    let plane = *better_planes
        .get(0)
        .unwrap_or_else(|| &compatible_planes[0]);

    let (better_planes, compatible_planes): (
        Vec<control::plane::Handle>,
        Vec<control::plane::Handle>,
    ) = planes
        .iter()
        .filter(|&&plane| {
            card.get_plane(plane)
                .map(|plane_info| {
                    let compatible_crtcs = res.filter_crtcs(plane_info.possible_crtcs());
                    compatible_crtcs.contains(&crtc.handle())
                        && plane_info.formats().contains(&0x3231564e /*NV12*/)
                })
                .unwrap_or(false)
        })
        .partition(|&&plane| {
            if let Ok(props) = card.get_properties(plane) {
                for (&id, &val) in props.iter() {
                    if let Ok(info) = card.get_property(id) {
                        if info.name().to_str().map(|x| x == "type").unwrap_or(false) {
                            return val == (drm::control::PlaneType::Overlay as u32).into();
                        }
                    }
                }
            }
            false
        });
    let nv12_plane = *better_planes
        .get(0)
        .unwrap_or_else(|| &compatible_planes[0]);

    let con_props = card.get_properties(con.handle())?.as_hashmap(&card)?;
    let crtc_props = card.get_properties(crtc.handle())?.as_hashmap(&card)?;
    let plane_props = card.get_properties(plane)?.as_hashmap(&card)?;
    let nv12_props = card.get_properties(nv12_plane)?.as_hashmap(&card)?;

    let mut atomic_req = atomic::AtomicModeReq::new();
    atomic_req.add_property(
        con.handle(),
        con_props["CRTC_ID"].handle(),
        property::Value::CRTC(Some(crtc.handle())),
    );
    let blob = card
        .create_property_blob(&mode)
        .expect("Failed to create blob");
    atomic_req.add_property(crtc.handle(), crtc_props["MODE_ID"].handle(), blob);
    atomic_req.add_property(
        crtc.handle(),
        crtc_props["ACTIVE"].handle(),
        property::Value::Boolean(true),
    );
    atomic_req.add_property(
        plane,
        plane_props["FB_ID"].handle(),
        property::Value::Framebuffer(Some(fb)),
    );
    atomic_req.add_property(
        plane,
        plane_props["CRTC_ID"].handle(),
        property::Value::CRTC(Some(crtc.handle())),
    );
    atomic_req.add_property(
        plane,
        plane_props["SRC_X"].handle(),
        property::Value::UnsignedRange(0),
    );
    atomic_req.add_property(
        plane,
        plane_props["SRC_Y"].handle(),
        property::Value::UnsignedRange(0),
    );
    atomic_req.add_property(
        plane,
        plane_props["SRC_W"].handle(),
        property::Value::UnsignedRange((width as u64) << 16),
    );
    atomic_req.add_property(
        plane,
        plane_props["SRC_H"].handle(),
        property::Value::UnsignedRange((height as u64) << 16),
    );
    atomic_req.add_property(
        plane,
        plane_props["CRTC_X"].handle(),
        property::Value::SignedRange(0),
    );
    atomic_req.add_property(
        plane,
        plane_props["CRTC_Y"].handle(),
        property::Value::SignedRange(0),
    );
    atomic_req.add_property(
        plane,
        plane_props["CRTC_W"].handle(),
        property::Value::UnsignedRange(width as u64),
    );
    atomic_req.add_property(
        plane,
        plane_props["CRTC_H"].handle(),
        property::Value::UnsignedRange(height as u64),
    );
    atomic_req.add_property(
        plane,
        plane_props["zpos"].handle(),
        property::Value::UnsignedRange(0),
    );

    //decoder.poll()?;
    let frame = decoder
        .finish_vp8_decode(format)
        .expect("Unable to finish decoding our VP8");
    let handle = card
        .prime_fd_to_buffer(fd.as_raw_fd())
        .expect("Unable to convert the dmabuf fd to a GEM handle");
    let buffer = DmaBuf::new(handle, &frame);
    let nv12_fb = card
        .add_planar_framebuffer(&buffer, control::FbCmd2Flags::MODIFIERS)
        .expect("Could not create planar FB");

    let screen_ratio = width as f32 / height as f32;
    let image_ratio = frame.width as f32 / frame.height as f32;
    let (image_width, image_height) = if screen_ratio > image_ratio {
        (frame.width * height as u32 / frame.height, height as u32)
    } else {
        (width as u32, frame.height * width as u32 / frame.width)
    };

    atomic_req.add_property(
        nv12_plane,
        nv12_props["FB_ID"].handle(),
        property::Value::Framebuffer(Some(nv12_fb)),
    );
    atomic_req.add_property(
        nv12_plane,
        nv12_props["CRTC_ID"].handle(),
        property::Value::CRTC(Some(crtc.handle())),
    );
    atomic_req.add_property(
        nv12_plane,
        nv12_props["SRC_X"].handle(),
        property::Value::UnsignedRange(0),
    );
    atomic_req.add_property(
        nv12_plane,
        nv12_props["SRC_Y"].handle(),
        property::Value::UnsignedRange(0),
    );
    atomic_req.add_property(
        nv12_plane,
        nv12_props["SRC_W"].handle(),
        property::Value::UnsignedRange((frame.width as u64) << 16),
    );
    atomic_req.add_property(
        nv12_plane,
        nv12_props["SRC_H"].handle(),
        property::Value::UnsignedRange((frame.height as u64) << 16),
    );
    atomic_req.add_property(
        nv12_plane,
        nv12_props["CRTC_X"].handle(),
        property::Value::SignedRange((width as i64 - image_width as i64) / 2),
    );
    atomic_req.add_property(
        nv12_plane,
        nv12_props["CRTC_Y"].handle(),
        property::Value::SignedRange((height as i64 - image_height as i64) / 2),
    );
    atomic_req.add_property(
        nv12_plane,
        nv12_props["CRTC_W"].handle(),
        property::Value::UnsignedRange(image_width as u64),
    );
    atomic_req.add_property(
        nv12_plane,
        nv12_props["CRTC_H"].handle(),
        property::Value::UnsignedRange(image_height as u64),
    );
    atomic_req.add_property(
        nv12_plane,
        nv12_props["zpos"].handle(),
        property::Value::UnsignedRange(1),
    );

    // Set the crtc
    // On many setups, this requires root access.
    card.atomic_commit(AtomicCommitFlags::ALLOW_MODESET, atomic_req)
        .expect("Failed to set mode");

    let five_seconds = std::time::Duration::from_millis(5000);
    std::thread::sleep(five_seconds);

    card.destroy_framebuffer(fb).unwrap();
    card.destroy_dumb_buffer(db).unwrap();

    Ok(())
}