// SPDX-FileCopyrightText: Peter Pentchev <roam@ringlet.net>
// SPDX-License-Identifier: BSD-2-Clause
//! Execute a program using the `smol` crate, send signals to it.

use core::time::Duration;

use std::os::unix::process::ExitStatusExt as _;

use anyhow::{Context as _, Result, anyhow};
use log::{debug, info, warn};
use nix::sys::signal;
use nix::unistd::Pid;
use oneshot::Sender as OneShotSender;
use smol::Timer;
use smol::channel::{self, Sender};
use smol::process::{Child, ChildStderr, ChildStdin, ChildStdout, Command};
use sysexits::ExitCode;

use crate::defs::Config;
use crate::exitcode::ExitCodeReport;

/// A message sent from a task to the main limit loop.
#[derive(Debug)]
enum LimitMsg {
    /// The child process completed, or we could not get its status.
    WaitResult(ExitCodeReport),

    /// The warning timer expired.
    TimerWarn,

    /// The kill timer expired too.
    TimerKill,
}

/// The I/O streams of the child process, if requested by the caller.
#[derive(Debug)]
pub struct ChildIo {
    /// The standard input stream.
    stdin: Option<ChildStdin>,

    /// The standard output stream.
    stdout: Option<ChildStdout>,

    /// The standard error stream.
    stderr: Option<ChildStderr>,
}

impl ChildIo {
    /// Take the standard input stream if (still) available.
    #[inline]
    #[must_use]
    pub const fn take_stdin(&mut self) -> Option<ChildStdin> {
        self.stdin.take()
    }

    /// Take the standard output stream if (still) available.
    #[inline]
    #[must_use]
    pub const fn take_stdout(&mut self) -> Option<ChildStdout> {
        self.stdout.take()
    }

    /// Take the standard error stream if (still) available.
    #[inline]
    #[must_use]
    pub const fn take_stderr(&mut self) -> Option<ChildStderr> {
        self.stderr.take()
    }
}

/// Use the `smol` crate's exports to run a program limiting its execution time.
#[derive(Debug)]
pub struct Limit {
    /// The runtime configuration - limits and signals.
    cfg: Config,
}

impl Limit {
    /// Prepare a [`Limit`] object for running programs.
    #[inline]
    #[must_use]
    pub const fn new(cfg: Config) -> Self {
        Self { cfg }
    }

    /// Execute a program and limit things.
    #[inline]
    #[must_use]
    pub async fn run(&self, cmd: Command) -> ExitCodeReport {
        self.run_with_io(cmd, None).await
    }

    /// Spawn the child process.
    fn spawn_child(
        mut cmd: Command,
        tx_opt: Option<OneShotSender<ChildIo>>,
    ) -> Result<(Child, u32, i32), ExitCodeReport> {
        let mut child = match cmd.spawn().context("Could not run the specified command") {
            Ok(child) => child,
            Err(err) => return Err((ExitCode::OsErr, err).into()),
        };
        let pid = child.id();
        let pid_i32 = match i32::try_from(pid)
            .with_context(|| format!("Could not represent pid {pid} as i32"))
        {
            Ok(pid_i32) => pid_i32,
            Err(err) => return Err((ExitCode::OsErr, err).into()),
        };

        if let Some(tx) = tx_opt {
            let child_io = ChildIo {
                stdin: child.stdin.take(),
                stdout: child.stdout.take(),
                stderr: child.stderr.take(),
            };
            if tx.send(child_io).is_err() {
                warn!("Could not tell the caller about the child I/O streams");
            }
        }

        Ok((child, pid, pid_i32))
    }

    /// Wait for the child to end, tell the main thread how it ended.
    async fn wait_for_child(mut child: Child, tx: Sender<LimitMsg>) {
        let res = match child
            .status()
            .await
            .context("Could not examine the exit status of the command")
        {
            Ok(stat) => stat.code().map_or_else(
                || {
                    stat.signal().map_or_else(
                        || {
                            (
                                ExitCode::OsErr,
                                anyhow!("Unexpected command exit status {stat:?}"),
                            )
                                .into()
                        },
                        ExitCodeReport::from_raw_signal,
                    )
                },
                Into::into,
            ),
            Err(err) => (ExitCode::OsErr, err).into(),
        };
        if tx.send(LimitMsg::WaitResult(res)).await.is_err() {
            warn!("Could not send the wait status to the main thread");
        }
    }

    /// Send the signals to the child process at the specified intervals.
    async fn send_signals(cfg: Config, tx: Sender<LimitMsg>) {
        Timer::after(Duration::from_secs_f32(cfg.warntime())).await;
        debug!("Signalling the main thread that the warning timer expired");
        if tx.send(LimitMsg::TimerWarn).await.is_err() {
            warn!("Could not tell the main thread about the warning timer");
            return;
        }

        Timer::after(Duration::from_secs_f32(cfg.killtime())).await;
        debug!("Signalling the main thread that the kill timer expired");
        if tx.send(LimitMsg::TimerKill).await.is_err() {
            warn!("Could not tell the main thread about the kill timer");
        }
    }

    /// Execute a program and limit things, sending stdin/stdout/stderr back.
    #[inline]
    #[must_use]
    pub async fn run_with_io(
        &self,
        cmd: Command,
        tx_opt: Option<OneShotSender<ChildIo>>,
    ) -> ExitCodeReport {
        let (child, pid, pid_i32) = match Self::spawn_child(cmd, tx_opt) {
            Ok(values) => values,
            Err(report) => return report,
        };
        let (tx, rx) = channel::bounded(16_usize);
        debug!("Spawned the command as {pid}");
        let timer_task = {
            let tx_timer = tx.clone();
            let cfg = self.cfg;
            smol::spawn(async move { Self::send_signals(cfg, tx_timer).await })
        };
        let wait_task = {
            let tx_wait = tx.clone();
            smol::spawn(async move { Self::wait_for_child(child, tx_wait).await })
        };
        let loop_res = loop {
            if let Ok(msg) = rx.recv().await {
                match msg {
                    LimitMsg::WaitResult(res) => break res,
                    LimitMsg::TimerWarn => {
                        let sig = self.cfg.warnsig();
                        info!("Sending {sig} to process {pid}");
                        if let Err(err) = signal::kill(Pid::from_raw(pid_i32), sig.sig())
                            .with_context(|| format!("Could not send {sig} to process {pid}"))
                        {
                            break (ExitCode::OsErr, err).into();
                        }
                    }
                    LimitMsg::TimerKill => {
                        let sig = self.cfg.killsig();
                        info!("Sending {sig} to process {pid}");
                        if let Err(err) = signal::kill(Pid::from_raw(pid_i32), sig.sig())
                            .with_context(|| format!("Could not send {sig} to process {pid}"))
                        {
                            break (ExitCode::OsErr, err).into();
                        }
                    }
                }
            } else {
                break (
                    ExitCode::OsErr,
                    anyhow!("timelimit internal error: what happened to the wait task?"),
                )
                    .into();
            }
        };
        wait_task.cancel().await;
        timer_task.cancel().await;
        loop_res
    }
}
