// Copyright 2021 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

//! Clidoc generates documentation for host tool commands consisting of their --help output.
use anyhow::{bail, Context, Result};
use argh::FromArgs;
use flate2::write::GzEncoder;
use flate2::Compression;
use lazy_static::lazy_static;
use std::collections::{HashMap, HashSet};
use std::env;
use std::ffi::{OsStr, OsString};
use std::fs::{self, File};
use std::io::{BufWriter, Write};
use std::path::{Path, PathBuf};
use std::process::Command;
use std::sync::Once;
use tar::Builder;
use tracing::{debug, info};

mod ffx_doc;

enum HelpError {
    Ignore,
    Fail(String),
}

/// CliDoc generates documentation for core Fuchsia developer tools.
#[derive(Debug, FromArgs)]
struct Opt {
    // Default input dir is parent dir of this tool, containing host tools exes
    // $FUCHSIA_DIR/out/default/host_x64 or $FUCHSIA_DIR/out/default/host-tools
    /// set the input folder
    #[argh(
        option,
        short = 'i',
        default = "env::current_exe().unwrap().parent().unwrap().to_path_buf()"
    )]
    in_dir: PathBuf,

    /// set the output directory
    #[argh(option, short = 'o', default = "PathBuf::from(\".\".to_string())")]
    out_dir: PathBuf,

    /// reduce text output
    #[argh(switch)]
    quiet: bool,

    /// increase text output
    #[argh(switch, short = 'v')]
    verbose: bool,

    /// path for tarball archive of the output will be compressed as a tarball
    /// and intermediate files will be cleaned up
    /// For example: "clidoc_out.tar.gz". Note that .tar.gz is not automatically
    /// added as a file extension.
    #[argh(option)]
    archive_path: Option<PathBuf>,

    /// path for depfile if clidoc is invoked as a BUILD action
    /// if set, will output a depfile based on inputs at specified location
    /// depfile is only supported with archive_path flag.
    /// depfile is a Ninja term and does not need to be split into 2 words.
    #[argh(option)]
    depfile: Option<PathBuf>,

    /// path to sdk manifest used to add files read by ffx to the depfile.
    #[argh(option)]
    sdk_manifest: Option<PathBuf>,

    /// root of the sdk for ffx
    #[argh(option)]
    sdk_root: Option<PathBuf>,

    /// isolate-dir to use with ffx
    #[argh(option)]
    isolate_dir: Option<PathBuf>,

    /// commands to run, otherwise defaults to internal list of commands.
    /// relative paths are on the input_path. Absolute paths are used as-is.
    #[argh(positional)]
    cmd_list: Vec<PathBuf>,
}

// Formatting styles for codeblocks.
const CODEBLOCK_START: &str = "```none {: style=\"white-space: break-spaces;\" \
        .devsite-disable-click-to-copy}\n";
const CODEBLOCK_END: &str = "```\n";

const HEADER: &str = r#"<!--  DO NOT EDIT THIS FILE DIRECTLY

 This file is generated using clidoc by parsing the help output of this tool.
 Please edit the help output or clidoc's processing of that output to make changes
 to this file.

 -->
"#;

// TODO(https://fxbug.dev/42148459): Move allow list to its own separate config file.
const ALLOW_LIST: &'static [&'static str] = &[
    "blobfs-compression",
    "bootserver",
    "cmc",
    "configc",
    "ffx",
    "fidl-format",
    "fidlc",
    "fidlcat",
    "fidlgen_cpp",
    "fidlgen_hlcpp",
    "fidlgen_rust",
    "fpublish",
    "fremote",
    "fserve",
    "fssh",
    "funnel",
    "minfs",
    "pm",
    "symbolizer",
    "triage",
    "zbi",
    "zxdb",
];

// This Hashset includes all sdk tools which return stderr and non-zero error codes
// when invoking --help.
lazy_static! {
    static ref IGNORE_ERR_CODE: HashSet<&'static str> = {
        let h = HashSet::from(["bootserver", "fssh", "fremote", "minfs", "symbolizer", "zxdb"]);
        h
    };
}

