//! Utility for building and running tests against mdbook.

use mdbook_core::utils::fs;
use mdbook_driver::MDBook;
use mdbook_driver::init::BookBuilder;
use snapbox::IntoData;
use std::collections::BTreeMap;
use std::path::{Path, PathBuf};
use std::process::Command;
use std::sync::atomic::{AtomicU32, Ordering};

/// Test number used for generating unique temp directory names.
static NEXT_TEST_ID: AtomicU32 = AtomicU32::new(0);

#[derive(Clone, Copy, Eq, PartialEq)]
enum StatusCode {
    Success,
    Failure,
    Code(i32),
}

/// Main helper for driving mdbook tests.
pub struct BookTest {
    /// The temp directory where the test should perform its work.
    pub dir: PathBuf,
    /// The original source directory if created from [`BookTest::from_dir`].
    original_source: Option<PathBuf>,
    /// Snapshot assertion support.
    pub assert: snapbox::Assert,
    /// This indicates whether or not the book has been built.
    built: bool,
}

impl BookTest {
    /// Creates a new test, copying the contents from the given directory into
    /// a temp directory.
    pub fn from_dir(dir: &str) -> BookTest {
        // Copy this test book to a temp directory.
        let dir = Path::new("tests/testsuite").join(dir);
        assert!(dir.exists(), "{dir:?} should exist");
        let tmp = Self::new_tmp();
        mdbook_core::utils::fs::copy_files_except_ext(
            &dir,
            &tmp,
            true,
            Some(&PathBuf::from("book")),
            &[],
        )
        .unwrap_or_else(|e| panic!("failed to copy test book {dir:?} to {tmp:?}: {e:?}"));
        Self::new(tmp, Some(dir))
    }

    /// Creates a new test with an empty temp directory.
    pub fn empty() -> BookTest {
        Self::new(Self::new_tmp(), None)
    }

    /// Creates a new test with the given function to initialize a new book.
    ///
    /// The book itself is not built.
    pub fn init(f: impl Fn(&mut BookBuilder)) -> BookTest {
        let tmp = Self::new_tmp();
        let mut bb = MDBook::init(&tmp);
        f(&mut bb);
        bb.build()
            .unwrap_or_else(|e| panic!("failed to initialize book at {tmp:?}: {e:?}"));
        Self::new(tmp, None)
    }

    fn new_tmp() -> PathBuf {
        let id = NEXT_TEST_ID.fetch_add(1, Ordering::SeqCst);
        let tmp = Path::new(env!("CARGO_TARGET_TMPDIR"))
            .join("ts")
            .join(format!("t{id}"));
        if tmp.exists() {
            std::fs::remove_dir_all(&tmp)
                .unwrap_or_else(|e| panic!("failed to remove {tmp:?}: {e:?}"));
        }
        fs::create_dir_all(&tmp).unwrap();
        tmp
    }

    fn new(dir: PathBuf, original_source: Option<PathBuf>) -> BookTest {
        let assert = assert(&dir);
        BookTest {
            dir,
            original_source,
            assert,
            built: false,
        }
    }

    /// Checks the contents of an HTML file that it has the given contents
    /// between the `<main>` tag.
    ///
    /// Normally the contents outside of the `<main>` tag aren't interesting,
    /// and they add a significant amount of noise.
    #[track_caller]
    pub fn check_main_file(&mut self, path: &str, expected: impl IntoData) -> &mut Self {
        if !self.built {
            self.build();
        }
        let full_path = self.dir.join(path);
        let actual = read_to_string(&full_path);
        let start = actual
            .find("<main>")
            .unwrap_or_else(|| panic!("didn't find <main> for `{full_path:?}` in:\n{actual}"));
        let end = actual.find("</main>").unwrap();
        let contents = actual[start + 6..end - 7].trim();
        self.assert.eq(contents, expected);
        self
    }

