//! Provides utilities for headlessly running and taking screenshots of X11
//! programs in unit and integration tests.
//!
//! This uses [Xvfb](https://x.org/releases/X11R7.7/doc/man/man1/Xvfb.1.xhtml).
//!
//! Tests of this module invoke `xeyes` and `xwud` subprocesses. For all tests
//! to pass, you may need to install these X11 utilities and make them available
//! on your `PATH`.

use std::{
    ffi::CString,
    fs, io,
    process::{Child, Command, ExitStatus, Stdio},
    thread,
    time::{Duration, Instant},
};
use tempdir::TempDir;

use super::xwd;
use super::Lock;

/// Arguments for `Xvfb`. Use `XvfbArgs::default()` for default values.
#[derive(Clone, Copy, Debug)]
pub struct XvfbArgs {
    /// Width of the X11 screen.
    pub width: u32,
    /// Height of the X11 screen.
    pub height: u32,
    /// Bit depth of the X11 screen.
    ///
    /// This should probably be one of 8, 16, or 24. (And if you don't choose 24
    /// bits, you should have a good reason why.)
    pub depth: u8,
}

impl Default for XvfbArgs {
    fn default() -> Self {
        XvfbArgs {
            width: 640,
            height: 480,
            depth: 24,
        }
    }
}

/// A captive `Xvfb` process, with affordances for running subprocesses that
/// connect to it and taking screenshots.
///
/// When dropped, child processes will be killed automatically.
pub struct XvfbProcess {
    lock: Lock,
    framebuffer_directory: TempDir,
    process: Child,
    subprocesses: Vec<(Child, SubprocessMonitor)>,
    is_shutdown: bool,
}

impl XvfbProcess {
    /// Attempts to start `Xvfb` with default options.
    ///
    /// Returns `None` if an error occurs while starting the `Xvfb` process.
    pub fn start_default(lock: Lock) -> SubprocessResult<Self> {
        XvfbProcess::start(lock, XvfbArgs::default())
    }

    /// Starts an `Xvfb` process with options specified by `args`.
    ///
    /// Returns `None` if an error occurs while starting the `Xvfb` process.
    ///
    /// This function connects to the `Xvfb` server briefly to try to ensure
    /// that Xvfb is ready for clients to connect to it.
    pub fn start(lock: Lock, args: XvfbArgs) -> SubprocessResult<Self> {
        let framebuffer_directory = TempDir::new("wings_rs_xvfb").map_err(|e| {
            SubprocessError::new("Xvfb temp dir".into(), SubprocessErrorType::Spawn(e))
        })?;

        #[cfg(target_os = "linux")]
        unsafe {
            // Kill children if this parent process dies catastrophically (e.g.,
            // a test raises SIGABRT).
            libc::prctl(libc::PR_SET_PDEATHSIG, libc::SIGKILL);
        }
        let process = Command::new("Xvfb")
            .arg(format!(":{}", lock.display))
            .arg("-screen")
            .arg("0")
            .arg(format!("{}x{}x{}", args.width, args.height, args.depth))
            .arg("-fbdir")
            .arg(format!("{}", framebuffer_directory.path().display()))
            .stderr(Stdio::piped())
            .spawn()
            .map_err(|e| SubprocessError::new("Xvfb".into(), SubprocessErrorType::Spawn(e)))?;

        let process = XvfbProcess {
            lock,
            framebuffer_directory,
            process,
            subprocesses: Vec::new(),
            is_shutdown: false,
        };

        let now = Instant::now();
        let display_name = CString::new(
            format!(":{}", process.lock.display)
        ).expect("could not construct display name");
        while now.elapsed() < Duration::from_secs(1) {
            unsafe {
                let display = x11::xlib::XOpenDisplay(display_name.as_ptr());
                if !display.is_null() {
                    x11::xlib::XCloseDisplay(display);
                    break;
                }
            }
            thread::sleep(Duration::from_millis(10));
        }
        return Ok(process);
    }

    /// Returns the number in the `DISPLAY` environment variable to connect to
    /// this X11 server.
    pub fn display(&self) -> u16 {
        self.lock.display
    }

