use std::{
    io::Write,
    os::fd::AsFd,
    sync::{Arc, RwLock},
};

use color_eyre::{Result, eyre::eyre};
use image::{EncodableLayout, ImageReader};
use tracing::{debug, error, trace};
use wayland_client::{
    Dispatch,
    protocol::{
        wl_buffer, wl_compositor, wl_region,
        wl_registry::{self, WlRegistry},
        wl_shm, wl_shm_pool, wl_surface,
    },
};
use wayland_protocols_wlr::layer_shell::v1::client::{zwlr_layer_shell_v1, zwlr_layer_surface_v1};

fn main() -> Result<()> {
    tracing_subscriber::fmt::init();
    color_eyre::install()?;

    debug!("hello world!");

    let wayland_connection = wayland_client::Connection::connect_to_env()?;

    let display = wayland_connection.display();

    let mut queue = wayland_connection.new_event_queue();
    let queue_handle = queue.handle();

    display.get_registry(&queue_handle, ());

    let mut state = AppState::new();

    if let Err(e) = queue.roundtrip(&mut state) {
        error!("Roundtrip failed: {e}");
        Err(e)?
    }

    let Some(ref layer_shell) = state.layer_shell else {
        Err(eyre!("No zwlr_layer_shell found"))?
    };
    let Some(ref compositor) = state.compositor else {
        Err(eyre!("No wl_compositor found"))?
    };
    let Some(shm) = state.shm.clone() else {
        Err(eyre!("No wl_shm found"))?
    };

    let img = ImageReader::open("crosshair.png")?.decode()?;
    let rgba_image = img.as_rgba8().unwrap();
    let img_bytes = &rgba_image.as_bytes();

    let surface = compositor.create_surface(&queue_handle, ());
    let input_region = compositor.create_region(&queue_handle, ());
    surface.set_input_region(Some(&input_region));

    let layer_surface_data = Arc::new(RwLock::new((false, 0, 0)));
    let layer_surface = layer_shell.get_layer_surface(
        &surface,
        None,
        zwlr_layer_shell_v1::Layer::Overlay,
        "crosshair".to_string(),
        &queue_handle,
        layer_surface_data.clone(),
    );
    layer_surface.set_size(img.width(), img.height());
    layer_surface.set_exclusive_zone(-1); // ignore other exclusive zones

    surface.commit();

    while !layer_surface_data.read().unwrap().0 {
        queue.blocking_dispatch(&mut state)?;
    }

    let mut memfile = {
        let memfile = memfd::MemfdOptions::default()
            .allow_sealing(true)
            .create("crosshair_image_data")?;

        memfile.as_file().set_len(img_bytes.len() as u64)?;

        memfile.add_seals(&[
            memfd::FileSeal::SealShrink,
            memfd::FileSeal::SealGrow,
            memfd::FileSeal::SealSeal,
        ])?;

        memfile.into_file()
    };
    memfile.write_all(img_bytes)?;

    let shm_pool = shm.create_pool(memfile.as_fd(), (img_bytes.len()) as i32, &queue_handle, ());

    let buffer = shm_pool.create_buffer(
        0,
        img.width() as i32,
        img.height() as i32,
        (img.width() * 4) as i32,
        wl_shm::Format::Abgr8888,
        &queue_handle,
        (),
    );

    surface.attach(Some(&buffer), 0, 0);
    surface.commit();

    loop {
        debug!("dispatched!");
        queue.blocking_dispatch(&mut state)?;
    }
}

#[derive(Default)]
struct AppState {
    layer_shell: Option<zwlr_layer_shell_v1::ZwlrLayerShellV1>,
    shm: Option<wl_shm::WlShm>,
    compositor: Option<wl_compositor::WlCompositor>,
    formats: Vec<wl_shm::Format>,
}

impl AppState {
    fn new() -> Self {
        Self::default()
    }
}

impl Dispatch<wl_registry::WlRegistry, ()> for AppState {
    fn event(
        state: &mut Self,
        proxy: &WlRegistry,
        event: <WlRegistry as wayland_client::Proxy>::Event,
        _data: &(),
        _conn: &wayland_client::Connection,
        qhandle: &wayland_client::QueueHandle<Self>,
    ) {
        trace!("event: {event:?}");
        if let wl_registry::Event::Global {
            name,
            interface,
            version,
        } = event
        {
            match interface.as_str() {
                "zwlr_layer_shell_v1" => {
                    state.layer_shell =
                        Some(proxy.bind::<zwlr_layer_shell_v1::ZwlrLayerShellV1, _, _>(
                            name,
                            version,
                            qhandle,
                            (),
                        ))
                }
                "wl_compositor" => {
                    state.compositor = Some(proxy.bind::<wl_compositor::WlCompositor, _, _>(
                        name,
                        version,
                        qhandle,
                        (),
                    ))
                }
                "wl_shm" => {
                    state.shm = Some(proxy.bind::<wl_shm::WlShm, _, _>(name, version, qhandle, ()))
                }
                _ => (),
            }
        }
    }
}