    /// Verifies the HTML output of all chapters in a book.
    ///
    /// This calls [`BookTest::check_main_file`] for every `.html` file
    /// generated by building the book. All of expected files are stored in a
    /// director called "expected" in the original book source directory.
    ///
    /// This only works when created with [`BookTest::from_dir`].
    ///
    /// `404.html`, `print.html`, and `toc.html` are not validated. The root
    /// `index.html` is also not validated (since it is often duplicated with
    /// the first chapter). If you need to validate it, call
    /// [`BookTest::check_main_file`] directly.
    #[track_caller]
    pub fn check_all_main_files(&mut self) -> &mut Self {
        if !self.built {
            self.build();
        }
        let book_root = self.dir.join("book");
        let mut files = list_all_files(&book_root);
        files.retain(|file| {
            file.extension().is_some_and(|ext| ext == "html")
                && !matches!(
                    file.to_str().unwrap(),
                    "index.html" | "404.html" | "print.html" | "toc.html"
                )
        });
        let expected_path = self
            .original_source
            .as_ref()
            .expect("created with BookTest::from_dir")
            .join("expected");
        let mut expected_list = list_all_files(&expected_path);
        for file in &files {
            let expected = expected_path.join(file);
            let data = snapbox::Data::read_from(&expected, None);
            self.check_main_file(book_root.join(file).to_str().unwrap(), data);
            if let Some(i) = expected_list.iter().position(|p| p == file) {
                expected_list.remove(i);
            }
        }
        // Verify there aren't any unused expected files.
        if !expected_list.is_empty() {
            panic!(
                "extra expected files found in `{expected_path:?}:\n\
                {expected_list:#?}\n\
                Verify that these files are no longer needed and delete them."
            );
        }
        self
    }

    /// Checks the summary contents of `toc.js` against the expected value.
    #[track_caller]
    pub fn check_toc_js(&mut self, expected: impl IntoData) -> &mut Self {
        if !self.built {
            self.build();
        }
        let inner = self.toc_js_html();
        // Would be nice if this were prettified, but a primitive wrapping will do for now.
        let inner = inner.replace("><", ">\n<");
        self.assert.eq(inner, expected);
        self
    }

    /// Returns the summary contents from `toc.js`.
    #[track_caller]
    pub fn toc_js_html(&self) -> String {
        let toc_path = glob_one(&self.dir, "book/toc*.js");
        let actual = read_to_string(&toc_path);
        let inner = actual
            .lines()
            .filter_map(|line| {
                let line = line.trim().strip_prefix("this.innerHTML = '")?;
                let line = line.strip_suffix("';")?;
                Some(line)
            })
            .next()
            .expect("should have innerHTML");
        inner.to_string()
    }

    /// Checks that the contents of the given file matches the expected value.
    ///
    /// The path can use glob-style wildcards, but it must match only a single file.
    #[track_caller]
    pub fn check_file(&mut self, path_pattern: &str, expected: impl IntoData) -> &mut Self {
        if !self.built {
            self.build();
        }
        let path = glob_one(&self.dir, path_pattern);
        let actual = read_to_string(&path);
        self.assert.eq(actual, expected);
        self
    }

    /// Checks that the given file contains the given [`snapbox::Assert`] pattern somewhere.
    ///
    /// The path can use glob-style wildcards, but it must match only a single file.
    #[track_caller]
    pub fn check_file_contains(&mut self, path_pattern: &str, expected: &str) -> &mut Self {
        if !self.built {
            self.build();
        }
        let path = glob_one(&self.dir, path_pattern);
        let actual = read_to_string(&path);
        let expected = format!("...\n[..]{expected}[..]\n...\n");
        self.assert.eq(actual, expected);
        self
    }

    /// Checks that the given file does not contain the given string anywhere.
    ///
    /// Beware that using this is fragile, as it may be unable to catch
    /// regressions (it can't tell the difference between success, or the
    /// string being looked for changed).
    ///
    /// The path can use glob-style wildcards, but it must match only a single file.
    #[track_caller]
    pub fn check_file_doesnt_contain(&mut self, path_pattern: &str, string: &str) -> &mut Self {
        if !self.built {
            self.build();
        }
        let path = glob_one(&self.dir, path_pattern);
        let actual = read_to_string(&path);
        assert!(
            !actual.contains(string),
            "Unexpectedly found {string:?} in {path:?}\n\n{actual}",
        );
        self
    }

