use std::{
    env,
    ffi::OsStr,
    os::fd::{BorrowedFd, OwnedFd},
    os::unix::{
        io::{AsRawFd, RawFd},
        net::UnixStream,
        process::CommandExt,
    },
    process::{Child, Command},
    sync::{Arc, Mutex},
    thread,
};

use tracing::{error, info, trace};
use wayland_server::backend::{ClientData, ClientId, DisconnectReason};
use wayland_server::{Client, DisplayHandle};

use crate::{utils::user_data::UserDataMap, wayland::compositor::CompositorClientState};

use super::x11_sockets::{prepare_x11_sockets, X11Lock};

/// A handle to a running XWayland process. Using XWayland as an xserver for
/// X11-based clients requires two connections: one wayland socket, where
/// XWayland creates surfaces for its clients, and one X11 socket, where a
/// compositor's implementation of an X11 window manager connects and handles
/// X11 events.
///
/// ```text
///            ┌───────────────────┐
///            │                   │
///            │                   │ X11 Clients
///            │     Xwayland      │ ◄─────
///            │                   │
///            │                   │
///            └──┬──────────────▲─┘
///               │              │
/// Wayland Socket│              │ X11 Socket
///               │              │
///               │              │
///  ┌────────────▼──────────────┼────────────┐
///  │  Compositor               │            │
///  │                         ┌─┴───────┐    │
///  │                         │ XWM     │    │
///  │                         │         │    │
///  │                         └─────────┘    │
///  │                                        │
///  └────────────────────────────────────────┘
/// ```
///
/// This struct handles integrating the XWayland process itself into the event
/// loop, but a [separate X11 window manager implementation](crate::xwayland::xwm::X11Wm)
/// is needed as well to support X11 clients.
///
/// To shut down the instance, dropping this handle is generally sufficient,
/// along with the X11 connection that was passed to the window manager
/// implementation, if any. The process will die once the connections to it
/// are closed.
#[derive(Debug)]
pub struct XWayland {
    inner: Arc<Mutex<Instance>>,
    source: calloop::generic::Generic<calloop::generic::FdWrapper<RawFd>>,

    // So we can disconnect the client on drop.
    dh: DisplayHandle,
    client: Client,
}

/// Events generated by an [XWayland] instance.
#[derive(Debug)]
pub enum XWaylandEvent {
    /// The XWayland server is ready
    Ready {
        /// A privileged X11 connection to XWayland.
        x11_socket: UnixStream,

        /// The display number the XWayland server is using.
        ///
        /// This can be useful to set the `DISPLAY` variable manually when
        /// spawning processes that may use XWayland.
        display_number: u32,
    },

    /// The XWayland server exited unexpectedly during startup.
    Error,
}

