use clap::Parser;
use eyre::{Context, ContextCompat, bail, ensure, eyre};
use image::{DynamicImage, GrayAlphaImage, ImageReader, imageops::overlay};
use nix::ioctl_write_ptr;
use std::{
    fs::File,
    os::fd::AsRawFd,
    path::{Path, PathBuf},
};

#[derive(Parser)]
struct Opt {
    /// Path to the image-blob.
    pub path: PathBuf,

    /// Stamp an image onto the final image. "<x>:<y>:<path>"
    #[clap(long)]
    pub stamp: Vec<String>,
}

const SCREEN_W: usize = 1872;
const SCREEN_H: usize = 1404;

/// Image format is a buffer of `SCREEN_W * SCREEN_H + 1` bytes.
/// The 4 least-significant bits of each byte is one pixel. The 4 most-significant bits are unused.
/// Images are stored in landscape format, i.e. the first byte is the top-left pixel, and the
/// second-to-last byte is the bottom-right pixel. The final byte is a null-terminator, and must be 0.
const IOCTL_BUFFER_SIZE: usize = SCREEN_W * SCREEN_H + 1;

// const SPI_IOC_MAGIC: u8 = b'k'; // Defined in linux/spi/spidev.h
// const SPI_IOC_TYPE_MESSAGE: u8 = 0;
// ioctl_write_buf!(spi_transfer, SPI_IOC_MAGIC, SPI_IOC_TYPE_MESSAGE, spi_ioc_transfer);
// #define DRM_IOCTL_ROCKCHIP_EBC_OFF_SCREEN	DRM_IOWR(DRM_COMMAND_BASE + 0x01, struct drm_rockchip_ebc_off_screen)

#[repr(C)]
struct DrmRockchipEbcOffScreen {
    info1: u64,
    ptr_screen_content: *const u8,
}

const DRM_IOCTL_BASE: u8 = b'd';
const DRM_COMMAND_BASE: u8 = 0x40;
const ROCKCHIP_EBC_FILE: &str = "/dev/dri/by-path/platform-fdec0000.ebc-card";

ioctl_write_ptr!(
    set_off_screen,
    DRM_IOCTL_BASE,
    DRM_COMMAND_BASE + 1,
    DrmRockchipEbcOffScreen
);

/// Convert [DynamicImage] into a [GrayAlphaImage].
/// Note that we use [GrayAlphaImage] instead of `GrayImage` because we need the alpha channel to
/// be able to add stamps.
fn to_grayscale(path: &Path, image: DynamicImage) -> GrayAlphaImage {
    if let DynamicImage::ImageLumaA8(image) = image {
        return image;
    }

    eprintln!("Image {path:?} is not 16-bit grayscale+alpha. Converting...");
    image.to_luma_alpha8()
}

fn load_wallpaper(path: &Path) -> eyre::Result<GrayAlphaImage> {
    eprintln!("Opening image at {path:?}");

    let mut image = ImageReader::open(path)?.decode()?;

    let image_dimensions = (image.width() as usize, image.height() as usize);

    match image_dimensions {
        (SCREEN_W, SCREEN_H) => {}
        (SCREEN_H, SCREEN_W) => image = image.rotate90().flipv(),
        _ => bail!("Image must be {SCREEN_W}x{SCREEN_H}"),
    }

    Ok(to_grayscale(path, image))
}

fn main() -> eyre::Result<()> {
    let opt = Opt::parse();
    color_eyre::install()?;

    let mut image = load_wallpaper(&opt.path)?;

    for stamp in &opt.stamp {
        apply_stamp(stamp, &mut image)?;
    }

    let pixels = image_to_offimage_buffer(image)?;

    let set_off_screen_data = DrmRockchipEbcOffScreen {
        info1: 0, // TODO: what is this?
        ptr_screen_content: pixels.as_ptr(),
    };

    let driver_file = File::options()
        .write(true)
        .open(ROCKCHIP_EBC_FILE)
        .wrap_err_with(|| eyre!("Failed to open rockchip ebc file at {ROCKCHIP_EBC_FILE:?}"))?;

    unsafe { set_off_screen(driver_file.as_raw_fd(), &set_off_screen_data) }
        .wrap_err("ioctl error")?;

    Ok(())
}

fn image_to_offimage_buffer(image: GrayAlphaImage) -> eyre::Result<Box<[u8; IOCTL_BUFFER_SIZE]>> {
    let mut buf: Box<[u8; IOCTL_BUFFER_SIZE]> = vec![0u8; IOCTL_BUFFER_SIZE]
        .into_boxed_slice()
        .try_into()
        .expect("Buffer is the correct size");

    let image = image.into_vec();

    ensure!(
        image.len() / 2 + 1 == IOCTL_BUFFER_SIZE,
        "Invalid size of image backing buffer. Bug?",
    );

    // Iterate over pixels and copy them into `buf`.
    for (i, pixel) in image.chunks_exact(2).enumerate() {
        let &[gray, alpha] = pixel else {
            unreachable!()
        };

        // ignore the alpha channel.
        let _ = alpha;

        // Convert 8-bit colorspace to 4-bits by truncating the 4 least significant bits.
        // Yes, this means that the most-significant 4 bits of each byte are unused.
        // Not idea why they designed it this way. Maybe it's faster?
        buf[i] = gray >> 4;
    }

    // Sanity check buffer length
    assert_eq!(buf.len(), IOCTL_BUFFER_SIZE);

    // Sanity-check that last byte is null.
    assert_eq!(buf.last(), Some(&0u8));

    Ok(buf)
}

fn parse_stamp(stamp: &str) -> Option<(i64, i64, &Path)> {
    let (x, stamp) = stamp.split_once(':')?;
    let (y, path) = stamp.split_once(':')?;

    let x = x.parse().ok()?;
    let y = y.parse().ok()?;
    let path = Path::new(path);
    Some((x, y, path))
}

fn apply_stamp(stamp: &str, onto: &mut GrayAlphaImage) -> eyre::Result<()> {
    let (x, y, path) =
        parse_stamp(stamp).wrap_err_with(|| eyre!("Invalid stamp format: {stamp:?}"))?;

    let image = load_stamp(path)?;

    overlay(onto, &image, x, y);

    Ok(())
}

fn load_stamp(path: &Path) -> eyre::Result<GrayAlphaImage> {
    eprintln!("Opening image at {path:?}");

    let image = ImageReader::open(path)?.decode()?;

    Ok(to_grayscale(path, image))
}