// Copyright Kani Contributors
// SPDX-License-Identifier: Apache-2.0 OR MIT
//! Utilities to interact with the `Litani` build accumulator.

use pulldown_cmark::escape::StrWrite;
use serde::Deserialize;
use std::collections::HashMap;
use std::path::Path;
use std::process::{Child, Command};

/// Data structure representing a full `litani` run.
/// The same representation is used to represent a run
/// in the `run.json` (cache) file generated by `litani`
///
/// Deserialization is performed automatically for most
/// attributes in such files, but it may require you to
/// extend it if advanced features are used (e.g., pools)
#[derive(Debug, Deserialize)]
pub struct LitaniRun {
    pub aux: Option<HashMap<String, String>>,
    pub project: String,
    pub version: String,
    pub version_major: u32,
    pub version_minor: u32,
    pub version_patch: u32,
    pub release_candidate: bool,
    pub run_id: String,
    pub start_time: String,
    pub parallelism: LitaniParalellism,
    pub latest_symlink: Option<String>,
    pub end_time: String,
    pub pipelines: Vec<LitaniPipeline>,
}

impl LitaniRun {
    pub fn get_pipelines(self) -> Vec<LitaniPipeline> {
        self.pipelines
    }
}

#[derive(Debug, Deserialize)]
pub struct LitaniParalellism {
    pub trace: Vec<LitaniTrace>,
    pub max_paralellism: Option<u32>,
    pub n_proc: u32,
}

#[derive(Debug, Deserialize)]
pub struct LitaniTrace {
    pub running: u32,
    pub finished: u32,
    pub total: u32,
    pub time: String,
}

#[derive(Debug, Deserialize)]
pub struct LitaniPipeline {
    pub name: String,
    pub ci_stages: Vec<LitaniStage>,
    pub url: String,
    pub status: String,
}

impl LitaniPipeline {
    pub fn get_name(&self) -> &String {
        &self.name
    }

    pub fn get_status(&self) -> bool {
        match self.status.as_str() {
            "fail" => false,
            "success" => true,
            _ => panic!("pipeline status is not \"fail\" nor \"success\""),
        }
    }
}

#[derive(Debug, Deserialize)]
pub struct LitaniStage {
    pub jobs: Vec<LitaniJob>,
    pub progress: u32,
    pub complete: bool,
    pub status: String,
    pub url: String,
    pub name: String,
}

// Some attributes in litani's `jobs` are not always included
// or they are null, so we use `Option<...>` to deserialize them
#[derive(Debug, Deserialize)]
pub struct LitaniJob {
    pub wrapper_arguments: LitaniWrapperArguments,
    pub complete: bool,
    pub start_time: Option<String>,
    pub timeout_reached: Option<bool>,
    pub command_return_code: Option<i32>,
    pub memory_trace: Option<HashMap<String, String>>,
    pub loaded_outcome_dict: Option<HashMap<String, String>>,
    pub outcome: Option<String>,
    pub wrapper_return_code: Option<i32>,
    pub stdout: Option<Vec<String>>,
    pub stderr: Option<Vec<String>>,
    pub end_time: Option<String>,
    pub duration_str: Option<String>,
    pub duration: Option<u32>,
}

// Some attributes in litani's `wrapper_arguments` are not always included
// or they are null, so we use `Option<...>` to deserialize them
#[derive(Debug, Deserialize)]
pub struct LitaniWrapperArguments {
    pub subcommand: String,
    pub verbose: bool,
    pub very_verbose: bool,
    pub inputs: Vec<String>,
    pub command: String,
    pub outputs: Option<Vec<String>>,
    pub pipeline_name: String,
    pub ci_stage: String,
    pub cwd: Option<String>,
    pub timeout: Option<u32>,
    pub timeout_ok: Option<bool>,
    pub timeout_ignore: Option<bool>,
    pub ignore_returns: Option<bool>,
    pub ok_returns: Vec<String>,
    pub outcome_table: Option<HashMap<String, String>>,
    pub interleave_stdout_stderr: bool,
    pub stdout_file: Option<String>,
    pub stderr_file: Option<String>,
    pub pool: Option<u32>,
    pub description: String,
    pub profile_memory: bool,
    pub profile_memory_interval: u32,
    pub phony_outputs: Option<Vec<String>>,
    pub tags: Option<String>,
    pub status_file: String,
    pub job_id: String,
}

/// Data structure representing a `Litani` build.
pub struct Litani {
    /// A buffer of the `spawn`ed Litani jobs so far. `Litani` takes some time
    /// to execute each `add-job` command and executing thousands of them
    /// sequentially takes a considerable amount of time. To speed up the
    /// execution of those commands, we spawn those commands sequentially (as
    /// normal). However, instead of `wait`ing for each process to terminate,
    /// we add its handle to a buffer of the `spawn`ed processes and continue
    /// with our program. Once we are done adding jobs, we wait for all of them
    /// to terminate before we run the `run-build` command.
    spawned_commands: Vec<Child>,
}

impl Litani {
    /// Sets up a new [`Litani`] run.
    pub fn init(
        project_name: &str,
        stage_names: &[&str],
        output_prefix: &Path,
        output_symlink: &Path,
    ) -> Self {
        Command::new("litani")
            .args([
                "init",
                "--project-name",
                project_name,
                "--output-prefix",
                output_prefix.to_str().unwrap(),
                "--output-symlink",
                output_symlink.to_str().unwrap(),
                "--stages",
            ])
            .args(stage_names)
            .spawn()
            .unwrap()
            .wait()
            .unwrap();
        Self { spawned_commands: Vec::new() }
    }

    /// Adds a single command with its dependencies.
    #[allow(clippy::too_many_arguments)]
    pub fn add_job(
        &mut self,
        command: &Command,
        inputs: &[&Path],
        outputs: &[&Path],
        description: &str,
        pipeline: &str,
        stage: &str,
        exit_status: i32,
        timeout: u32,
    ) {
        let mut job = Command::new("litani");
        // The given command may contain additional env vars. Prepend those vars
        // to the command before passing it to Litani.
        let job_envs: HashMap<_, _> = job.get_envs().collect();
        let mut new_envs = String::new();
        command.get_envs().fold(&mut new_envs, |fmt, (k, v)| {
            if !job_envs.contains_key(k) {
                fmt.write_fmt(format_args!(
                    "{}=\"{}\" ",
                    k.to_str().unwrap(),
                    v.unwrap().to_str().unwrap()
                ))
                .unwrap();
            }
            fmt
        });
        job.args([
            "add-job",
            "--command",
            &format!("{new_envs}{command:?}"),
            "--description",
            description,
            "--pipeline-name",
            pipeline,
            "--ci-stage",
            stage,
            "--ok-returns",
            &exit_status.to_string(),
            "--timeout",
            &timeout.to_string(),
        ]);
        if !inputs.is_empty() {
            job.arg("--inputs").args(inputs);
        }
        if !outputs.is_empty() {
            job.arg("--outputs").args(outputs).arg("--phony-outputs").args(outputs);
        }
        // Start executing the command, but do not wait for it to terminate.
        self.spawned_commands.push(job.spawn().unwrap());
    }

    /// Starts a [`Litani`] run.
    pub fn run_build(&mut self) {
        // Wait for all spawned processes to terminate.
        for command in self.spawned_commands.iter_mut() {
            command.wait().unwrap();
        }
        self.spawned_commands.clear();
        // Run `run-build` command and wait for it to finish.
        Command::new("litani")
            .args(["run-build", "--no-pipeline-dep-graph"])
            .spawn()
            .unwrap()
            .wait()
            .unwrap();
    }
}