//https://gaultier.github.io/blog/wayland_from_scratch.html
use std::{ffi::CStr, io::Read, ptr::null_mut, u32};

use poppable::Poppable;
use rustix::{
    fs::{ftruncate, Mode},
    mm::{mmap, MapFlags, ProtFlags},
    shm,
};
use wayland::WAYLAND_XDG_BASE_EVENT_PING;

mod wayland {
    use anyhow::Result;
    use poppable::{PopFromNE, Pushable};
    use rustix::fd::OwnedFd;
    use std::ffi::c_void;
    use std::{env, io::Write, os::unix::net::UnixStream};

    pub const WAYLAND_DISPLAY_OBJECT_ID: u32 = 1;
    pub const WAYLAND_WL_DISPLAY_GET_REGISTRY_OPCODE: u16 = 1;
    pub const WAYLAND_WL_REGISTRY_BIND_OPCODE: u16 = 0;
    pub const COLOR_CHANNELS: u32 = 4;
    pub const WAYLAND_WL_REGISTRY_EVENT_GLOBAL: u16 = 0;
    pub const WAYLAND_WL_DISPLAY_ERROR_EVENT: u16 = 0;
    pub const WAYLAND_XDG_BASE_EVENT_PING: u16 = 0;
    pub const WAYLAND_WL_COMPOSITOR_CREATE_SURFACE_OPCODE: u16 = 0;
    pub const WAYLAND_XDG_SURFACE_EVENT_CONFIGURE: u16 = 0;
    pub const WAYLAND_WL_SURFACE_COMMIT_OPCODE: u16 = 6;
    pub const WAYLAND_XDG_WM_BASE_GET_XDG_SURFACE_OPCODE: u16 = 2;
    pub const WAYLAND_XDG_SURFACE_GET_TOPLEVEL_OPCODE: u16 = 1;
    pub const WAYLAND_XDG_WM_BASE_PONG_OPCODE: u16 = 3;
    pub const WAYLAND_XDG_SURFACE_ACK_CONFIGURE_OPCODE: u16 = 4;
    pub const WAYLAND_WL_SHM_POOL_CREATE_BUFFER_OPCODE: u16 = 0;
    pub const WAYLAND_FORMAT_XRGB8888: u32 = 1;
    pub const WAYLAND_WL_SURFACE_ATTACH_OPCODE: u16 = 1;
    pub const WAYLAND_WL_SHM_CREATE_POOL_OPCODE: u16 = 0;

    pub fn write_string(buf: &mut Vec<u8>, src: &str) -> Result<usize> {
        let mut n_write = 0;
        (src.len() as u32 + 1).push_ne_into(buf)?;
        for b in src.bytes().chain([0_u8]) {
            buf.push(b);
            n_write += 1;
        }
        let required_size = nearest_multiple_4(n_write as i32) as usize;
        while n_write < required_size {
            buf.push(0);
            n_write += 1;
        }

        return Ok(n_write);
    }
    #[inline(always)]
    pub fn nearest_multiple_4(n: i32) -> i32 {
        return ((n) + 3) & -4;
    }

    #[derive(Default, Debug, PartialEq)]
    pub enum StateEnum {
        #[default]
        None,
        SurfaceAckedConfigure,
        SurfaceAttached,
    }
    #[derive(Default)]
    pub struct State {
        pub wl_registry: u32,
        pub wl_shm: u32,
        pub wl_shm_pool: u32,
        pub wl_buffer: u32,
        pub xdg_wm_base: u32,
        pub xdg_surface: u32,
        pub wl_compositor: u32,
        pub wl_surface: u32,
        pub xdg_toplevel: u32,
        pub stride: u32,
        pub w: u32,
        pub h: u32,
        pub shm_pool_size: u32,
        pub shm_fd: Option<OwnedFd>,
        pub shm_pool_data: Option<*mut c_void>,

        pub state: StateEnum,
    }

    pub struct Wayland {
        pub sock: UnixStream,
        pub resource_counter: u32,
    }
    pub fn connect() -> Result<Wayland> {
        let xdg_runtime_dir = env::var("XDG_RUNTIME_DIR")?;
        println!("Using xdg runtime dir: {}", xdg_runtime_dir);
        let wayland_display = env::var("WAYLAND_DISPLAY").unwrap_or("wayland-0".into());
        let wayland_path = xdg_runtime_dir + "/" + &wayland_display;
        println!("Opening UnixStream: {wayland_path}");
        let sock = UnixStream::connect(&wayland_path)?;
        Ok(Wayland {
            sock,
            resource_counter: 1,
        })
    }