    /// Checks that the list of files at the given path matches the given value.
    #[track_caller]
    pub fn check_file_list(&mut self, path: &str, expected: impl IntoData) -> &mut Self {
        let mut all_paths: Vec<_> = walkdir::WalkDir::new(&self.dir.join(path))
            .into_iter()
            // Skip the outer directory.
            .skip(1)
            .map(|e| {
                e.unwrap()
                    .into_path()
                    .strip_prefix(&self.dir)
                    .unwrap()
                    .to_str()
                    .unwrap()
                    .replace('\\', "/")
            })
            .collect();
        all_paths.sort();
        let actual = all_paths.join("\n");
        self.assert.eq(actual, expected);
        self
    }

    /// Loads an [`MDBook`] from the temp directory.
    pub fn load_book(&self) -> MDBook {
        MDBook::load(&self.dir).unwrap_or_else(|e| panic!("book failed to load: {e:?}"))
    }

    /// Builds the book in the temp directory.
    pub fn build(&mut self) -> &mut Self {
        let book = self.load_book();
        book.build()
            .unwrap_or_else(|e| panic!("book failed to build: {e:?}"));
        self.built = true;
        self
    }

    /// Runs the `mdbook` binary in the temp directory.
    ///
    /// This runs `mdbook` with the given args. The args are split on spaces
    /// (if you need args with spaces, use the `args` method). The given
    /// callback receives a [`BookCommand`] for you to customize how the
    /// executable is run.
    pub fn run(&mut self, args: &str, f: impl Fn(&mut BookCommand)) -> &mut Self {
        let mut cmd = BookCommand {
            assert: self.assert.clone(),
            dir: self.dir.clone(),
            args: split_args(args),
            env: BTreeMap::new(),
            expect_status: StatusCode::Success,
            expect_stderr_data: None,
            expect_stdout_data: None,
            debug: None,
        };
        f(&mut cmd);
        cmd.run();
        // Ensure that `built` gets set if a build command is used so that all
        // the `check` methods do not overwrite the contents of what was just
        // built.
        if cmd.args.first().map(String::as_str) == Some("build") {
            self.built = true
        }
        self
    }

    /// Change a file's contents in the given path.
    pub fn change_file(&mut self, path: impl AsRef<Path>, body: &str) -> &mut Self {
        let path = self.dir.join(path);
        fs::write(&path, body).unwrap();
        self
    }

    /// Removes a file or directory relative to the test root.
    pub fn rm_r(&mut self, path: impl AsRef<Path>) -> &mut Self {
        let path = self.dir.join(path.as_ref());
        let meta = match path.symlink_metadata() {
            Ok(meta) => meta,
            Err(e) => panic!("failed to remove {path:?}, could not read: {e:?}"),
        };
        // There is a race condition between fetching the metadata and
        // actually performing the removal, but we don't care all that much
        // for our tests.
        if meta.is_dir() {
            if let Err(e) = std::fs::remove_dir_all(&path) {
                panic!("failed to remove {path:?}: {e:?}");
            }
        } else if let Err(e) = std::fs::remove_file(&path) {
            panic!("failed to remove {path:?}: {e:?}")
        }
        self
    }

    /// Builds a Rust program with the given src.
    ///
    /// The given path should be the path where to output the executable in
    /// the temp directory.
    pub fn rust_program(&mut self, path: &str, src: &str) -> &mut Self {
        let rs = self.dir.join(path).with_extension("rs");
        let parent = rs.parent().unwrap();
        if !parent.exists() {
            fs::create_dir_all(&parent).unwrap();
        }
        fs::write(&rs, src).unwrap();
        let status = std::process::Command::new("rustc")
            .arg(&rs)
            .current_dir(&parent)
            .status()
            .expect("rustc should run");
        assert!(status.success());
        self
    }
}