fn main() -> Result<()> {
    let opt: Opt = argh::from_env();
    run(opt)
}

static INIT_LOGGER: Once = Once::new();

fn set_up_logger(opt: &Opt) {
    INIT_LOGGER.call_once(|| {
        if opt.verbose {
            tracing_subscriber::fmt().compact().with_max_level(tracing::Level::DEBUG).init();
            debug!("Debug logging enabled.");
        } else if opt.quiet {
            tracing_subscriber::fmt().compact().with_max_level(tracing::Level::WARN).init();
        } else {
            tracing_subscriber::fmt().compact().with_max_level(tracing::Level::INFO).init();
        }
    });
}

fn run(opt: Opt) -> Result<()> {
    if opt.quiet && opt.verbose {
        bail!("cannot use --quiet and --verbose together");
    }

    set_up_logger(&opt);

    // Set the directory for the command executables.
    let input_path = &opt.in_dir;
    info!("Input dir: {}", input_path.display());

    // Set the directory to output documentation to.
    let output_path = &opt.out_dir;
    info!("Output dir: {}", output_path.display());

    let mut cmd_paths: Vec<PathBuf>;

    if opt.cmd_list.is_empty() {
        debug!("Building cmd list from defaults");
        // Create a set of SDK tools to generate documentation for.
        let allow_list: HashSet<OsString> =
            ALLOW_LIST.iter().cloned().map(OsString::from).collect();
        // Create a vector of full paths to each command in the allow_list.
        cmd_paths = get_command_paths(&input_path, &allow_list)?;
    } else {
        // Use the commands passed on the command line. If they are relative paths,
        // make them absolute based on the input_path.
        cmd_paths = Vec::new();
        for p in opt.cmd_list {
            if p.is_absolute() {
                cmd_paths.push(p);
            } else {
                cmd_paths.push(input_path.join(p));
            }
        }
        debug!("Using cmds from opt.cmd_list: {:?}", cmd_paths);
    }

    // Create the directory for doc files if it doesn't exist.
    create_output_dir(&output_path)
        .context(format!("Unable to create output directory {:?}", output_path))?;

    // Write documentation output for each command.
    for cmd_path in cmd_paths.iter() {
        // ffx can export the help info in JSON format.
        if cmd_path.ends_with("ffx") {
            ffx_doc::write_formatted_output_for_ffx(
                &cmd_path,
                output_path,
                &opt.sdk_root,
                &opt.sdk_manifest,
                &opt.isolate_dir,
            )
            .context(format!(
                "Unable to write generate doc for {:?} to {:?}",
                cmd_path, output_path
            ))?;
        } else {
            write_formatted_output(&cmd_path, output_path).context(format!(
                "Unable to write generate doc for {:?} to {:?}",
                cmd_path, output_path
            ))?;
        }
    }

    info!("Generated documentation at dir: {}", &output_path.display());

    if let Some(tardir) = opt.archive_path {
        // First check if depfile is needed as well since this will probably be invoked
        // as part of a BUILD action.
        if let Some(depfile_path) = opt.depfile {
            let mut f = File::create(depfile_path).expect("Unable to create file");
            let output_filename = tardir.display();
            // Documented tools live in host_$ARCH path of build directory
            let input_path_last = input_path.file_name().expect("input path trailing element");
            let rel_path: PathBuf = PathBuf::from(input_path_last);

            for cmd_path in cmd_paths {
                let tool = cmd_path.file_name();
                let filename = rel_path.join(tool.expect("get toolname"));
                info!("depfile: {output_filename}: {filename:?}");
                write!(f, "{output_filename}: {}\n", filename.display())?;
            }
        }

        info!("Tarballing output at {:?}", tardir);
        let tar_gz = File::create(tardir)?;
        let enc = GzEncoder::new(tar_gz, Compression::default());
        let mut tar = Builder::new(enc);
        tar.append_dir_all("clidoc/", output_path.to_str().expect("Get file name of outdir"))?;

        info!("Cleaning up {:?}", output_path);
        fs::remove_dir_all(output_path)?
    }
    Ok(())
}