    #[derive(Debug)]
    struct RawWaylandMsg {
        resource_id: u32,
        opcode: u16,
        size: u16,
        payload: Vec<u8>,
    }
    impl RawWaylandMsg {
        const HEADER_SIZE: u16 =
            (std::mem::size_of::<u32>() + (2 * std::mem::size_of::<u16>())) as u16;
        pub fn new(resource_id: u32, opcode: u16, payload: Vec<u8>) -> Self {
            RawWaylandMsg {
                resource_id,
                opcode,
                size: Self::HEADER_SIZE + payload.len() as u16,
                payload,
            }
        }
    }

    impl Wayland {
        pub fn wl_compositor_create_surface(&mut self, state: &mut State) -> Result<u32> {
            assert!(state.wl_compositor > 0);
            let resource_id = state.wl_compositor;
            let opcode = WAYLAND_WL_COMPOSITOR_CREATE_SURFACE_OPCODE;

            self.resource_counter += 1;
            let object_id = self.resource_counter;
            let mut payload = Vec::new();
            object_id.push_ne_into(&mut payload)?;

            let size = RawWaylandMsg::HEADER_SIZE + payload.len() as u16;

            let msg = RawWaylandMsg {
                resource_id,
                opcode,
                size,
                payload,
            };
            println!(
                "-> wl_compositor@{}.create_surface: wl_surface={}",
                state.wl_compositor, object_id
            );
            self.raw_send(msg)?;
            Ok(object_id)
        }
        pub fn xdg_wm_base_get_xdg_surface(&mut self, state: &mut State) -> Result<u32> {
            assert!(state.xdg_wm_base > 0);
            assert!(state.wl_surface > 0);

            self.resource_counter += 1;
            let object_id = self.resource_counter;
            let mut payload = object_id.to_ne_bytes().to_vec();
            state.wl_surface.push_ne_into(&mut payload)?;
            let msg: RawWaylandMsg = RawWaylandMsg::new(
                state.xdg_wm_base,
                WAYLAND_XDG_WM_BASE_GET_XDG_SURFACE_OPCODE,
                payload,
            );
            println!(
                "-> xdg_wm_base@{}.get_xdg_surface: xdg_surface={} wl_surface={}",
                state.xdg_wm_base, object_id, state.wl_surface
            );
            self.raw_send(msg)?;
            Ok(object_id)
        }
        pub fn xdg_surface_get_toplevel(&mut self, state: &mut State) -> Result<u32> {
            assert!(state.xdg_surface > 0);
            self.resource_counter += 1;
            let payload = self.resource_counter.to_ne_bytes().to_vec();
            let msg = RawWaylandMsg::new(
                state.xdg_surface,
                WAYLAND_XDG_SURFACE_GET_TOPLEVEL_OPCODE,
                payload,
            );
            println!(
                "-> xdg_surface@{}.get_toplevel: xdg_toplevel={}",
                state.xdg_surface, self.resource_counter
            );
            self.raw_send(msg)?;
            Ok(self.resource_counter)
        }
        pub fn wl_surface_commit(&mut self, state: &mut State) -> Result<()> {
            assert!(state.wl_surface > 0);
            let msg =
                RawWaylandMsg::new(state.wl_surface, WAYLAND_WL_SURFACE_COMMIT_OPCODE, vec![]);
            println!("-> wl_surface@{}.commit: ", state.wl_surface);
            self.raw_send(msg)?;
            Ok(())
        }
        fn raw_send(&mut self, msg: RawWaylandMsg) -> Result<()> {
            //println!("msg: {msg:?}");
            let mut buffer = vec![];
            msg.resource_id.push_ne_into(&mut buffer)?;
            msg.opcode.push_ne_into(&mut buffer)?;
            msg.size.push_ne_into(&mut buffer)?;
            buffer.push_vec(msg.payload)?;

            /*
            println!("msg_size: {}", buffer.len());
            println!("msg_announced_size: {}", msg.size);
            for i in buffer.iter() {
                print!("{} ", *i);
            }
            println!("");
            */

            self.sock.write(&buffer)?;
            Ok(())
        }

