A KVM switch emulator using UDP/IP
use crate::io::*;

use futures::{ready, Sink, Stream};
use log::warn;
use tokio::io::PollEvented;
use x11_dl::{
    xfixes,
    xinput2,
    xlib,
    xtest,
};

use std::cell::Cell;
use std::io;
use std::mem::MaybeUninit;
use std::os::raw::{c_char, c_int, c_uchar, c_uint};
use std::os::unix::io::RawFd;
use std::pin::Pin;
use std::ptr;
use std::task::{Context, Poll};

pub struct Host {
    xlib: xlib::Xlib,
    display: *mut xlib::Display,
    io: PollEvented<Connection>,
    root: xlib::Window,

    xfixes: xfixes::XFixes,
    xfixes_event_base: c_int,
    clipboard: xlib::Atom,

    xinput_major_opcode: c_int,

    xtest: xtest::Xf86vmode,

    last_pos: Cell<(c_int, c_int)>,
    grabbed: Cell<bool>,
}

impl Host {
    pub fn open() -> Result<Self, x11_dl::error::OpenError> {
        unsafe {
            let xlib = xlib::Xlib::open()?;
            let display = (xlib.XOpenDisplay)(ptr::null());
            if display.is_null() {
                panic!("Failed to open display");
            }

            let io = PollEvented::new(Connection {
                fd: (xlib.XConnectionNumber)(display)
            }).unwrap();

            let root = (xlib.XDefaultRootWindow)(display);

            // Query XFixes
            let xfixes = xfixes::XFixes::open()?;
            let xfixes_event_base = {
                let mut event_base = MaybeUninit::uninit();
                let mut error_base = MaybeUninit::uninit();
                if (xfixes.XFixesQueryExtension)(
                    display,
                    event_base.as_mut_ptr(),
                    error_base.as_mut_ptr(),
                ) == xlib::False {
                    panic!("Failed to query XFixes");
                }

                event_base.assume_init()
            };

            // Setup selection events
            (xfixes.XFixesSelectSelectionInput)(
                display,
                root,
                xlib::XA_PRIMARY,
                xfixes::XFixesSetSelectionOwnerNotifyMask
            );

            let clipboard = (xlib.XInternAtom)(
                display,
                b"CLIPBOARD\0".as_ptr() as *const c_char,
                0
            );

            (xfixes.XFixesSelectSelectionInput)(
                display,
                root,
                clipboard,
                xfixes::XFixesSetSelectionOwnerNotifyMask);

            // Query XInput2
            let xinput2 = xinput2::XInput2::open()?;
            let xinput_major_opcode = {
                let mut major_opcode = MaybeUninit::uninit();
                let mut first_event = MaybeUninit::uninit();
                let mut first_error = MaybeUninit::uninit();
                if (xlib.XQueryExtension)(
                    display,
                    b"XInputExtension\0".as_ptr() as *const c_char,
                    major_opcode.as_mut_ptr(),
                    first_event.as_mut_ptr(),
                    first_error.as_mut_ptr(),
                ) == xlib::False {
                    panic!("Failed to query XInputExtension");
                };

                major_opcode.assume_init()
            };

            // Setup raw motion events
            let mut mask = [0; (xinput2::XI_LASTEVENT as usize + 7) / 8];
            xinput2::XISetMask(&mut mask, xinput2::XI_RawMotion);

            let mut events = [xinput2::XIEventMask {
                deviceid: xinput2::XIAllMasterDevices,
                mask_len: mask.len() as c_int,
                mask: mask.as_mut_ptr(),
            }];

            (xinput2.XISelectEvents)(
                display, root,
                &mut events[0] as *mut xinput2::XIEventMask, events.len() as c_int
            );

            let xtest = xtest::Xf86vmode::open()?;

            let host = Self {
                xlib,
                display,
                io,
                root,
                clipboard,

                xfixes,
                xfixes_event_base,

                xinput_major_opcode,

                xtest,

                last_pos: Cell::new((0, 0)),
                grabbed: Cell::new(false),
            };

            host.last_pos.set(host.cursor_pos());
            Ok(host)
        }
    }

    fn map_button_event(event: xlib::XButtonEvent, state: bool) -> ButtonEvent {
        ButtonEvent {
            button: event.button,
            state,
        }
    }

    fn map_key_event(&self, event: xlib::XKeyEvent, state: bool) -> KeyEvent {
        let key = unsafe {
            (self.xlib.XKeycodeToKeysym)(self.display, event.keycode as c_uchar, 0)
        };

        KeyEvent {
            key,
            state,
        }
    }