#[derive(Eq, Copy, Clone, Debug, Hash, PartialEq)]
enum Section {
    Top,
    Options,
    Commands,
}

/// Helper function for write_formatted_output.
///
/// Recursively calls `cmd_name`'s subcommands and writes to `output_writer`.
fn recurse_cmd_output<W: Write>(
    cmd_name: &str,
    cmd_path: &PathBuf,
    output_writer: &mut W,
    cmds_sequence: &Vec<&String>,
) -> Result<()> {
    // Create vector to collect subcommands.
    let mut cmds_list: Vec<String> = Vec::new();

    // Track command level starting from 0, to set command headers' formatting.
    let cmd_level = cmds_sequence.len();

    // Avoid something being horribly structured causing infinite recursion
    if cmd_level > 10 {
        panic!("Too many subcommand layers for {cmd_path:?}: {cmds_sequence:?}");
    }

    // Write out the header.
    let cmd_heading_formatting = "#".repeat(cmd_level + 1);

    // Get terminal output for cmd <subcommands> --help for a given command.
    let lines: Vec<String> = match help_output_for(&cmd_path, &cmds_sequence) {
        Ok(lines) => lines,
        Err(e) => match e {
            HelpError::Ignore => return Ok(()),
            HelpError::Fail(c) => bail!("Error running help: {}", c),
        },
    };

    debug!("Processing {:?} {:?}", cmd_path, cmds_sequence);

    writeln!(output_writer, "{} {}\n", cmd_heading_formatting, cmd_name)?;
    writeln!(output_writer, "{}", CODEBLOCK_START)?;

    // The reference page is broken into a couple sections, and based on the formatting of the
    // lines, these are output directly, or if there is no markdown in the lines, wrapped in a code fence.
    let mut current_section = Section::Top;
    let mut section_contents: HashMap<Section, Vec<String>> = HashMap::new();

    for line in lines {
        // TODO(https://fxbug.dev/42148593): Capture all section headers in addition to "Commands" and "Options".
        match line.to_lowercase().as_str() {
            "subcommands:" | "commands:" => {
                current_section = Section::Commands;
            }
            "options:" | "options" | "flags:" => {
                current_section = Section::Options;
            }
            // Command section ends at a blank line (or end of file).
            "" => {
                if let Some(&mut ref mut contents) = section_contents.get_mut(&current_section) {
                    // two blank lines resets the section to Top
                    if let Some(last_line) = contents.last() {
                        if last_line.is_empty() {
                            current_section = Section::Top;
                            continue;
                        }
                    }
                    contents.push(line);
                } else {
                    let contents = vec![line];
                    section_contents.insert(current_section, contents);
                }
            }
            // Collect sub-commands into a vector.
            _ if current_section == Section::Commands => {
                // Command name is the first word on the line.
                if let Some(command) = line.split_whitespace().next() {
                    match command.as_ref() {
                        "commands" | "subcommands" | "help" => {
                            debug!("skipping {:?} to avoid recursion", command);
                        }
                        _ => {
                            cmds_list.push(command.to_string());
                        }
                    }
                    if let Some(&mut ref mut contents) = section_contents.get_mut(&current_section)
                    {
                        contents.push(line)
                    } else {
                        let contents = vec![line];
                        section_contents.insert(current_section, contents);
                    }
                }
            }
            _ => {
                if line.contains(&cmd_path.as_path().display().to_string()) {
                    let line_no_path =
                        line.replace(&cmd_path.as_path().display().to_string(), &cmd_name);
                    // Write line after stripping full path preceding command name.
                    if let Some(&mut ref mut contents) = section_contents.get_mut(&current_section)
                    {
                        contents.push(line_no_path)
                    } else {
                        let contents = vec![line_no_path];
                        section_contents.insert(current_section, contents);
                    }
                } else if !line.contains("sdk WARN:") {
                    if let Some(&mut ref mut contents) = section_contents.get_mut(&current_section)
                    {
                        contents.push(line)
                    } else {
                        let contents = vec![line];
                        section_contents.insert(current_section, contents);
                    }
                }
            }
        }
    }

    if let Some(contents) = section_contents.get(&Section::Top) {
        for line in contents {
            writeln!(output_writer, "{}", line.trim_end())?;
        }
        writeln!(output_writer, "\n{CODEBLOCK_END}")?;
    }
    if let Some(contents) = section_contents.get(&Section::Options) {
        writeln!(output_writer, "__OPTIONS:__\n\n")?;
        // if there is a non-empty line of the contents starts with a * or `, then treat it as
        // markdown and do not print out the code fence.
        let is_markdown = contents.into_iter().any(|s| {
            !s.is_empty() && (s.trim_start().starts_with("*") || s.trim_start().starts_with("`"))
        });
        if !is_markdown {
            writeln!(output_writer, "{CODEBLOCK_START}")?;
        }
        for line in contents {
            writeln!(output_writer, "{}", line.trim_end())?;
        }
        if !is_markdown {
            writeln!(output_writer, "\n{CODEBLOCK_END}")?;
        }
    }
    if let Some(contents) = section_contents.get(&Section::Commands) {
        writeln!(output_writer, "__SUBCOMMANDS:__\n\n{CODEBLOCK_START}")?;
        for line in contents {
            writeln!(output_writer, "{}", line.trim_end())?;
        }
        writeln!(output_writer, "\n{CODEBLOCK_END}")?;
    }
    cmds_list.sort();

    for cmd in cmds_list {
        // Copy current command sequence and append newest command.
        let mut cmds_sequence = cmds_sequence.clone();
        cmds_sequence.push(&cmd);
        recurse_cmd_output(&cmd, &cmd_path, output_writer, &cmds_sequence)?;
    }

    Ok(())
}