        pub fn get_registry(&mut self) -> Result<u32> {
            let mut payload: Vec<u8> = Vec::new();
            self.resource_counter += 1;
            self.resource_counter.push_ne_into(&mut payload)?;
            let size = RawWaylandMsg::HEADER_SIZE + payload.len() as u16;
            let msg = RawWaylandMsg {
                resource_id: WAYLAND_DISPLAY_OBJECT_ID,
                opcode: WAYLAND_WL_DISPLAY_GET_REGISTRY_OPCODE,
                size,
                payload,
            };

            println!(
                "-> wl_display@{WAYLAND_DISPLAY_OBJECT_ID}.get_registry: wl_registry={}",
                self.resource_counter
            );
            self.raw_send(msg)?;

            return Ok(self.resource_counter);
        }
        pub fn registry_bind(
            &mut self,
            wl_registry: u32,
            name: u32,
            interface: &str,
            version: u32,
        ) -> Result<u32> {
            println!(
                "-> wl_registry@{}.bind: name={} interface={} version={}",
                wl_registry, name, interface, version
            );
            let resource_id = wl_registry;
            let opcode = WAYLAND_WL_REGISTRY_BIND_OPCODE;

            let mut payload = vec![];
            name.push_ne_into(&mut payload)?;
            write_string(&mut payload, interface)?;
            version.push_ne_into(&mut payload)?;
            self.resource_counter += 1;
            self.resource_counter.push_ne_into(&mut payload)?;

            let size = RawWaylandMsg::HEADER_SIZE + payload.len() as u16;

            let msg = RawWaylandMsg {
                resource_id,
                opcode,
                size,
                payload,
            };
            self.raw_send(msg)?;
            return Ok(self.resource_counter);
        }

        pub fn xdg_wm_base_pong(&mut self, state: &mut State, ping: u32) -> Result<()> {
            assert!(state.xdg_wm_base > 0);
            assert!(state.wl_surface > 0);

            let msg = RawWaylandMsg::new(
                state.xdg_wm_base,
                WAYLAND_XDG_WM_BASE_PONG_OPCODE,
                ping.to_ne_bytes().to_vec(),
            );
            println!("-> xdg_wm_base@{}.pong: ping={}", state.xdg_wm_base, ping);
            self.raw_send(msg)?;
            Ok(())
        }

        pub fn xdg_surface_ack_configure(
            &mut self,
            state: &mut State,
            configure: u32,
        ) -> Result<()> {
            assert!(state.xdg_surface > 0);

            let msg = RawWaylandMsg::new(
                state.xdg_surface,
                WAYLAND_XDG_SURFACE_ACK_CONFIGURE_OPCODE,
                configure.to_ne_bytes().to_vec(),
            );
            println!(
                "-> xdg_surface@{}.ack_configure: configure={}",
                state.xdg_surface, configure
            );
            self.raw_send(msg)?;
            Ok(())
        }

        pub fn wl_shm_create_pool(&mut self, state: &mut State) -> Result<u32> {
            assert!(state.shm_pool_size > 0);
            let _resource_id = state.wl_shm;
            let _opcode = WAYLAND_WL_SHM_CREATE_POOL_OPCODE;
            self.resource_counter += 1;
            let _object_id = self.resource_counter;
            /*
            uint16_t msg_announced_size = wayland_header_size +
                                            sizeof(wayland_current_id) +
                                            sizeof(state->shm_pool_size);

            assert(roundup_4(msg_announced_size) == msg_announced_size);
            buf_write_u16(msg, &msg_size, sizeof(msg), msg_announced_size);

            wayland_current_id++;
            buf_write_u32(msg, &msg_size, sizeof(msg), wayland_current_id);

            buf_write_u32(msg, &msg_size, sizeof(msg), state->shm_pool_size);

            assert(roundup_4(msg_size) == msg_size);

            // Send the file descriptor as ancillary data.
            // UNIX/Macros monstrosities ahead.
            char buf[CMSG_SPACE(sizeof(state->shm_fd))] = "";

            struct iovec io = {.iov_base = msg, .iov_len = msg_size};
            struct msghdr socket_msg = {
                .msg_iov = &io,
                .msg_iovlen = 1,
                .msg_control = buf,
                .msg_controllen = sizeof(buf),
            };

            struct cmsghdr *cmsg = CMSG_FIRSTHDR(&socket_msg);
            cmsg->cmsg_level = SOL_SOCKET;
            cmsg->cmsg_type = SCM_RIGHTS;
            cmsg->cmsg_len = CMSG_LEN(sizeof(state->shm_fd));

            *((int *)CMSG_DATA(cmsg)) = state->shm_fd;
            socket_msg.msg_controllen = CMSG_SPACE(sizeof(state->shm_fd));

            if (sendmsg(fd, &socket_msg, 0) == -1)
                exit(errno);

            printf("-> wl_shm@%u.create_pool: wl_shm_pool=%u\n", state->wl_shm,
                    wayland_current_id);

            return wayland_current_id;
            */

            // FIXME:
            println!("-x wl_shm@{}.create_pool: NOT IMPLEMENTED", state.wl_shm);
            todo!()
        }