    fn map_generic_event(&self, cookie: xlib::XGenericEventCookie) -> Option<HostOutputEvent> {
        if cookie.extension == self.xinput_major_opcode {
            match cookie.evtype {
                xinput2::XI_RawMotion => Some(HostOutputEvent::Motion(self.map_raw_motion())),
                _ => None,
            }
        } else {
            None
        }
    }

    fn map_raw_motion(&self) -> MotionEvent {
        let (x, y) = self.cursor_pos();
        let (last_x, last_y) = self.last_pos.get();
        let (dx, dy) = (x - last_x, y - last_y);
        self.last_pos.set((x, y));

        // Lock cursor to center when grabbed
        if self.grabbed.get() {
            let (width, height) = self.screen_size();
            let (x, y) = (width / 2, height / 2);
            self.start_send_position_event(PositionEvent { x, y });
        }

        MotionEvent { dx, dy }
    }

    fn map_selection_event(&self, event: xfixes::XFixesSelectionNotifyEvent) -> Option<SelectionEvent> {
        match event.subtype {
            xfixes::XFixesSetSelectionOwnerNotify => {
                if event.selection == xlib::XA_PRIMARY {
                    Some(SelectionEvent::Primary)
                } else if event.selection == self.clipboard {
                    Some(SelectionEvent::Clipboard)
                } else {
                    warn!("Unexpected selection source: {}", event.selection);
                    None
                }
            },
            subtype => {
                warn!("Unexpected XFixesSelection sub event: {}", subtype);
                None
            }
        }
    }

    fn start_send_position_event(&self, event: PositionEvent) {
        self.last_pos.set((event.x, event.y));
        unsafe {
            (self.xlib.XWarpPointer)(self.display, 0, self.root, 0, 0, 0, 0, event.x, event.y);
        }
    }

    fn start_send_grab_event(&self, event: GrabEvent) {
        if event.grab {
            if self.grabbed.get() {
                return;
            }

            let mask = (xlib::ButtonPressMask | xlib::ButtonReleaseMask) as c_uint;

            unsafe {
                (self.xlib.XGrabPointer)(
                    self.display,
                    self.root,
                    xlib::True,
                    mask,
                    xlib::GrabModeAsync,
                    xlib::GrabModeAsync,
                    0,
                    0,
                    xlib::CurrentTime,
                );

                (self.xlib.XGrabKeyboard)(
                    self.display,
                    self.root,
                    xlib::True,
                    xlib::GrabModeAsync,
                    xlib::GrabModeAsync,
                    xlib::CurrentTime,
                );

                (self.xfixes.XFixesHideCursor)(self.display, self.root);
            }

            self.grabbed.set(true);
        } else {
            if !self.grabbed.get() {
                return;
            }

            unsafe {
                (self.xlib.XUngrabPointer)(self.display, xlib::CurrentTime);
                (self.xlib.XUngrabKeyboard)(self.display, xlib::CurrentTime);
                (self.xfixes.XFixesShowCursor)(self.display, self.root);
            }

            self.grabbed.set(false);
        }
    }

    fn start_send_button_event(&self, event: ButtonEvent) {
        unsafe {
            (self.xtest.XTestFakeButtonEvent)(self.display, event.button, event.state as c_int, xlib::CurrentTime);
        }
    }

    fn start_send_key_event(&self, event: KeyEvent) {
        unsafe {
            let keycode = (self.xlib.XKeysymToKeycode)(self.display, event.key);
            (self.xtest.XTestFakeKeyEvent)(self.display, keycode as c_uint, event.state as c_int, 0);
        }
    }
}

impl Drop for Host {
    fn drop(&mut self) {
        unsafe {
            (self.xlib.XCloseDisplay)(self.display);
        }
    }
}

impl HostInterface for Host {
    fn screen_size(&self) -> (i32, i32) {
        let screen = unsafe {
            // TODO: Test if flushes output buffer
            &*(self.xlib.XDefaultScreenOfDisplay)(self.display)
        };

        assert!(screen.width > 0 && screen.height > 0);
        (screen.width, screen.height)
    }