/// Write output of cmd at `cmd_path` to new cmd.md file at `output_path`.
fn write_formatted_output(cmd_path: &PathBuf, output_path: &PathBuf) -> Result<()> {
    // Get name of command from full path to the command executable.
    let cmd_name = cmd_path.file_name().expect("Could not get file name for command");
    let output_md_path = md_path(&cmd_name, &output_path);

    debug!("Generating docs for {:?} to {:?}", cmd_path, output_md_path);

    // Create vector for commands to call in sequence.
    let cmd_sequence = Vec::new();

    // Create a buffer writer to format and write consecutive lines to a file.
    let file = File::create(&output_md_path).context(format!("create {:?}", output_md_path))?;
    let output_writer = &mut BufWriter::new(file);

    let cmd_name = cmd_name.to_str().expect("Could not convert cmd_name from OsStr to str");

    writeln!(output_writer, "{}", HEADER)?;

    // Write output for cmd and all of its subcommands.
    recurse_cmd_output(&cmd_name, &cmd_path, output_writer, &cmd_sequence)
}

/// Generate a vector of full paths to each command in the allow_list.
fn get_command_paths(input_path: &Path, allow_list: &HashSet<OsString>) -> Result<Vec<PathBuf>> {
    // Build a set of all file names in the input_path dir.
    let mut files = HashSet::new();
    if let Ok(paths) = fs::read_dir(&input_path) {
        for path in paths {
            if let Ok(path) = path {
                files.insert(path.file_name());
            }
        }
    }

    // Get the intersection of all files and commands in the allow_list.
    let commands: HashSet<_> = files.intersection(&allow_list).collect();
    info!("Including tools: {:?}", commands);

    // Build full paths to allowed commands found in the input_path dir.
    let mut cmd_paths = Vec::new();
    for c in commands.iter() {
        let path = Path::new(&input_path).join(c);
        cmd_paths.push(path);
    }
    Ok(cmd_paths)
}

/// Create the output dir if doesn't exist, recursively creating subdirs in path.
fn create_output_dir(path: &Path) -> Result<()> {
    if !path.exists() {
        fs::create_dir_all(path)
            .with_context(|| format!("Unable to create output directory {}", path.display()))?;
        info!("Created directory {}", path.display());
    }
    Ok(())
}