impl XWayland {
    /// Spawns an XWayland server instance. `Xwayland` must be on the `PATH` and
    /// executable.
    ///
    /// ## Arguments
    ///
    /// - `display` - if provided, only the given display number will be tested.
    ///   If you wish smithay to choose a display for you, pass `None`.
    /// - `envs` - Allows additionally environment variables to be set when
    ///   launching XWayland.
    /// - `open_abstract_socket` - Open an abstract socket as well as filesystem
    ///   sockets (only on available on Linux).
    /// - `stdout, stderr` - Allows redirecting stdout and stderr of the
    ///   XWayland process. XWayland output is rarely useful, so `Stdio::null()`
    ///   is a good choice if you're not sure.
    /// - `user_data` - Allows mutating the `XWaylandClientData::user_data`-map
    ///   before the client is added to the wayland display. Useful for
    ///   initializing state for global filters.
    ///
    /// Returns a handle to the XWayland instance and the
    /// [Client](wayland_server::Client) representing the XWayland server. The
    /// handle can be inserted in your event loop, and If everything goes well,
    /// you'll eventually receive an `XWaylandEvent::Ready`, indicating that
    /// it's time to start the X11 window manager.
    pub fn spawn<K, V, I, F>(
        dh: &DisplayHandle,
        display: impl Into<Option<u32>>,
        envs: I,
        open_abstract_socket: bool,
        stdout: impl Into<std::process::Stdio>,
        stderr: impl Into<std::process::Stdio>,
        user_data: F,
    ) -> std::io::Result<(Self, Client)>
    where
        I: IntoIterator<Item = (K, V)>,
        K: AsRef<OsStr>,
        V: AsRef<OsStr>,
        F: FnOnce(&UserDataMap),
    {
        let (x_wm_x11, x_wm_me) = UnixStream::pair()?;
        let (wl_x11, wl_me) = UnixStream::pair()?;

        let (lock, listen_sockets) = prepare_x11_sockets(display.into(), open_abstract_socket)?;
        let display_number = lock.display_number();

        // XWayland writes the the display number and a newline to this pipe when it's ready.
        let (displayfd_recv, displayfd_send) =
            rustix::pipe::pipe_with(rustix::pipe::PipeFlags::NONBLOCK | rustix::pipe::PipeFlags::CLOEXEC)?;

        let mut command = Command::new("Xwayland");

        command
            .stdout(stdout)
            .stderr(stderr)
            .arg(format!(":{}", display_number))
            .arg("-verbose")
            .arg("-rootless")
            .arg("-terminate")
            .arg("-wm")
            .arg(x_wm_x11.as_raw_fd().to_string())
            .arg("-displayfd")
            .arg(displayfd_send.as_raw_fd().to_string());

        for socket in &listen_sockets {
            command.arg("-listenfd").arg(socket.as_raw_fd().to_string());
        }

        // Setup the environment; clear everything except PATH and XDG_RUNTIME_DIR.
        command.env_clear();
        for (key, value) in env::vars_os() {
            if key.to_str() == Some("PATH") || key.to_str() == Some("XDG_RUNTIME_DIR") {
                command.env(key, value);
                continue;
            }
        }

        command.env("WAYLAND_SOCKET", format!("{}", wl_x11.as_raw_fd()));
        command.envs(envs);

        unsafe {
            let wayland_socket_fd = wl_x11.as_raw_fd();
            let wm_socket_fd = x_wm_x11.as_raw_fd();
            let pipe_fd = displayfd_send.as_raw_fd();
            let socket_fds: Vec<_> = listen_sockets.iter().map(|socket| socket.as_raw_fd()).collect();

            command.pre_exec(move || {
                // unset the CLOEXEC flag from the sockets we need to pass
                // to xwayland.
                unset_cloexec(wayland_socket_fd)?;
                unset_cloexec(wm_socket_fd)?;
                unset_cloexec(pipe_fd)?;
                for &socket in socket_fds.iter() {
                    unset_cloexec(socket)?;
                }

                Ok(())
            });
        }

        info!("spawning XWayland instance");

        let child = command.spawn()?;

        // SAFETY: RawFd's AsRawFd impl is infallible.
        let wrapper = unsafe { calloop::generic::FdWrapper::new(displayfd_recv.as_raw_fd()) };
        let source = calloop::generic::Generic::new(wrapper, calloop::Interest::READ, calloop::Mode::Level);
        let inner = Instance {
            display_lock: lock,
            display_fd: displayfd_recv,
            x11_socket: Some(x_wm_me),
        };

        let data_map = UserDataMap::new();
        user_data(&data_map);

        // Insert the client into the display handle. The order is important
        // here; XWayland never starts up at all unless it can roundtrip with
        // wayland.
        let inner = Arc::new(Mutex::new(inner));
        let mut dh = dh.clone();
        let client = dh.insert_client(
            wl_me,
            Arc::new(XWaylandClientData {
                #[cfg(feature = "wayland_frontend")]
                compositor_state: CompositorClientState::default(),
                data_map,
                child: Mutex::new(Some(child)),
            }),
        )?;

        Ok((
            Self {
                inner,
                source,
                dh,
                client: client.clone(),
            },
            client,
        ))
    }

    /// Returns the X11 display used by the instance, suitable for setting the
    /// `DISPLAY` environment variable.
    pub fn display_number(&self) -> u32 {
        self.inner.lock().unwrap().display_lock.display_number()
    }

    /// Returns a file descriptor which can be polled for readiness. When the fd
    /// is readable, the XWayland server's readiness can be checked with
    /// [take_socket](Self::take_socket).
    pub fn poll_fd(&self) -> BorrowedFd<'_> {
        let guard = self.inner.lock().unwrap();

        // SAFETY: we never mutate display_fd, and it lives as long as the instance.
        unsafe { BorrowedFd::borrow_raw(guard.display_fd.as_raw_fd()) }
    }

    /// Checks if the XWayland instance is ready. If XWayland has fully started,
    /// this will return the X11 socket connected to the running instance.
    ///
    /// Calling `take_socket` successfully transfers ownership of the connection
    /// to the caller. After returning `Some` the first time, it will always
    /// return `None`.
    ///
    /// An `Err` result is only returned if the XWayland instance has exited
    /// unexpectedly.
    ///
    /// This is a low-level method. Using the instance as an event source is the
    /// recommended way to interact with it.
    pub fn take_socket(&mut self) -> std::io::Result<Option<UnixStream>> {
        self.inner.lock().unwrap().take_socket()
    }
}