    /// Provides access to the `Xvfb` child process wrapped by `self`.
    pub fn process(&mut self) -> &mut Child {
        &mut self.process
    }

    /// Attempts to clean up all child processes. Returns any errors encountered, or `Ok` if none.
    ///
    /// Child processes are cleaned up with the following escalating steps:
    ///
    /// * Call [`SubprocessMonitor::on_xvfb_shutdown`] on each child's monitor.
    /// * Kill each of them directly (with [`Child::kill`]).
    /// * Kill the `Xvfb` process (which may cause any children still connected to that display server to exit).
    pub fn shutdown(&mut self) -> Result<(), Vec<SubprocessError>> {
        if self.is_shutdown {
            return Ok(());
        }
        self.is_shutdown = true;
        let mut errors = Vec::new();
        for (child, monitor) in self.subprocesses.iter_mut() {
            if let Err(e) = monitor.on_xvfb_shutdown(child) {
                errors.push(e);
            }
            match child.try_wait() {
                Ok(Some(status)) => {
                    if let Err(e) = monitor.check_subprocess_exit(status) {
                        errors.push(e);
                    }
                }
                Ok(None) => {
                    if let Err(e) = child.kill() {
                        errors.push(SubprocessError::new(
                            monitor.name.clone(),
                            SubprocessErrorType::Kill(e),
                        ));
                    }
                }
                Err(e) => errors.push(SubprocessError::new(
                    monitor.name.clone(),
                    SubprocessErrorType::Status(e),
                )),
            }
        }
        if let Err(e) = self.process.kill() {
            errors.push(SubprocessError::new(
                "Xvfb".into(),
                SubprocessErrorType::Kill(e),
            ));
        }
        if errors.is_empty() {
            Ok(())
        } else {
            Err(errors)
        }
    }

    /// Runs a simple command in this display server. No command line parsing is
    /// done, so `run("ls -la")` will probably not do what you expect (because
    /// there is no binary on your `PATH` called `ls -la`). Use [`XvfbProcess::run_args`]
    /// to provide arguments to `command`.
    pub fn run(&mut self, command: impl Into<String>) -> Result<(), SubprocessError> {
        let name = command.into();
        self.run_args(name.clone(), &name, |_c| {})
    }

    /// Runs a command that you can provide arguments to in this display
    /// server.
    ///
    /// The `f` parameter can be used to modify the `Command` before it is
    /// run (e.g., by adding arguments). For example:
    ///
    /// ```
    /// # use wings_rs_tests::headless::DisplayNumberRegistry;
    /// # use wings_rs_tests::headless::xvfb::XvfbProcess;
    /// # fn main() {
    /// # let mut xvfb = XvfbProcess::start_default(DisplayNumberRegistry::next_unused_display().unwrap()).unwrap();
    /// xvfb.run_args(
    ///     "demo of xeyes, angel-style",
    ///     "/usr/bin/xeyes",
    ///     |cmd| { cmd.arg("-biblicallyAccurate"); },
    /// ).unwrap();
    /// # }
    /// ```
    pub fn run_args(
        &mut self,
        name: impl Into<String>,
        command: impl AsRef<std::ffi::OsStr>,
        f: impl FnOnce(&mut Command),
    ) -> Result<(), SubprocessError> {
        let mut command = Command::new(command);
        command.env("DISPLAY", format!(":{}", self.lock.display));
        f(&mut command);
        let mut monitor = SubprocessMonitor::new(name.into(), command);
        monitor.with_subprocess_exit(SubprocessMonitor::require_clean_exit());
        let child = monitor.start()?;
        self.subprocesses.push((child, monitor));
        Ok(())
    }

    /// Checks all subprocesses started with the [`XvfbProcess::run`] family of
    /// methods and returns errors for any that have an error status, or `Ok` if
    /// none do.
    pub fn check(&mut self) -> Result<(), Vec<SubprocessError>> {
        let errors: Vec<_> = self
            .subprocesses
            .iter_mut()
            .filter_map(|(child, monitor)| match child.try_wait() {
                Ok(Some(status)) => monitor.check_subprocess_exit(status).err(),
                _ => None,
            })
            .collect();
        if errors.is_empty() {
            Ok(())
        } else {
            Err(errors)
        }
    }