/// Get cmd --help output when given a full path to a cmd.
fn help_output_for(tool: &Path, subcommands: &Vec<&String>) -> Result<Vec<String>, HelpError> {
    let output = Command::new(&tool)
        .args(&*subcommands)
        .arg("--help")
        .output()
        .context(format!("Command failed for {:?}", &tool.display()))
        .expect("get output");

    let stdout = output.stdout;
    let stderr = output.stderr;
    let exit_code = output.status.code().expect("get help status code");

    // Convert string outputs to vector of lines.
    let stdout_string = String::from_utf8(stdout).expect("Help string from utf8");
    let mut combined_lines = stdout_string.lines().map(String::from).collect::<Vec<_>>();

    let stderr_string = String::from_utf8(stderr).expect("Help string from utf8");
    let stderr_lines = stderr_string.lines().map(String::from).collect::<Vec<_>>();
    let stderr_empty = stderr_lines.is_empty();
    combined_lines.extend(stderr_lines);
    if !combined_lines.is_empty() {
        // TODO(https://fxbug.dev/42166743): This is a short term solution to prevent errantly documenting
        // run-on sentences as args with ffx. Long term solution involves using help-json.
        let first_line = &combined_lines[0];
        if first_line.contains("Unrecognized argument:") {
            // TODO(https://fxbug.dev/42172027): Create Error enums to better fit these errors.
            return Err(HelpError::Ignore);
        }
    }
    if !stderr_empty && exit_code != 0 {
        let tool_name = tool.file_name().expect("get tool name");
        if IGNORE_ERR_CODE.get(tool_name.to_str().expect("get tool str")) == None {
            let error_message = format!(
                "Unexpected non-zero error code with tool {:?}
            and subcommands {:?}.",
                tool.display(),
                subcommands
            );
            return Err(HelpError::Fail(error_message));
        }
    }

    Ok(combined_lines)
}

/// Given a cmd name and a dir, create a full path ending in cmd.md.
pub(crate) fn md_path(file_stem: &OsStr, dir: &PathBuf) -> PathBuf {
    let mut path = Path::new(dir).join(file_stem);
    path.set_extension("md");
    path
}

#[cfg(test)]
mod tests {
    use super::*;
    use flate2::read::GzDecoder;
    use tar::Archive;

    #[test]
    fn run_test_commands() {
        let tmp_dir = tempfile::Builder::new().prefix("clidoc-test-out").tempdir().unwrap();
        let argv = [
            "-v",
            "-o",
            &tmp_dir.path().to_str().unwrap(),
            "clidoc_test_data/tool_with_subcommands.sh",
        ];
        let cmd = "clidoc-test";
        let opt = Opt::from_args(&[cmd], &argv).unwrap();
        let generated = tmp_dir.path().join("tool_with_subcommands.md");
        let expected = &opt.in_dir.join("clidoc_test_data/tool_with_subcommands.md");

        run(opt).expect("tool_with_subcommands could not be generated");

        let generated_contents = fs::read_to_string(generated).unwrap();
        let expected_contents = fs::read_to_string(expected).unwrap();

        assert_eq!(generated_contents, expected_contents);
    }

    #[test]
    fn run_test_archive_and_cleanup() {
        let tmp_dir = tempfile::Builder::new().prefix("clidoc-tar-test").tempdir().unwrap();
        let argv = [
            "--archive-path",
            "clidoc_archive.tar.gz",
            "-v",
            "-o",
            &tmp_dir.path().to_str().unwrap(),
            "clidoc_test_data/tool_with_subcommands.sh",
        ];
        let cmd = "clidoc-test-archive";
        let opt = Opt::from_args(&[cmd], &argv).unwrap();
        run(opt).expect("tool_with_subcommands could not be generated");

        // With the archive-path flag set, the md file should be zipped
        // and not exist.
        assert!(!tmp_dir.path().join("tool_with_subcommands.md").exists());

        let tar_gz = File::open("clidoc_archive.tar.gz").expect("open tarball");
        let tar = GzDecoder::new(tar_gz);
        let mut archive = Archive::new(tar);
        archive.unpack(".").expect("extract tar");

        assert!(Path::new("clidoc/tool_with_subcommands.md").exists());
    }
}