/// A builder for preparing to run the `mdbook` executable.
///
/// By default, it expects the process to succeed.
pub struct BookCommand {
    pub dir: PathBuf,
    assert: snapbox::Assert,
    args: Vec<String>,
    env: BTreeMap<String, Option<String>>,
    expect_status: StatusCode,
    expect_stderr_data: Option<snapbox::Data>,
    expect_stdout_data: Option<snapbox::Data>,
    debug: Option<String>,
}

impl BookCommand {
    /// Indicates that the process should fail.
    pub fn expect_failure(&mut self) -> &mut Self {
        self.expect_status = StatusCode::Failure;
        self
    }

    /// Indicates the process should fail with the given exit code.
    pub fn expect_code(&mut self, code: i32) -> &mut Self {
        self.expect_status = StatusCode::Code(code);
        self
    }

    /// Verifies that stderr matches the given value.
    pub fn expect_stderr(&mut self, expected: impl snapbox::IntoData) -> &mut Self {
        self.expect_stderr_data = Some(expected.into_data());
        self
    }

    /// Verifies that stdout matches the given value.
    pub fn expect_stdout(&mut self, expected: impl snapbox::IntoData) -> &mut Self {
        self.expect_stdout_data = Some(expected.into_data());
        self
    }

    /// Adds arguments to the command to run.
    pub fn args(&mut self, args: &[&str]) -> &mut Self {
        self.args.extend(args.into_iter().map(|t| t.to_string()));
        self
    }

    /// Specifies an environment variable to set on the executable.
    pub fn env<T: Into<String>>(&mut self, key: &str, value: T) -> &mut Self {
        self.env.insert(key.to_string(), Some(value.into()));
        self
    }

    /// Sets the directory used for running the command.
    pub fn current_dir<S: AsRef<std::path::Path>>(&mut self, path: S) -> &mut Self {
        self.dir = self.dir.join(path.as_ref());
        self
    }

    /// Use this to debug a command.
    ///
    /// Pass the value that you would normally pass to `MDBOOK_LOG`, and this
    /// will enable logging, print the command that runs and its output.
    ///
    /// This will fail if you use it in CI.
    #[allow(unused)]
    pub fn debug(&mut self, value: &str) -> &mut Self {
        if std::env::var_os("CI").is_some() {
            panic!("debug is not allowed on CI");
        }
        self.debug = Some(value.into());
        self
    }

    /// Runs the command, and verifies the output.
    fn run(&mut self) {
        let mut cmd = Command::new(env!("CARGO_BIN_EXE_mdbook"));
        cmd.current_dir(&self.dir)
            .args(&self.args)
            .env_remove("MDBOOK_LOG")
            // Don't read the system git config which is out of our control.
            .env("GIT_CONFIG_NOSYSTEM", "1")
            .env("GIT_CONFIG_GLOBAL", &self.dir)
            .env("GIT_CONFIG_SYSTEM", &self.dir)
            .env_remove("GIT_AUTHOR_EMAIL")
            .env_remove("GIT_AUTHOR_NAME")
            .env_remove("GIT_COMMITTER_EMAIL")
            .env_remove("GIT_COMMITTER_NAME");

        if let Some(debug) = &self.debug {
            cmd.env("MDBOOK_LOG", debug);
        }

        for (k, v) in &self.env {
            match v {
                Some(v) => cmd.env(k, v),
                None => cmd.env_remove(k),
            };
        }

        if self.debug.is_some() {
            eprintln!("running {cmd:#?}");
        }
        let output = cmd.output().expect("mdbook should be runnable");
        let stdout = std::str::from_utf8(&output.stdout).expect("stdout is not utf8");
        let stderr = std::str::from_utf8(&output.stderr).expect("stderr is not utf8");
        let render_output = || format!("\n--- stdout\n{stdout}\n--- stderr\n{stderr}");
        match (self.expect_status, output.status.success()) {
            (StatusCode::Success, false) => {
                panic!("mdbook failed, but expected success{}", render_output())
            }
            (StatusCode::Failure, true) => {
                panic!("mdbook succeeded, but expected failure{}", render_output())
            }
            (StatusCode::Code(expected), _) => match output.status.code() {
                Some(actual) => assert_eq!(
                    actual, expected,
                    "process exit code did not match as expected"
                ),
                None => panic!("process exited via signal {:?}", output.status),
            },
            _ => {}
        }
        if self.debug.is_some() {
            eprintln!("{}", render_output());
        }
        self.expect_status = StatusCode::Success; // Reset to default.
        if let Some(expect_stderr_data) = &self.expect_stderr_data {
            if let Err(e) = self.assert.try_eq(
                Some(&"stderr"),
                stderr.into_data(),
                expect_stderr_data.clone(),
            ) {
                panic!("{e}");
            }
        }
        if let Some(expect_stdout_data) = &self.expect_stdout_data {
            if let Err(e) = self.assert.try_eq(
                Some(&"stdout"),
                stdout.into_data(),
                expect_stdout_data.clone(),
            ) {
                panic!("{e}");
            }
        }
    }
}