    fn cursor_pos(&self) -> (i32, i32) {
        unsafe {
            let mut root = MaybeUninit::uninit();
            let mut child = MaybeUninit::uninit();
            let mut root_x = MaybeUninit::uninit();
            let mut root_y = MaybeUninit::uninit();
            let mut child_x = MaybeUninit::uninit();
            let mut child_y = MaybeUninit::uninit();
            let mut mask = MaybeUninit::uninit();

            // NOTE: Blocks sending & receiving response from X server
            (self.xlib.XQueryPointer)(
                self.display, self.root,
                root.as_mut_ptr(), child.as_mut_ptr(),
                root_x.as_mut_ptr(), root_y.as_mut_ptr(),
                child_x.as_mut_ptr(), child_y.as_mut_ptr(),
                mask.as_mut_ptr());

            (root_x.assume_init(), root_y.assume_init())
        }
    }
}

impl Stream for Host {
    type Item = io::Result<HostOutputEvent>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
        ready!(self.io.poll_read_ready(cx, mio::Ready::readable()))?;

        loop {
            // TODO: Add XEventsQueued constants to xlib
            #[allow(non_upper_case_globals)]
            const QueuedAfterReading: c_int = 1;

            let num_events = unsafe {
                (self.xlib.XEventsQueued)(self.display, QueuedAfterReading)
            };

            if num_events <= 0 {
                self.io.clear_read_ready(cx, mio::Ready::readable())?;
                return Poll::Pending;
            }

            let event = unsafe {
                let mut event = MaybeUninit::uninit();
                (self.xlib.XNextEvent)(self.display, event.as_mut_ptr());
                event.assume_init()
            };

            let host_event = match event.get_type() {
                // Standard events
                xlib::ButtonPress => Some(HostOutputEvent::Button(Self::map_button_event(event.into(), true))),
                xlib::ButtonRelease => Some(HostOutputEvent::Button(Self::map_button_event(event.into(), false))),
                xlib::KeyPress => Some(HostOutputEvent::Key(self.map_key_event(event.into(), true))),
                xlib::KeyRelease => Some(HostOutputEvent::Key(self.map_key_event(event.into(), false))),
                xlib::MappingNotify => None,
                xlib::GenericEvent => self.map_generic_event(event.into()),

                // XFixes selection events
                event_type if event_type - self.xfixes_event_base == xfixes::XFixesSelectionNotify =>
                    self.map_selection_event(event.into()).map(HostOutputEvent::Selection),

                event_type => {
                    warn!("Unexpected X11 event: {}", event_type);
                    None
                }
            };

            if let Some(host_event) = host_event {
                return Poll::Ready(Some(Ok(host_event)));
            }
        }
    }
}

impl Sink<HostInputEvent> for Host {
    type Error = io::Error;

    fn poll_ready(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), Self::Error>> {
        ready!(self.io.poll_write_ready(cx))?;
        Poll::Ready(Ok(()))
    }

    fn start_send(self: Pin<&mut Self>, event: HostInputEvent) -> Result<(), Self::Error> {
        match event {
            HostInputEvent::Position(event) => self.start_send_position_event(event),
            HostInputEvent::Grab(event) => self.start_send_grab_event(event),
            HostInputEvent::Button(event) => self.start_send_button_event(event),
            HostInputEvent::Key(event) => self.start_send_key_event(event),
            HostInputEvent::Selection(_event) => todo!(),
        }

        Ok(())
    }

    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), Self::Error>> {
        match self.io.poll_write_ready(cx) {
            Poll::Ready(result) => Poll::Ready(result.map(|_| unsafe {
                // NOTE: Blocks until output buffer is completely flushed,
                // but let's assume the X server will always be responsive.
                // If it isn't, we have bigger problems.
                (self.xlib.XFlush)(self.display);
            })),
            Poll::Pending => Poll::Pending,
        }
    }

    fn poll_close(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), Self::Error>> {
        self.poll_flush(cx)
    }
}

struct Connection {
    fd: RawFd,
}

impl mio::Evented for Connection {
    fn register(&self, poll: &mio::Poll, token: mio::Token, interest: mio::Ready, opts: mio::PollOpt)
        -> io::Result<()>
    {
        mio::unix::EventedFd(&self.fd).register(poll, token, interest, opts)
    }

    fn reregister(&self, poll: &mio::Poll, token: mio::Token, interest: mio::Ready, opts: mio::PollOpt)
        -> io::Result<()>
    {
        mio::unix::EventedFd(&self.fd).reregister(poll, token, interest, opts)
    }

    fn deregister(&self, poll: &mio::Poll) -> io::Result<()> {
        mio::unix::EventedFd(&self.fd).deregister(poll)
    }
}