//! Watch processes through it's pid
use anyhow::Context;
use std::{fs, io, path::PathBuf, ptr, time::Duration};
use thiserror::Error;

use crate::{
    pipeline::{MeasurementPipeline, control::request, matching::SourceNamePattern},
    plugin::event::{EndConsumerMeasurement, StartConsumerMeasurement},
    resources::ResourceConsumer,
};

use super::{RunningAgent, builder::ShutdownError};

/// Error that can occur in [`watch_process`].
#[derive(Error, Debug)]
pub enum WatchError {
    /// The process could not be spawned.
    #[error("failed to use process with pid {0}")]
    PidValue(String),
    /// The process could not be spawned.
    #[error("failed to check process with pid {0}")]
    PidCheck(String, #[source] std::num::ParseIntError),
    /// The process has spawned but waiting for it has failed.
    #[error("failed to wait for pid {0}")]
    ProcessWait(i32, #[source] std::io::Error),
    /// An error occurred while waiting for the agent to shut down.
    #[error("error in shutdown")]
    Shutdown(#[source] ShutdownError),
}

const TRIGGER_TIMEOUT: Duration = Duration::from_secs(1);

/// Watch process that runs identified by it's pid until it's end.
///
/// The measurement sources are triggered before the process spawns and after it exits.
///
/// After the process exits, the pipeline must stop within `shutdown_timeout`, or an error is returned.
pub fn watch_process(agent: RunningAgent, pid: u32, shutdown_timeout: Duration) -> Result<(), WatchError> {
    // Check if we can convert the pid to i32, from u32 because we don't want a negative value
    let pid_i32 = if pid <= i32::MAX as u32 {
        pid as i32
    } else {
        return Err(WatchError::PidValue(String::from("Value exceeds i32::MAX")));
    };

    if let Err(e) = trigger_poll_now(&agent.pipeline) {
        log::error!("Could not trigger a first time poll before the child spawn: {e}");
    }
    log::info!("Watch process spawned with pid {pid_i32}.");
    crate::plugin::event::start_consumer_measurement()
        .publish(StartConsumerMeasurement(vec![ResourceConsumer::Process { pid }]));

    // Spawn the process and wait for it to exit.
    wait_child(pid_i32)?;
    log::info!("Watched process exited, Alumet will now stop.");

    // One last measurement.
    if let Err(e) = trigger_poll_now(&agent.pipeline) {
        log::error!("Could not trigger one last time poll after the child exit: {e}");
    }

    // Publish an event to perform a measurement at the end of the experiment
    log::info!("Publishing EndConsumerMeasurement event");
    crate::plugin::event::end_consumer_measurement().publish(EndConsumerMeasurement);

    // Stop the pipeline
    agent.pipeline.control_handle().shutdown();
    agent.wait_for_shutdown(shutdown_timeout).map_err(WatchError::Shutdown)
}

/// Spawns a child process and waits for it to exit.
fn wait_child(pid: i32) -> Result<(), WatchError> {
    match unsafe { nc::pidfd_open(pid, 0) } {
        Ok(pidfd) => wait_child_syscall(pid, pidfd),
        Err(e) => {
            log::warn!(
                "PIDFD_OPEN returned an error: {}, syscall may not be supported on this system: - use of another way",
                e
            );
            wait_child_sysfs(pid)
        }
    }
}

// Wait for process exit using sysfs file
fn wait_child_sysfs(pid: i32) -> Result<(), WatchError> {
    let path = PathBuf::from(format!("/proc/{pid}/"));
    loop {
        if !fs::metadata(&path).is_ok() {
            break;
        }
        // Wait 5s to avoid too much load
        std::thread::sleep(std::time::Duration::from_secs(5));
    }
    Ok(())
}

// Wait for process exit using a syscall
fn wait_child_syscall(pid: i32, pidfd: i32) -> Result<(), WatchError> {
    loop {
        // Attempt to read from the file descriptor
        let mut timeout = libc::timeval {
            tv_sec: Duration::from_secs(5).as_secs() as i64,
            tv_usec: 0,
        };
        let mut read_fds = unsafe { std::mem::zeroed::<libc::fd_set>() }; // Initialize fd_set
        unsafe { libc::FD_SET(pidfd, &mut read_fds) }; // Add the pidfd to the read set
        let result = unsafe { libc::select(pidfd + 1, &mut read_fds, ptr::null_mut(), ptr::null_mut(), &mut timeout) };

        if result < 0 {
            // An error occurred
            log::error!("Error occurred in select: {}", io::Error::last_os_error());
            break;
        } else if result == 0 {
            // Timeout occurred - It means the process is still running
            continue;
        } else {
            // At least one file descriptor is ready
            if unsafe { libc::FD_ISSET(pidfd, &read_fds) } {
                // The process has terminated
                log::info!("Process {} has terminated", pid);
                break;
            }
        }
        // Reset the read file descriptor set for the next select call
        unsafe { libc::FD_SET(pidfd, &mut read_fds) };
    }
    unsafe { libc::close(pidfd) };
    Ok(())
}

// Triggers one poll (on all sources that support manual trigger).
fn trigger_poll_now(pipeline: &MeasurementPipeline) -> anyhow::Result<()> {
    let control_handle = pipeline.control_handle();
    let send_task = control_handle.send_wait(
        request::source(SourceNamePattern::wildcard()).trigger_now(),
        TRIGGER_TIMEOUT,
    );
    pipeline
        .async_runtime()
        .block_on(send_task)
        .context("failed to send TriggerMessage")
}
