pandora_lib_promethion/src/runners.rs

use std::{
    fs::File,
    io::{BufRead, BufReader, Write},
    process::{Child, Command, Stdio},
    sync::{Arc, Mutex},
    thread,
};

use chrono::{DateTime, Utc};
use log::{info, warn};
use serde::{Deserialize, Serialize};
use uuid::Uuid;

use crate::DOCKER_ID;


pub trait Run {
    fn run(&mut self) -> anyhow::Result<()>;
}

pub trait Wait {
    fn wait(&mut self) -> anyhow::Result<()>;
}

pub trait RunWait: Run + Wait + Log {}
impl<T: Run + Wait + Log> RunWait for T {}

pub fn run_wait<T: RunWait>(item: &mut T) -> anyhow::Result<RunReport> {
    let mut run_report = RunReport {
        start: Utc::now(),
        ..Default::default()
    };
    item.run()?;
    item.wait()?;
    run_report.log = item.log();
    run_report.end = Utc::now();
    Ok(run_report)
}

#[derive(Debug, Default, Serialize, Deserialize)]
pub struct RunReport {
    pub start: DateTime<Utc>,
    pub end: DateTime<Utc>,
    pub log: String,
}

impl RunReport {
    /// Serialize the RunReport to a JSON string
    pub fn save_to_file(&self, file_prefix: &str) -> std::io::Result<()> {
        let json_data = serde_json::to_string_pretty(self).expect("Failed to serialize RunReport");
        let uuid = Uuid::new_v4().to_string()[..5].to_string();
        let file_path = format!("{}{}.log", file_prefix, uuid);
        let mut file = File::create(&file_path)?;
        file.write_all(json_data.as_bytes())?;
        Ok(())
    }
}

pub trait Log {
    fn log(&self) -> String;
}

#[derive(Debug, Default)]
pub struct DockerRun {
    pub args: Vec<String>,
    pub container_id: Arc<Mutex<Option<String>>>,
    pub start: DateTime<Utc>,
    pub logs: Arc<Mutex<String>>,
}

impl DockerRun {
    pub fn new(args: &[&str]) -> Self {
        DockerRun {
            args: args.iter().map(|e| e.to_string()).collect(),
            start: Utc::now(),
            ..Default::default()
        }
    }
}

impl Run for DockerRun {
    fn run(&mut self) -> anyhow::Result<()> {
        // Setup Ctrl-C handler cant be defined two times
        let _ = ctrlc::try_set_handler(move || {
            if let Ok(container_id) = DOCKER_ID.lock() {
                if let Some(ref id) = *container_id {
                    warn!("Stopping Docker container...");
                    let _ = Command::new("docker").args(["stop", id]).status();
                }
            }
            std::process::exit(1);
        });

        // Spawn the main command
        let output = Command::new("docker")
            .args(&self.args)
            .stdout(Stdio::piped())
            .stderr(Stdio::inherit())
            .output()
            .expect("Failed to run Docker command");

        // add id to Arc
        let id = String::from_utf8_lossy(&output.stdout).trim().to_string();
        {
            let mut container_id_lock = DOCKER_ID.lock().unwrap();
            *container_id_lock = Some(id.clone());
        }

        // Spawn a thread to follow the logs
        let log_id = id.clone();
        let logs_clone = Arc::clone(&self.logs);
        let _loger_thread = thread::spawn(move || {
            let output = Command::new("docker")
                .args(["inspect", "--format='{{.Config.Image}}'", &log_id])
                .output()
                .expect("Failed to execute Docker inspect");
            let container_name = if output.status.success() {
                String::from_utf8(output.stdout)
                    .expect("Failed to convert output to string")
                    .trim() // Trim to remove any trailing newlines
                    .to_string()
            } else {
                "?".to_string()
            };
            let mut child = Command::new("docker")
                .args(["logs", "--follow", &log_id])
                .stdout(Stdio::piped())
                .stderr(Stdio::inherit())
                .spawn()
                .expect("Failed to follow Docker logs");

            if let Some(stdout) = child.stdout.take() {
                let reader = BufReader::new(stdout);
                for line in reader.lines().map_while(Result::ok) {
                    info!("[{container_name}] {line}");
                    let mut logs_lock = logs_clone.lock().unwrap();
                    logs_lock.push_str(&line);
                    logs_lock.push('\n');
                }
            }
        });
        Ok(())
    }
}

impl Wait for DockerRun {
    fn wait(&mut self) -> anyhow::Result<()> {
        if let Ok(container_id) = DOCKER_ID.lock() {
            if let Some(ref id) = *container_id {
                let status = Command::new("docker")
                    .args(["wait", id])
                    .status()
                    .expect("Failed to wait on Docker container");
                if !status.success() {
                    warn!("Docker command failed with status: {}", status);
                }
            }
        }
        let mut container_id_lock = DOCKER_ID.lock().unwrap();
        *container_id_lock = None;

        Ok(())
    }
}

impl Log for DockerRun {
    fn log(&self) -> String {
        let logs_lock = self.logs.lock().unwrap();
        logs_lock.to_string()
    }
}

pub struct CommandRun {
    pub bin: String,
    pub args: Vec<String>,
    pub child: Option<Child>,
    pub log: String,
}

impl CommandRun {
    pub fn new(bin: &str, args: &[&str]) -> Self {
        CommandRun {
            bin: bin.to_string(),
            args: args.iter().map(|e| e.to_string()).collect(),
            child: None,
            log: String::default(),
        }
    }
}

impl Run for CommandRun {
    fn run(&mut self) -> anyhow::Result<()> {
        info!("Running command: {} {}", &self.bin, &self.args.join(" "));
        let mut child = Command::new(&self.bin)
            .args(&self.args)
            .stdout(Stdio::inherit())
            .stderr(Stdio::piped())
            .spawn()
            .expect("Failed to spawn");

        if let Some(stderr) = child.stderr.take() {
            let reader = BufReader::new(stderr);
            for line in reader.lines().map_while(Result::ok) {
                info!("[{}] {line}", self.bin);
                self.log.push_str(&line);
                self.log.push('\n');
            }
        }
        self.child = Some(child);
        Ok(())
    }
}

impl Wait for CommandRun {
    fn wait(&mut self) -> anyhow::Result<()> {
        if let Some(child) = &mut self.child {
            child.wait().unwrap();
        }
        Ok(())
    }
}

impl Log for CommandRun {
    fn log(&self) -> String {
        self.log.clone()
    }
}