        pub fn wl_shm_pool_create_buffer(&mut self, state: &mut State) -> Result<u32> {
            assert!(state.wl_shm_pool > 0);
            let resource_id = state.wl_shm_pool;
            let opcode = WAYLAND_WL_SHM_POOL_CREATE_BUFFER_OPCODE;

            self.resource_counter += 1;
            let mut payload = self.resource_counter.to_ne_bytes().to_vec();
            0_u32.push_ne_into(&mut payload)?;
            state.w.push_ne_into(&mut payload)?;
            state.h.push_ne_into(&mut payload)?;
            state.stride.push_ne_into(&mut payload)?;
            WAYLAND_FORMAT_XRGB8888.push_ne_into(&mut payload)?;

            let msg = RawWaylandMsg::new(resource_id, opcode, payload);

            println!(
                "-> wl_shm_pool@{}.create_buffer: wl_buffer={}",
                state.wl_shm_pool, self.resource_counter
            );
            self.raw_send(msg)?;
            Ok(self.resource_counter)
        }

        pub fn wl_surface_attach(&mut self, state: &mut State) -> Result<()> {
            assert!(state.wl_surface > 0);
            assert!(state.wl_buffer > 0);

            let resource_id = state.wl_surface;
            let opcode = WAYLAND_WL_SURFACE_ATTACH_OPCODE;

            let mut payload = Vec::new();
            state.wl_buffer.push_ne_into(&mut payload)?;
            let (x, y) = (0_u32, 0_u32);
            x.push_ne_into(&mut payload)?;
            y.push_ne_into(&mut payload)?;

            let msg = RawWaylandMsg::new(resource_id, opcode, payload);
            println!(
                "-> wl_surface@{}.attach: wl_buffer={}",
                resource_id, state.wl_buffer
            );
            self.raw_send(msg)?;

            Ok(())
        }
    }
}

fn gen_rand_name(len: usize) -> String {
    let mut s = String::from("/");
    for _ in 0..(len - 1) {
        let mut r: f64 = rand::random();
        let a = ('a' as u8) as f64;
        r = (r * 26.0) + a;
        let c: char = r as u8 as char;
        s.push(c);
    }
    return s;
}

fn wayland_handle_message(
    fd: &mut wayland::Wayland,
    state: &mut wayland::State,
    msg: &mut &[u8],
) -> anyhow::Result<()> {
    if msg.len() < 8 {
        anyhow::bail!("msg too small for mesage header");
    }
    let object_id: u32 = msg.try_pop_ne()?;
    let opcode: u16 = msg.try_pop_ne()?;
    let msglen: u16 = msg.try_pop_ne()?;
    let payloadlen = msglen as usize - 8;

    let mut payload = &msg[0..payloadlen as usize];
    *msg = &msg[payloadlen..];

    if object_id == state.wl_registry && opcode == wayland::WAYLAND_WL_REGISTRY_EVENT_GLOBAL {
        let name: u32 = payload.try_pop_ne()?;
        let interface_len: u32 = payload.try_pop_ne()?;
        let padded_interface_len = wayland::nearest_multiple_4(interface_len as i32) as usize;

        let interface = CStr::from_bytes_until_nul(&payload[..padded_interface_len])?.to_str()?;
        payload = &payload[padded_interface_len..];

        let version: u32 = payload.try_pop_ne()?;
        println!(
            "<- wl_registry@{}.global: name={name} interface={interface} version={version}",
            state.wl_registry
        );

        use std::mem::size_of_val;
        assert_eq!(
            msglen as usize,
            size_of_val(&object_id)
                + size_of_val(&opcode)
                + size_of_val(&msglen)
                + size_of_val(&name)
                + size_of_val(&interface_len)
                + padded_interface_len
                + size_of_val(&version)
        );

        match interface {
            "wl_shm" => {
                state.wl_shm = fd.registry_bind(state.wl_registry, name, interface, version)?;
            }
            "xdg_wm_base" => {
                state.xdg_wm_base =
                    fd.registry_bind(state.wl_registry, name, interface, version)?;
            }
            "wl_compositor" => {
                state.wl_compositor =
                    fd.registry_bind(state.wl_registry, name, interface, version)?;
            }
            _ => (),
        };

        return Ok(());
    } else if object_id == wayland::WAYLAND_DISPLAY_OBJECT_ID
        && opcode == wayland::WAYLAND_WL_DISPLAY_ERROR_EVENT
    {
        let target_object_id: u32 = payload.try_pop_ne()?;
        let code: u32 = payload.try_pop_ne()?;
        let error_len: u32 = payload.try_pop_ne()?;
        let padded_err_len: usize = wayland::nearest_multiple_4(error_len as i32) as usize;

        let error = CStr::from_bytes_until_nul(&payload[..padded_err_len])?.to_str()?;
        //payload = &payload[padded_err_len..];
        anyhow::bail!(
            "fatal error: target_object_id={} code={} error={}",
            target_object_id,
            code,
            error
        );
    } else if object_id == state.xdg_wm_base && opcode == WAYLAND_XDG_BASE_EVENT_PING {
        let ping: u32 = payload.try_pop_ne()?;
        println!("<- xdg_wm_base@{}.ping: ping={}", state.xdg_wm_base, ping);
        fd.xdg_wm_base_pong(state, ping)?;
    } else if object_id == state.xdg_surface
        && opcode == wayland::WAYLAND_XDG_SURFACE_EVENT_CONFIGURE
    {
        let configure: u32 = payload.try_pop_ne()?;
        println!(
            "<- xdg_surface@{}.configure: configure={}",
            state.xdg_surface, configure
        );
        fd.xdg_surface_ack_configure(state, configure)?;
        // TODO: move into xdg_surface_ack_configure
        state.state = wayland::StateEnum::SurfaceAckedConfigure;
    }

    Ok(())
}

