//! Control a container

use std::{
  ffi::OsString,
  process::{Command, ExitStatus, Stdio},
};

use shared_child::SharedChild;

use crate::{Error, Result};

/// Network mode, refer to the [docker documentation](https://docs.docker.com/engine/network/)
/// for an explanation on the different modes.
pub enum Network
{
  /// The default network driver.
  Bridge,
  /// Remove network isolation between the container and the Docker host.
  Host,
  /// Completely isolate a container from the host and other containers.
  None,
}

//   ____             __ _                       _
//  / ___|___  _ __  / _(_) __ _ _   _ _ __ __ _| |_ ___  _ __
// | |   / _ \| '_ \| |_| |/ _` | | | | '__/ _` | __/ _ \| '__|
// | |__| (_) | | | |  _| | (_| | |_| | | | (_| | || (_) | |
//  \____\___/|_| |_|_| |_|\__, |\__,_|_|  \__,_|\__\___/|_|
//                         |___/

fn generate_name(prefix: String) -> String
{
  use rand::Rng;
  const CHARSET: &[u8] = b"abcdefghijklmnopqrstuvwxyz0123456789";
  let mut rng = rand::thread_rng();
  let rand_string: String = (0..10)
    .map(|_| {
      let idx = rng.gen_range(0..CHARSET.len());
      CHARSET[idx] as char
    })
    .collect();
  format!("{}-{}", prefix, rand_string)
}

/// Configuration of the container
pub struct Configurator
{
  image: String,
  capture_stdio: bool,
  interactive: bool,
  tty: bool,
  daemon: bool,
  priviliged: bool,
  x11_forwarding: bool,
  network: Network,
  mounts: Vec<(String, String)>,
  env: Vec<(String, String)>,
  username: Option<String>,
  clean_up: bool,
  command: Vec<String>,
  name: String,
}

impl Configurator
{
  /// Create the container
  pub fn create(self) -> Container
  {
    Container {
      config: self,
      process: None,
    }
  }
  /// Set the command for the docker
  pub fn set_command(mut self, command: impl IntoIterator<Item = impl Into<String>>) -> Self
  {
    self.command = command
      .into_iter()
      .map(|x| -> String { x.into() })
      .collect();
    self
  }
  /// Capture stdin/err/out
  pub fn set_capture_stdio(mut self, capture: bool) -> Self
  {
    self.capture_stdio = capture;
    self
  }
  /// Set the docker in interactive mode.
  pub fn set_interactive(mut self, interactive: bool) -> Self
  {
    self.interactive = interactive;
    self
  }
  /// Set the docker in TTY mode.
  pub fn set_tty(mut self, tty: bool) -> Self
  {
    self.tty = tty;
    self
  }
  /// Enable/disable priviliged (see [docker documentation](https://docs.docker.com/engine/containers/run/#runtime-privilege-and-linux-capabilities)).
  pub fn set_privilged(mut self, priviliged: bool) -> Self
  {
    self.priviliged = priviliged;
    self
  }
  /// Forward X11 to the container
  pub fn set_x11_forwarding(mut self, x11_forwarding: bool) -> Self
  {
    self.x11_forwarding = x11_forwarding;
    self
  }
  /// Set the network mode
  pub fn set_network(mut self, network: Network) -> Self
  {
    self.network = network;
    self
  }
  /// Mount a drive from the host to the container
  pub fn mount(mut self, host_dir: impl Into<String>, container_dir: impl Into<String>) -> Self
  {
    self.mounts.push((host_dir.into(), container_dir.into()));
    self
  }
  /// Set the username in the container, this is required for X11 forwarding
  pub fn set_username(mut self, username: impl Into<String>) -> Self
  {
    self.username = Some(username.into());
    self
  }
  /// Set an environment variable
  pub fn set_env_variable(mut self, name: impl Into<String>, variable: impl Into<String>) -> Self
  {
    self.env.push((name.into(), variable.into()));
    self
  }
  /// Set if the container should be cleaned up after execution (set --rm)
  pub fn set_clean_up(mut self, clean_up: bool) -> Self
  {
    self.clean_up = clean_up;
    self
  }
  /// set the name of the container
  pub fn set_name(mut self, name: impl Into<String>) -> Self
  {
    self.name = name.into();
    self
  }
  /// generate a name with the given prefix
  pub fn generate_name(mut self, prefix: impl Into<String>) -> Self
  {
    self.name = generate_name(prefix.into());
    self
  }
}

//   ____            _        _
//  / ___|___  _ __ | |_ __ _(_)_ __   ___ _ __
// | |   / _ \| '_ \| __/ _` | | '_ \ / _ \ '__|
// | |__| (_) | | | | || (_| | | | | |  __/ |
//  \____\___/|_| |_|\__\__,_|_|_| |_|\___|_|

/// Handle to a container
pub struct Container
{
  config: Configurator,
  process: Option<SharedChild>,
}

macro_rules! config_to_arg {
  ($check:expr, $args:ident, $arg:expr) => {
    if $check
    {
      $args.push($arg.into());
    }
  };
}

impl Container
{
  /// Call this function to configure a new container
  pub fn configure(image: impl Into<String>) -> Configurator
  {
    Configurator {
      image: image.into(),
      capture_stdio: false,
      interactive: false,
      tty: false,
      daemon: false,
      priviliged: false,
      x11_forwarding: false,
      network: Network::Bridge,
      mounts: Default::default(),
      env: Default::default(),
      username: None,
      clean_up: true,
      name: generate_name("unknown".to_string()),
      command: Default::default(),
    }
  }