#[derive(Debug)]
struct Instance {
    display_lock: X11Lock,
    x11_socket: Option<UnixStream>,
    display_fd: OwnedFd,
}

impl calloop::EventSource for XWayland {
    type Event = XWaylandEvent;
    type Metadata = ();
    type Ret = ();
    type Error = std::io::Error;

    #[profiling::function]
    fn process_events<F>(
        &mut self,
        readiness: calloop::Readiness,
        token: calloop::Token,
        mut callback: F,
    ) -> std::io::Result<calloop::PostAction>
    where
        F: FnMut(Self::Event, &mut Self::Metadata) -> Self::Ret,
    {
        let mut guard = self.inner.lock().unwrap();

        self.source.process_events(readiness, token, |_, _| {
            let x11_socket = match guard.take_socket() {
                Ok(Some(sockets)) => sockets,
                Ok(None) => return Ok(calloop::PostAction::Continue),
                Err(_) => {
                    callback(XWaylandEvent::Error, &mut ());
                    return Ok(calloop::PostAction::Disable);
                }
            };

            callback(
                XWaylandEvent::Ready {
                    x11_socket,
                    display_number: guard.display_lock.display_number(),
                },
                &mut (),
            );

            Ok(calloop::PostAction::Disable)
        })
    }

    fn register(
        &mut self,
        poll: &mut calloop::Poll,
        factory: &mut calloop::TokenFactory,
    ) -> calloop::Result<()> {
        self.source.register(poll, factory)
    }

    fn reregister(
        &mut self,
        poll: &mut calloop::Poll,
        factory: &mut calloop::TokenFactory,
    ) -> calloop::Result<()> {
        self.source.reregister(poll, factory)
    }

    fn unregister(&mut self, poll: &mut calloop::Poll) -> calloop::Result<()> {
        self.source.unregister(poll)
    }
}

impl Drop for XWayland {
    fn drop(&mut self) {
        self.dh
            .backend_handle()
            .kill_client(self.client.id(), DisconnectReason::ConnectionClosed);
    }
}

impl Instance {
    fn take_socket(&mut self) -> std::io::Result<Option<UnixStream>> {
        trace!("checking for XWayland readiness");

        if self.x11_socket.is_none() {
            return Ok(None);
        }

        let mut buf = [0; 64];
        loop {
            let res = rustix::io::read(&self.display_fd, &mut buf);
            trace!(?res, "read from XWayland displayfd");

            match res {
                Ok(0) => return Ok(None),
                Ok(len) if (buf[..len]).contains(&b'\n') => return Ok(self.x11_socket.take()),
                Err(err) if err.kind() == std::io::ErrorKind::WouldBlock => return Ok(None),
                Err(err) => return Err(err.into()),
                _ => continue,
            }
        }
    }
}

/// Inner `ClientData`-type of an xwayland client.
#[derive(Debug)]
pub struct XWaylandClientData {
    /// client state of the [`crate::wayland::compositor`] module
    #[cfg(feature = "wayland_frontend")]
    pub compositor_state: CompositorClientState,
    data_map: UserDataMap,
    child: Mutex<Option<Child>>,
}

impl ClientData for XWaylandClientData {
    fn disconnected(&self, _client_id: ClientId, reason: DisconnectReason) {
        if let DisconnectReason::ProtocolError(err) = reason {
            error!("Xwayland disconnected: {}", err);
        }

        let mut child = self.child.lock().unwrap().take().unwrap();
        thread::spawn(move || {
            if let Ok(status) = child.wait() {
                if !status.success() {
                    error!("Xwayland terminated: {}", status);
                }
            }
        });
    }
}

impl XWaylandClientData {
    /// Access user_data map for a xwayland client
    pub fn user_data(&self) -> &UserDataMap {
        &self.data_map
    }
}

/// Removes the `O_CLOEXEC` flag from a `RawFd`, causing it to leak when we
/// `exec` XWayland
unsafe fn unset_cloexec(fd: RawFd) -> std::io::Result<()> {
    let fd = BorrowedFd::borrow_raw(fd);
    rustix::io::fcntl_setfd(fd, rustix::io::FdFlags::empty())?;
    Ok(())
}