fn split_args(s: &str) -> Vec<String> {
    s.split_whitespace()
        .map(|arg| {
            if arg.contains(&['"', '\''][..]) {
                panic!("shell-style argument parsing is not supported");
            }
            String::from(arg)
        })
        .collect()
}

static LITERAL_REDACTIONS: &[(&str, &str)] = &[
    // Unix message for an entity was not found
    ("[NOT_FOUND]", "No such file or directory (os error 2)"),
    // Windows message for an entity was not found
    (
        "[NOT_FOUND]",
        "The system cannot find the file specified. (os error 2)",
    ),
    (
        "[NOT_FOUND]",
        "The system cannot find the path specified. (os error 3)",
    ),
    ("[NOT_FOUND]", "program not found"),
    // Unix message for exit status
    ("[EXIT_STATUS]", "exit status"),
    // Windows message for exit status
    ("[EXIT_STATUS]", "exit code"),
    ("[TAB]", "\t"),
    ("[EXE]", std::env::consts::EXE_SUFFIX),
];

fn assert(root: &Path) -> snapbox::Assert {
    let mut subs = snapbox::Redactions::new();
    subs.insert("[ROOT]", root.to_path_buf()).unwrap();
    subs.insert("[VERSION]", mdbook_core::MDBOOK_VERSION)
        .unwrap();

    subs.extend(LITERAL_REDACTIONS.into_iter().cloned())
        .unwrap();

    snapbox::Assert::new()
        .action_env(snapbox::assert::DEFAULT_ACTION_ENV)
        .redact_with(subs)
}

/// Helper to read a string from the filesystem.
#[track_caller]
pub fn read_to_string<P: AsRef<Path>>(path: P) -> String {
    let path = path.as_ref();
    fs::read_to_string(path).unwrap()
}

/// Returns the first path from the given glob pattern.
pub fn glob_one<P: AsRef<Path>>(path: P, pattern: &str) -> PathBuf {
    let path = path.as_ref();
    let mut matches = glob::glob(path.join(pattern).to_str().unwrap()).unwrap();
    let Some(first) = matches.next() else {
        panic!("expected at least one file at `{path:?}` with pattern `{pattern}`, found none");
    };
    let first = first.unwrap();
    if let Some(next) = matches.next() {
        panic!(
            "expected only one file for pattern `{pattern}` in `{path:?}`, \
             found `{first:?}` and `{:?}`",
            next.unwrap()
        );
    }
    first
}

/// Lists all files at the given directory.
///
/// Recursively walks the tree. Paths are relative to the directory.
pub fn list_all_files(dir: &Path) -> Vec<PathBuf> {
    walkdir::WalkDir::new(dir)
        .sort_by_file_name()
        .into_iter()
        // Skip the outer directory.
        .skip(1)
        .map(|entry| {
            let entry = entry.unwrap();
            let path = entry.path();
            path.strip_prefix(dir).unwrap().to_path_buf()
        })
        .collect()
}
