use anyhow::Result;
use std::process::Child;
use std::time::{Duration, Instant};
use std::thread;

pub trait Workload: Send + Sync {
    fn start(&mut self, threads: usize, load_percent: usize) -> Result<()>;
    fn stop(&mut self) -> Result<()>;
    fn get_throughput(&self) -> Result<f64>;
}

pub struct StressNg {
    child: Option<Child>,
}

impl StressNg {
    pub fn new() -> Self {
        Self { child: None }
    }
}

impl Workload for StressNg {
    fn start(&mut self, threads: usize, load_percent: usize) -> Result<()> {
        self.stop()?;
        
        let child = std::process::Command::new("stress-ng")
            .args([
                "--cpu", &threads.to_string(),
                "--cpu-load", &load_percent.to_string(),
                "--quiet"
            ])
            .spawn()?;
        
        self.child = Some(child);
        Ok(())
    }

    fn stop(&mut self) -> Result<()> {
        if let Some(mut child) = self.child.take() {
            // Try SIGTERM first
            #[cfg(unix)]
            {
                use libc::{kill, SIGTERM};
                unsafe { kill(child.id() as i32, SIGTERM); }
            }

            let start = Instant::now();
            let mut exited = false;
            while start.elapsed() < Duration::from_secs(2) {
                if let Ok(Some(_)) = child.try_wait() {
                    exited = true;
                    break;
                }
                thread::sleep(Duration::from_millis(100));
            }

            if !exited {
                let _ = child.kill();
                let _ = child.wait();
            }
        }
        Ok(())
    }

    fn get_throughput(&self) -> Result<f64> {
        Ok(0.0)
    }
}

impl Drop for StressNg {
    fn drop(&mut self) {
        let _ = self.stop();
    }
}