systemd-transient 0.1.3

use std::time::Duration;

use futures_util::StreamExt;
use tracing::warn;
use zbus::{
  Connection,
  zvariant::{ObjectPath, OwnedObjectPath, Value},
};

use crate::{
  ReusableTransientUnit, TransientUnitHandle, UnitExit,
  dbus::{
    manager::{JobRemovedStream, ManagerProxy},
    unit::UnitProxy,
  },
  error::SystemdTransientError,
  props::{UnitProperty, Usec, get_property},
  reusable::Opts,
};

const JOB_MODE_RPELACE: &str = "replace";
const ADDITIONAL_TIMEOUT_BUFFER: Duration = Duration::from_secs(5);

#[derive(Clone)]
pub struct SystemdTransientManager {
  conn: Connection,
}

impl SystemdTransientManager {
  pub async fn connect() -> Result<Self, SystemdTransientError> {
    let conn = Connection::session().await?;
    Ok(Self { conn })
  }

  pub fn new_reusable(&self, name: String, props: Vec<UnitProperty>) -> ReusableTransientUnit {
    ReusableTransientUnit {
      opts: Opts { name, props },
      unit_handle: None,
      systemd_manager: self.clone(),
    }
  }

  pub async fn start(
    &self,
    name: String,
    props: Vec<UnitProperty>,
  ) -> Result<TransientUnitHandle, SystemdTransientError> {
    let name = format!("{}.service", name);

    let proxy = ManagerProxy::new(&self.conn).await?;

    let start_timeout = if let Some(timeout) = get_property(&props, |prop| {
      if let UnitProperty::TimeoutStartUSec(usec) = prop {
        Some(Duration::from(*usec))
      } else {
        None
      }
    }) {
      timeout
    } else {
      Usec(proxy.default_timeout_start_usec().await?).into()
    } + ADDITIONAL_TIMEOUT_BUFFER;

    let stop_timeout = if let Some(timeout) = get_property(&props, |prop| {
      if let UnitProperty::TimeoutStopUSec(usec) = prop {
        Some(Duration::from(*usec))
      } else {
        None
      }
    }) {
      timeout
    } else {
      Usec(proxy.default_timeout_start_usec().await?).into()
    } + ADDITIONAL_TIMEOUT_BUFFER;

    let pairs: Vec<_> = props.into_iter().map(UnitProperty::into_pair).collect();
    let tuples: Vec<(&str, &Value<'_>)> = pairs.iter().map(|(k, v)| (k.as_ref(), v)).collect();
    let refs: Vec<&(&str, &Value<'_>)> = tuples.iter().collect();

    let action = async |proxy: &ManagerProxy<'_>| {
      proxy
        .start_transient_unit(&name, JOB_MODE_RPELACE, &refs, &[])
        .await
    };

    run_job(&proxy, action, start_timeout).await?;

    Ok(TransientUnitHandle {
      name,
      stop_timeout,
      manager: self.clone(),
    })
  }

  pub(crate) async fn stop(
    &self,
    name: &str,
    timeout: Duration,
  ) -> Result<(), SystemdTransientError> {
    let proxy = ManagerProxy::new(&self.conn).await?;
    let action = async |proxy: &ManagerProxy<'_>| proxy.stop_unit(name, JOB_MODE_RPELACE).await;
    run_job(&proxy, action, timeout).await
  }

  pub(crate) async fn wait_for_exit(&self, name: &str) -> Result<UnitExit, SystemdTransientError> {
    let manager_proxy = ManagerProxy::new(&self.conn).await?;

    let unit_path = manager_proxy.load_unit(name).await?;

    let unit_proxy = UnitProxy::builder(&self.conn)
      .path(unit_path)?
      .build()
      .await?;

    let mut changes = unit_proxy.receive_active_state_changed().await;

    // also check if the unit is exited to be begin with
    let mut state = unit_proxy.active_state().await?;

    loop {
      match state.as_str() {
        "inactive" => return Ok(UnitExit::Inactive),
        "failed" => return Ok(UnitExit::Failed),
        _ => {}
      }

      state = changes
        .next()
        .await
        .ok_or(SystemdTransientError::UnitStreamEnded)?
        .get()
        .await?;
    }
  }
}

async fn run_job<F>(
  proxy: &ManagerProxy<'_>,
  action: F,
  timeout: Duration,
) -> Result<(), SystemdTransientError>
where
  F: AsyncFnOnce(&ManagerProxy<'_>) -> zbus::Result<OwnedObjectPath>,
{
  let stream = proxy.receive_job_removed().await?;
  let job = action(proxy).await?;

  tokio::time::timeout(timeout, wait_for_job(stream, job.as_ref()))
    .await
    .map_err(|_| SystemdTransientError::JobTimedOut {
      job: job.clone(),
      timeout,
    })?
}

async fn wait_for_job(
  mut stream: JobRemovedStream,
  job_path: ObjectPath<'_>,
) -> Result<(), SystemdTransientError> {
  while let Some(signal) = stream.next().await {
    let args = match signal.args() {
      Ok(args) => args,
      Err(err) => {
        warn!("unparsable JobRemoved signal: {}", err);
        continue;
      }
    };

    if *args.job() != job_path {
      continue;
    }

    return match *args.result() {
      "done" => Ok(()),
      r => Err(SystemdTransientError::JobWaitFailed {
        unit: args.unit().to_string(),
        reason: r.to_owned(),
      }),
    };
  }

  Err(SystemdTransientError::JobStreamEnded)
}