impl Dispatch<zwlr_layer_shell_v1::ZwlrLayerShellV1, ()> for AppState {
    fn event(
        _state: &mut Self,
        _proxy: &zwlr_layer_shell_v1::ZwlrLayerShellV1,
        _event: <zwlr_layer_shell_v1::ZwlrLayerShellV1 as wayland_client::Proxy>::Event,
        _data: &(),
        _conn: &wayland_client::Connection,
        _qhandle: &wayland_client::QueueHandle<Self>,
    ) {
        // ignore
    }
}

impl Dispatch<wl_compositor::WlCompositor, ()> for AppState {
    fn event(
        _state: &mut Self,
        _proxy: &wl_compositor::WlCompositor,
        _event: <wl_compositor::WlCompositor as wayland_client::Proxy>::Event,
        _data: &(),
        _conn: &wayland_client::Connection,
        _qhandle: &wayland_client::QueueHandle<Self>,
    ) {
        // ignore
    }
}

impl Dispatch<wl_surface::WlSurface, ()> for AppState {
    fn event(
        _state: &mut Self,
        _proxy: &wl_surface::WlSurface,
        _event: <wl_surface::WlSurface as wayland_client::Proxy>::Event,
        _data: &(),
        _conn: &wayland_client::Connection,
        _qhandle: &wayland_client::QueueHandle<Self>,
    ) {
        // ignore
    }
}

impl Dispatch<zwlr_layer_surface_v1::ZwlrLayerSurfaceV1, Arc<RwLock<(bool, u32, u32)>>>
    for AppState
{
    fn event(
        _state: &mut Self,
        proxy: &zwlr_layer_surface_v1::ZwlrLayerSurfaceV1,
        event: <zwlr_layer_surface_v1::ZwlrLayerSurfaceV1 as wayland_client::Proxy>::Event,
        data: &Arc<RwLock<(bool, u32, u32)>>,
        _conn: &wayland_client::Connection,
        _qhandle: &wayland_client::QueueHandle<Self>,
    ) {
        trace!("layer surface event: {event:?}");
        if let zwlr_layer_surface_v1::Event::Configure {
            serial,
            width,
            height,
        } = event
        {
            proxy.ack_configure(serial);
            let mut data = data.write().unwrap();
            *data = (true, width, height);
        }
    }
}

impl Dispatch<wl_shm::WlShm, ()> for AppState {
    fn event(
        state: &mut Self,
        _proxy: &wl_shm::WlShm,
        event: <wl_shm::WlShm as wayland_client::Proxy>::Event,
        _data: &(),
        _conn: &wayland_client::Connection,
        _qhandle: &wayland_client::QueueHandle<Self>,
    ) {
        trace!("wl_shm event: {event:?}");
        if let wl_shm::Event::Format { format } = event {
            match format {
                wayland_client::WEnum::Value(f) => state.formats.push(f),
                wayland_client::WEnum::Unknown(_) => (),
            }
        }
    }
}

impl Dispatch<wl_shm_pool::WlShmPool, ()> for AppState {
    fn event(
        _state: &mut Self,
        _proxy: &wl_shm_pool::WlShmPool,
        _event: <wl_shm_pool::WlShmPool as wayland_client::Proxy>::Event,
        _data: &(),
        _conn: &wayland_client::Connection,
        _qhandle: &wayland_client::QueueHandle<Self>,
    ) {
        // ignore
    }
}

impl Dispatch<wl_buffer::WlBuffer, ()> for AppState {
    fn event(
        _state: &mut Self,
        _proxy: &wl_buffer::WlBuffer,
        _event: <wl_buffer::WlBuffer as wayland_client::Proxy>::Event,
        _data: &(),
        _conn: &wayland_client::Connection,
        _qhandle: &wayland_client::QueueHandle<Self>,
    ) {
        // ignore
    }
}

impl Dispatch<wl_region::WlRegion, ()> for AppState {
    fn event(
        _state: &mut Self,
        _proxy: &wl_region::WlRegion,
        _event: <wl_region::WlRegion as wayland_client::Proxy>::Event,
        _data: &(),
        _conn: &wayland_client::Connection,
        _qhandle: &wayland_client::QueueHandle<Self>,
    ) {
        // ignore
    }
}