    /// Returns a screenshot of the current framebuffer contents in `xwd(1)`
    /// format.
    ///
    /// For the same in a more generally useful format, try [`XvfbProcess::png_screenshot`].
    pub fn xwd_screenshot(&self) -> io::Result<xwd::XwdImage> {
        let path = self.framebuffer_directory.path().join("Xvfb_screen0");
        Ok(xwd::XwdImage::read(&mut fs::read(path)?.into_iter()))
    }

    /// Returns a screenshot of the current framebuffer contents in PNG format.
    ///
    /// Panics on errors.
    pub fn png_screenshot(&self) -> Vec<u8> {
        self.xwd_screenshot().unwrap().into_png().unwrap()
    }
}

impl Drop for XvfbProcess {
    fn drop(&mut self) {
        self.shutdown().unwrap();
    }
}

pub type SubprocessResult<T> = Result<T, SubprocessError>;

/// Errors that may occur when a subprocess fails to run or exit cleanly.
#[derive(Debug)]
pub struct SubprocessError {
    /// Meaningful name that can be used to determine what subprocess failed.
    pub name: String,
    pub err: SubprocessErrorType,
}

impl SubprocessError {
    pub fn new(name: String, err: SubprocessErrorType) -> Self {
        SubprocessError { name, err }
    }
}

/// Error details for `SubprocessError`.
#[derive(Debug)]
pub enum SubprocessErrorType {
    /// Subprocess exited with bad exit code.
    BadExit(ExitStatus),
    /// Could not spawn child process (with `io::Error` as the reason).
    Spawn(io::Error),
    /// Could not kill child process (with `io::Error` as the reason).
    Kill(io::Error),
    /// Could not query child process for exit status (with `io::Error` as the reason).
    Status(io::Error),
}

/// Manages spawning a subprocess, monitoring it during execution, and killing it when `Xvfb` shuts down.
pub struct SubprocessMonitor {
    name: String,
    command: Command,
    check_subprocess_exit: Option<Box<dyn FnMut(&str, ExitStatus) -> SubprocessResult<()>>>,
    on_xvfb_shutdown: Option<Box<dyn FnMut(&mut Child) -> SubprocessResult<()>>>,
}

impl SubprocessMonitor {
    /// Creates a new monitor but does not start any subprocess yet.
    ///
    /// This monitor may be configured further by calling mutators like
    /// [`SubprocessMonitor::with_subprocess_exit`]. Once configuration is
    /// completely, start the subprocess with [`SubprocessMonitor::start`].
    pub fn new(name: impl Into<String>, command: Command) -> Self {
        SubprocessMonitor {
            name: name.into(),
            command,
            check_subprocess_exit: None,
            on_xvfb_shutdown: None,
        }
    }

    /// Sets the callback to execute when [`SubprocessMonitor::check_subprocess_exit`] is called.
    pub fn with_subprocess_exit(
        &mut self,
        f: Box<dyn FnMut(&str, ExitStatus) -> SubprocessResult<()>>,
    ) -> &mut Self {
        self.check_subprocess_exit = Some(f);
        self
    }

    /// Sets the callback to execute when
    /// [`SubprocessMonitor::on_xvfb_shutdown`] is called. It may check process
    /// status, clean up, etc.
    pub fn with_xvfb_shutdown(
        &mut self,
        f: Box<dyn FnMut(&mut Child) -> SubprocessResult<()>>,
    ) -> &mut Self {
        self.on_xvfb_shutdown = Some(f);
        self
    }

    /// Attempts to spawn a child process from the `Command` provided to
    /// [`SubprocessMonitor::new`].
    pub fn start(&mut self) -> SubprocessResult<Child> {
        self.command
            .spawn()
            .map_err(|e| SubprocessError::new(self.name.clone(), SubprocessErrorType::Spawn(e)))
    }