  /// Call this function to start the container.
  pub fn start(&mut self) -> Result<()>
  {
    if self.is_running()?
    {
      return Err(Error::ContainerRunning);
    }

    let mut args = Vec::<OsString>::new();
    args.push("run".into());
    args.push("--name".into());
    args.push(self.config.name.to_owned().into());

    config_to_arg!(self.config.daemon, args, "-d");
    config_to_arg!(self.config.tty, args, "-t");
    config_to_arg!(self.config.interactive, args, "-i");
    config_to_arg!(self.config.priviliged, args, "--privileged");
    config_to_arg!(self.config.clean_up, args, "--rm");

    args.push("--network".into());
    args.push(
      match self.config.network
      {
        Network::Bridge => "bridge",
        Network::Host => "host",
        Network::None => "none",
      }
      .into(),
    );

    if self.config.x11_forwarding
    {
      args.push("-e".into());
      args.push(format!("DISPLAY={}", std::env::var("DISPLAY")?).into());
      args.push("-v".into());
      args.push(
        format!(
          "/home/{}/.Xauthority:/home/{}/.Xauthority",
          whoami::username(),
          self
            .config
            .username
            .as_ref()
            .ok_or_else(|| Error::MissingUsername)?
        )
        .into(),
      );
      args.push("-v".into());
      args.push("/tmp/.X11-unix:/tmp/.X11-unix".into());
      args.push("-h".into());
      args.push(hostname::get()?.into())
    }

    for (host_dir, container_dir) in self.config.mounts.iter()
    {
      args.push("-v".into());
      args.push(format!("{}:{}", host_dir, container_dir).into());
    }

    // Add environment
    for (env_name, env_value) in self.config.env.iter()
    {
      args.push("-e".into());
      args.push(format!("{}={}", env_name, env_value).into());
    }

    // Set the image
    args.push(self.config.image.to_owned().into());

    // Set the command arguments
    let mut command = self
      .config
      .command
      .iter()
      .map(|x| -> OsString { x.into() })
      .collect();
    args.append(&mut command);

    // Prepare the process
    let mut cmd = Command::new("docker");
    cmd.args(args);

    if self.config.capture_stdio
    {
      cmd.stdin(Stdio::piped());
      cmd.stdout(Stdio::piped());
      cmd.stderr(Stdio::piped());
    }

    // Spawn
    self.process = Some(SharedChild::spawn(&mut cmd)?);
    Ok(())
  }
  /// Call this function to wait on the container
  pub fn wait(&self) -> Result<ExitStatus>
  {
    if let Some(process) = &self.process
    {
      Ok(process.wait()?)
    }
    else
    {
      Err(Error::ContainerNotRunning)
    }
  }
  /// Call this function to stop the container
  pub fn stop(&self) -> Result<()>
  {
    if let Some(_) = &self.process
    {
      Command::new("docker")
        .args(["kill", self.config.name.to_owned().as_str()])
        .output()?;
      Ok(())
    }
    else
    {
      Err(Error::ContainerNotRunning)
    }
  }
  /// Check if running
  pub fn is_running(&self) -> Result<bool>
  {
    if let Some(p) = &self.process
    {
      if p.try_wait()?.is_none()
      {
        return Ok(true);
      }
    }
    Ok(false)
  }
  /// Get stdin
  pub fn take_stdin(&self) -> Result<std::process::ChildStdin>
  {
    Ok(
      self
        .process
        .as_ref()
        .ok_or_else(|| Error::ContainerNotRunning)?
        .take_stdin()
        .ok_or_else(|| Error::StdIONotPiped)?,
    )
  }
  /// Get stdout
  pub fn take_stdout(&self) -> Result<std::process::ChildStdout>
  {
    Ok(
      self
        .process
        .as_ref()
        .ok_or_else(|| Error::ContainerNotRunning)?
        .take_stdout()
        .ok_or_else(|| Error::StdIONotPiped)?,
    )
  }
  /// Get stderr
  pub fn take_stderr(&self) -> Result<std::process::ChildStderr>
  {
    Ok(
      self
        .process
        .as_ref()
        .ok_or_else(|| Error::ContainerNotRunning)?
        .take_stderr()
        .ok_or_else(|| Error::StdIONotPiped)?,
    )
  }
}

impl Drop for Container
{
  fn drop(&mut self)
  {
    if self
      .is_running()
      .expect("to successfully query for running")
    {
      self.stop().expect("to successfully stop");
      self.wait().expect("to successfully wait");
    }
  }
}

#[cfg(test)]
mod tests
{
  use std::io::Read;

  use super::*;

  #[test]
  fn start_wait_stop_containers()
  {
    let mut container = Container::configure("alpine")
      .set_command(["sleep", "1"])
      .create();
    // Try to start and stop container
    container.start().expect("To start.");
    assert!(container.is_running().unwrap());
    container.stop().expect("To stop.");
    let wait_status = container.wait().unwrap();
    assert!(wait_status.success());

    assert!(!container.is_running().unwrap());

    // Run container and wait
    container.start().expect("To start.");
    assert!(container.is_running().unwrap());
    let wait_status = container.wait().unwrap();
    assert!(wait_status.success());
    assert_eq!(wait_status.code().unwrap(), 0);
    assert!(!container.is_running().unwrap());
  }
  #[test]
  fn interactive_containers()
  {
    use std::io::Write;

    let mut container = Container::configure("alpine")
      .set_interactive(true)
      .set_capture_stdio(true)
      .create();
    container.start().unwrap();
    let mut stdin = container.take_stdin().unwrap();
    stdin.write_all(b"echo Hello World!").unwrap();
    drop(stdin);
    let mut buf = vec![];
    container
      .take_stdout()
      .unwrap()
      .read_to_end(&mut buf)
      .unwrap();
    assert_eq!(buf, b"Hello World!\n");
  }
}