fn main() -> anyhow::Result<()> {
    let mut display = wayland::connect()?;
    println!("Connected to wayland display");
    let reg = display.get_registry()?;

    let mut state = wayland::State::default();
    state.wl_registry = reg;
    state.w = 117;
    state.h = 150;
    state.stride = state.w * wayland::COLOR_CHANNELS;

    // Single buffering.
    state.shm_pool_size = state.h * state.stride;

    let shm_path = gen_rand_name(255);
    let fd = shm::open(
        &shm_path,
        shm::OFlags::RDWR | shm::OFlags::EXCL | shm::OFlags::CREATE,
        Mode::WUSR | Mode::RUSR,
    )?;

    // file descriptor remains valid since our process still has the file open
    shm::unlink(shm_path)?;
    ftruncate(&fd, state.shm_pool_size as u64)?;

    let ptr = unsafe {
        mmap(
            null_mut(),
            state.shm_pool_size as usize,
            ProtFlags::READ | ProtFlags::WRITE,
            MapFlags::SHARED,
            &fd,
            0,
        )?
    };
    state.shm_pool_data = Some(ptr);
    state.shm_fd = Some(fd);

    let mut buf = [0_u8; 4096];
    loop {
        let read_bytes = display.sock.read(&mut buf)?;
        let mut msgs = &buf[..read_bytes];
        while msgs.len() > 0 {
            wayland_handle_message(&mut display, &mut state, &mut msgs)?;

            if state.wl_compositor != 0
                && state.wl_shm != 0
                && state.xdg_wm_base != 0
                && state.wl_surface == 0
            {
                assert_eq!(state.state, wayland::StateEnum::None);

                state.wl_surface = display.wl_compositor_create_surface(&mut state)?;
                state.xdg_surface = display.xdg_wm_base_get_xdg_surface(&mut state)?;
                state.xdg_toplevel = display.xdg_surface_get_toplevel(&mut state)?;
                display.wl_surface_commit(&mut state)?;
            }
            if state.state == wayland::StateEnum::SurfaceAckedConfigure {
                assert_ne!(state.wl_surface, 0);
                assert_ne!(state.xdg_surface, 0);
                assert_ne!(state.xdg_toplevel, 0);

                if state.wl_shm_pool == 0 {
                    state.wl_shm_pool = display.wl_shm_create_pool(&mut state)?;
                }
                if state.wl_buffer == 0 {
                    state.wl_buffer = display.wl_shm_pool_create_buffer(&mut state)?;
                }

                let n_pixels = state.w as usize * state.h as usize;
                let pixels;
                let ptr = state.shm_pool_data.unwrap() as *mut u32;
                unsafe {
                    pixels = std::slice::from_raw_parts_mut(ptr, n_pixels);
                }
                for pixel in pixels {
                    let r: u32 = 255;
                    let g: u32 = 0;
                    let b: u32 = 0;
                    *pixel = (r << 16) | (g << 8) | b;
                }

                display.wl_surface_attach(&mut state)?;
                display.wl_surface_commit(&mut state)?;
                state.state = wayland::StateEnum::SurfaceAttached;
            }
        }
    }

    //Ok(())
}