    /// Called when subprocess exit status is polled. This may be used to signal
    /// an error because the exit status is bad.
    pub fn check_subprocess_exit(&mut self, status: ExitStatus) -> SubprocessResult<()> {
        if let Some(f) = &mut self.check_subprocess_exit {
            (f)(&self.name, status)
        } else {
            Ok(())
        }
    }

    /// Called immediately before the parent display server is being shut
    /// down. This may be used to send `SIGINT` or otherwise clean up the child
    /// before it is killed more forcibly.
    pub fn on_xvfb_shutdown(&mut self, child: &mut Child) -> SubprocessResult<()> {
        if let Some(f) = &mut self.on_xvfb_shutdown {
            (f)(child)
        } else {
            Ok(())
        }
    }

    /// Returns a callback suitable for passing to [`SubprocessMonitor::check_subprocess_exit`]
    /// which requires that the process has exited cleanly.
    pub fn require_clean_exit() -> Box<dyn FnMut(&str, ExitStatus) -> SubprocessResult<()>> {
        Box::new(|name: &str, status: ExitStatus| {
            if status.success() {
                Ok(())
            } else {
                Err(SubprocessError::new(
                    String::from(name),
                    SubprocessErrorType::BadExit(status),
                ))
            }
        })
    }
}

#[cfg(test)]
mod tests {
    use crate::headless::DisplayNumberRegistry;
    use super::XvfbProcess;
    use insta_image::assert_png_snapshot;
    use std::{
        thread,
        time::Duration,
    };

    /// Starts Xvfb and returns a managing object, or panics.
    fn start_xvfb() -> XvfbProcess {
        XvfbProcess::start_default(
            DisplayNumberRegistry::next_unused_display().expect("cannot find a value for DISPLAY"),
        )
        .expect("cannot start Xvfb")
    }

    /// Reads a single frame of PNG data from the PNG image in `data`. Panics on
    /// errors.
    fn read_png_metadata(data: &[u8]) -> png::OutputInfo {
        let mut reader = png::Decoder::new(std::io::Cursor::new(data))
            .read_info()
            .unwrap();
        let mut buf = vec![0; reader.output_buffer_size().unwrap()];
        reader.next_frame(&mut buf).unwrap()
    }

    /// Launches `xeyes(1)` with colors to make sure we're doing the right thing
    /// with our color channels when encoding a screenshot tas PNG.
    fn xeyes_color(xvfb: &mut XvfbProcess) {
        xvfb.run_args("xeyes", "xeyes", |c| {
            c.arg("-outline")
                .arg("blue")
                .arg("-center")
                .arg("red")
                .arg("-fg")
                .arg("green");
        })
        .unwrap();
        thread::sleep(Duration::from_millis(100));
    }

    #[test]
    fn png_blank_screenshot() {
        let xvfb = start_xvfb();
        let xwd = xvfb.xwd_screenshot().unwrap();

        assert_eq!(
            xwd.visual_class(),
            crate::headless::xwd::VisualClass::TrueColor
        );

        let args = crate::headless::xvfb::XvfbArgs::default();
        assert_eq!(xwd.header.depth, args.depth as u32);
        assert_eq!(xwd.header.width, args.width);
        assert_eq!(xwd.header.height, args.height);

        assert_png_snapshot!("blank_screen", xwd.into_png().unwrap());
    }

    #[test]
    fn png_xeyes_screenshot() {
        let mut xvfb = start_xvfb();
        xeyes_color(&mut xvfb);

        let compressed = xvfb.png_screenshot();
        let image_data = read_png_metadata(&compressed);

        let args = crate::headless::xvfb::XvfbArgs::default();
        assert_eq!(image_data.width, args.width);
        assert_eq!(image_data.height, args.height);

        assert_png_snapshot!("xeyes", compressed);
    }

    #[test]
    fn png_lamp_image() {
        let mut xvfb = start_xvfb();
        xvfb.run_args("xwud", "xwud", |c| {
            c.arg("-in").arg("src/headless/snowlamp.xwd");
        })
        .unwrap();
        thread::sleep(Duration::from_millis(100));

        let compressed = xvfb.png_screenshot();

        assert_png_snapshot!("snowlamp_in_window", compressed);
    }
}
