#[macro_use]
extern crate lazy_static;
pub mod deqp_command;
pub mod gtest_command;
pub mod mock_deqp;
pub mod mock_gtest;
pub mod mock_piglit;
mod parse_deqp;
pub mod parse_piglit;
pub mod piglit_command;
mod runner_results;

use anyhow::bail;
pub use runner_results::*;

use anyhow::{Context, Result};
use log::*;
use parse_deqp::DeqpTestResult;
use piglit_command::*;
use rand::rngs::StdRng;
use rand::seq::SliceRandom;
use rand::SeedableRng;
use rayon::prelude::*;
use regex::RegexSet;
use std::collections::HashMap;
use std::fs::File;
use std::io::prelude::*;
use std::io::BufReader;
use std::path::{Path, PathBuf};
use std::sync::mpsc::{channel, Receiver};
use std::time::Duration;
use std::time::Instant;
use structopt::StructOpt;

pub fn parse_key_val<T, U>(s: &str) -> Result<(T, U), Box<dyn std::error::Error>>
where
    T: std::str::FromStr,
    T::Err: std::error::Error + 'static,
    U: std::str::FromStr,
    U::Err: std::error::Error + 'static,
{
    let pos = s
        .find('=')
        .ok_or_else(|| format!("invalid KEY=value: no `=` found in `{}`", s))?;
    Ok((s[..pos].parse()?, s[pos + 1..].parse()?))
}

#[derive(Debug, StructOpt)]
pub struct CommandLineRunOptions {
    #[structopt(long = "output", help = "path to output directory")]
    pub output_dir: PathBuf,

    #[structopt(
        long,
        help = "path to baseline results (such as output/failures.csv from another run)"
    )]
    pub baseline: Option<PathBuf>,

    #[structopt(
        long,
        help = "path to file of regexes of tests to skip running (for runtime or stability reasons)"
    )]
    pub skips: Vec<PathBuf>,

    #[structopt(
        short = "t",
        long = "include-tests",
        help = "regexes of tests to include (non-matching tests are skipped)"
    )]
    pub include: Vec<String>,

    #[structopt(
        long,
        help = "path to file of regexes of tests to assume any failures in those tests are flaky results (but still run them, for long-term status tracking)"
    )]
    pub flakes: Vec<PathBuf>,

    #[structopt(
        long,
        default_value = "60.0",
        help = "per-test timeout in floating point seconds"
    )]
    pub timeout: f32,

    #[structopt(
        short = "j",
        long,
        default_value = "0",
        help = "Number of processes to invoke in parallel (default 0 = number of CPUs in system)"
    )]
    pub jobs: usize,

    #[structopt(long, default_value = "1", help = "Runs 1 out of every N tests.")]
    pub fraction: usize,

    #[structopt(
        long,
        default_value = "1",
        help = "Skips the first N-1 tests in the test list before applying --fraction (useful for running N/M fraciton of the test list across multiple devices)."
    )]
    pub fraction_start: usize,

    #[structopt(
        parse(try_from_str = parse_key_val),
        long = "env",
        help = "Environment variables to set when invoking the test process"
    )]
    pub env: Vec<(String, String)>,

    #[structopt(
        long,
        default_value = "25",
        help = "Number of fails or flakes to print in the summary line (0 = no limit)"
    )]
    pub summary_limit: usize,

    #[structopt(
        parse(from_occurrences),
        short = "v",
        long,
        help = "Enable verbose mode (-v, -vv, -vvv, etc)"
    )]
    pub verbose: usize,

    #[structopt(long, help = "Enable log timestamps (sec, ms, ns)")]
    pub timestamp: Option<stderrlog::Timestamp>,
    #[structopt(
        long,
        help = "Saves log files for expected failures along with new ones"
    )]
    pub save_xfail_logs: bool,
}

impl CommandLineRunOptions {
    pub fn setup(&self) -> Result<()> {
        stderrlog::new()
            .module(module_path!())
            .verbosity(self.verbose)
            .timestamp(self.timestamp.unwrap_or(stderrlog::Timestamp::Off))
            .init()
            .unwrap();

        if self.jobs > 0 {
            rayon::ThreadPoolBuilder::new()
                .num_threads(self.jobs)
                .build_global()
                .unwrap();
        }

        if self.fraction < 1 {
            eprintln!("--fraction must be >= 1.");
            std::process::exit(1);
        }
        if self.fraction_start < 1 {
            eprintln!("--fraction_start must be >= 1.");
            std::process::exit(1);
        }

        std::fs::create_dir_all(&self.output_dir).context("creating output directory")?;

        Ok(())
    }

    pub fn baseline(&self) -> Result<RunnerResults> {
        read_baseline(self.baseline.as_ref())
    }

    pub fn skips_regex(&self) -> Result<RegexSet> {
        parse_regex_set(read_lines(&self.skips)?).context("compiling skips regexes")
    }

    pub fn flakes_regex(&self) -> Result<RegexSet> {
        parse_regex_set(read_lines(&self.flakes)?).context("compiling flakes regexes")
    }

    pub fn includes_regex(&self) -> Result<RegexSet> {
        if self.include.is_empty() {
            RegexSet::new(vec![""]).context("compiling all-tests include RE")
        } else {
            parse_regex_set(&self.include).context("compiling include filters")
        }
    }
}

pub struct TestConfiguration {
    pub output_dir: PathBuf,
    pub skips: RegexSet,
    pub flakes: RegexSet,
    pub baseline: RunnerResults,
    pub timeout: Duration,
    pub env: HashMap<String, String>,
    pub save_xfail_logs: bool,
}

pub trait TestCommand {
    fn config(&self) -> &TestConfiguration;

    fn run<S: AsRef<TestCase>, I: IntoIterator<Item = S>>(
        &self,
        caselist_state: &CaselistState,
        tests: I,
    ) -> Result<Vec<RunnerResult>>;

    fn see_more(&self, _name: &str, _caselist_state: &CaselistState) -> String {
        "".to_string()
    }

    fn skips(&self) -> &RegexSet {
        &self.config().skips
    }

    fn flakes(&self) -> &RegexSet {
        &self.config().flakes
    }

    fn baseline(&self) -> &RunnerResults {
        &self.config().baseline
    }

    fn baseline_status<S: AsRef<str>>(&self, test: S) -> Option<RunnerStatus> {
        self.baseline().tests.get(test.as_ref()).map(|x| x.status)
    }

    fn translate_result(
        &self,
        result: &DeqpTestResult,
        caselist_state: &CaselistState,
    ) -> RunnerStatus {
        let mut status = RunnerStatus::from_deqp(result.status)
            .with_baseline(self.baseline_status(&result.name));

        if !status.is_success() && self.flakes().is_match(&result.name) {
            status = RunnerStatus::Flake;
        }

        if !status.is_success() || status == RunnerStatus::Flake {
            error!(
                "Test {}: {}: {}",
                &result.name,
                status,
                self.see_more(&result.name, &caselist_state)
            );
        }

        status
    }

    fn skip_test(&self, test: &str) -> bool {
        self.skips().is_match(test)
    }

    fn run_caselist_and_flake_detect(
        &self,
        caselist: &[TestCase],
        caselist_state: &mut CaselistState,
    ) -> Result<Vec<RunnerResult>> {
        // Sort the caselists within test groups.  dEQP runs tests in sorted order, and when one
        // is debugging a failure in one case in a caselist, it can be nice to be able to easily trim
        // all of the caselist appearing after the failure, to reduce runtime.
        let mut caselist: Vec<_> = caselist.iter().collect();
        caselist.sort_by(|x, y| x.name().cmp(y.name()));

        caselist_state.run_id += 1;
        let mut results = self.run(&caselist_state, &caselist)?;
        // If we made no more progress on the whole caselist,
        // then dEQP doesn't know about some of our tests and they'll report Missing.
        if results.is_empty() {
            anyhow::bail!(
                "No results parsed.  Is your caselist out of sync with your deqp binary?"
            );
        }

        // If any results came back with an unexpected failure, run the caselist again
        // to see if we get the same results, and mark any changing results as flaky tests.
        if results.iter().any(|x| !x.status.is_success()) {
            caselist_state.run_id += 1;
            let retest_results = self.run(&caselist_state, &caselist)?;
            for pair in results.iter_mut().zip(retest_results.iter()) {
                if pair.0.status != pair.1.status {
                    pair.0.status = RunnerStatus::Flake;
                }
            }
        }

        Ok(results)
    }

    fn process_caselist<S: AsRef<TestCase>, I: IntoIterator<Item = S>>(
        &self,
        tests: I,
        caselist_id: u32,
    ) -> Result<Vec<RunnerResult>> {
        let mut caselist_results: Vec<RunnerResult> = Vec::new();
        let mut remaining_tests = Vec::new();
        for test in tests {
            let test = test.as_ref().clone();
            if self.skip_test(format!("{}{}", self.prefix(), &test.name()).as_str()) {
                caselist_results.push(RunnerResult {
                    test: format!("{}{}", self.prefix(), test.name()),
                    status: RunnerStatus::Skip,
                    duration: Default::default(),
                    subtest: false,
                });
            } else {
                remaining_tests.push(test);
            }
        }

        let mut caselist_state = CaselistState {
            caselist_id,
            run_id: 0,
        };

        while !remaining_tests.is_empty() {
            let results = self.run_caselist_and_flake_detect(&remaining_tests, &mut caselist_state);

            match results {
                Ok(results) => {
                    for result in results {
                        /* Remove the reported test from our list of tests to run.  If it's not in our list, then it's
                         * a subtest.
                         */
                        if let Some(position) = remaining_tests
                            .iter()
                            .position(|x| x.name() == result.test.trim_start_matches(self.prefix()))
                        {
                            remaining_tests.swap_remove(position);
                        } else if !result.subtest {
                            error!(
                                "Top-level test result for {} not found in list of tests to run.",
                                &result.test
                            );
                        }

                        caselist_results.push(result);
                    }
                }
                Err(e) => {
                    error!(
                        "Failure getting run results: {:#} ({})",
                        e,
                        self.see_more("", &caselist_state)
                    );

                    for test in remaining_tests {
                        caselist_results.push(RunnerResult {
                            test: format!("{}{}", self.prefix(), &test.name()),
                            status: RunnerStatus::Missing,
                            duration: Default::default(),
                            subtest: false,
                        });
                    }
                    break;
                }
            }
        }

        Ok(caselist_results)
    }

    fn split_tests_to_groups(
        &self,
        mut tests: Vec<TestCase>,
        tests_per_group: usize,
        min_tests_per_group: usize,
    ) -> Result<Vec<(&Self, Vec<TestCase>)>> {
        if tests_per_group < 1 {
            bail!("tests_per_group must be >= 1.");
        }

        // If you haven't requested the scaling-down behavior, make all groups
        // the same size.
        let min_tests_per_group = if min_tests_per_group == 0 {
            tests_per_group
        } else {
            min_tests_per_group
        };

        let rayon_threads = rayon::current_num_threads();
        let tests_per_group = usize::max(
            1,
            usize::min(
                (tests.len() + rayon_threads - 1) / rayon_threads,
                tests_per_group,
            ),
        );

        // Shuffle the test groups using a deterministic RNG so that every run gets the same shuffle.
        tests.shuffle(&mut StdRng::from_seed([0x3bu8; 32]));

        // Make test groups of tests_per_group() (512) tests, or if
        // min_tests_per_group() is lower than that, then 1/32nd of the
        // remaining tests down to that limit.
        let mut test_groups: Vec<(&Self, Vec<TestCase>)> = Vec::new();
        let mut remaining = tests.len();
        while remaining != 0 {
            let min = usize::min(min_tests_per_group, remaining);
            let group_len = usize::min(usize::max(remaining / 32, min), tests_per_group);
            remaining -= group_len;

            test_groups.push((self, tests.split_off(remaining)));
        }

        Ok(test_groups)
    }

    fn caselist_file_path(&self, caselist_state: &CaselistState, suffix: &str) -> Result<PathBuf> {
        // deqp must be run from its directory, so make sure all the filenames we pass in are absolute.
        let output_dir = self.config().output_dir.canonicalize()?;

        Ok(output_dir.join(format!(
            "c{}.r{}.{}",
            caselist_state.caselist_id, caselist_state.run_id, suffix
        )))
    }

    fn prefix(&self) -> &str {
        ""
    }
}

#[derive(Clone, Debug, PartialEq)]
pub enum TestCase {
    Deqp(String),
    GTest(String),
    Piglit(PiglitTest),
}

impl TestCase {
    pub fn name(&self) -> &str {
        match self {
            TestCase::Deqp(name) => &name,
            TestCase::GTest(name) => &name,
            TestCase::Piglit(test) => &test.name,
        }
    }
}

impl AsRef<str> for TestCase {
    fn as_ref(&self) -> &str {
        self.name()
    }
}

impl AsRef<TestCase> for TestCase {
    fn as_ref(&self) -> &TestCase {
        self
    }
}

fn results_collection<W: Write>(
    status_output: &mut W,
    run_results: &mut RunnerResults,
    total_tests: u32,
    receiver: Receiver<Result<Vec<RunnerResult>>>,
) {
    let update_interval = Duration::new(2, 0);

    run_results.status_update(status_output, total_tests);
    let mut last_status_update = Instant::now();

    for group_results in receiver {
        match group_results {
            Ok(group_results) => {
                for result in group_results {
                    if run_results.tests.contains_key(&result.test) {
                        error!(
                            "Duplicate test result for {}, marking test failed",
                            &result.test
                        );
                        let mut fail_result = result;
                        fail_result.status = RunnerStatus::Fail;
                        run_results.record_result(fail_result);
                    } else {
                        run_results.record_result(result);
                    }
                }
            }
            Err(e) => {
                println!("Error: {}", e);
            }
        }
        if last_status_update.elapsed() >= update_interval {
            run_results.status_update(status_output, total_tests);
            last_status_update = Instant::now();
        }
    }

    // Always print the final results
    run_results.status_update(status_output, total_tests);
}

// Splits the list of tests to groups and parallelize them across all cores, collecting results in
// a separate thread
pub fn parallel_test<D, W>(
    status_output: W,
    test_groups: Vec<(&D, Vec<TestCase>)>,
) -> Result<RunnerResults>
where
    D: TestCommand,
    D: Sync,
    W: Write,
    W: Sync,
    W: Send,
{
    let test_count = test_groups.iter().map(|x| x.1.len() as u32).sum();

    let mut run_results = RunnerResults::new();

    // Make a channel for the parallel iterator to send results to, which is what will be
    // printing the console status output but also computing the run_results.
    let (sender, receiver) = channel::<Result<Vec<RunnerResult>>>();

    let mut status_output = status_output;

    crossbeam_utils::thread::scope(|s| {
        // Spawn the results collection in a crossbeam scope, so that it doesn't
        // take a slot in rayon's thread pool.
        s.spawn(|_| results_collection(&mut status_output, &mut run_results, test_count, receiver));

        // Rayon parallel iterator takes our vector and runs it on its thread
        // pool.
        test_groups
            .into_iter()
            .enumerate()
            .par_bridge()
            .try_for_each_with(sender, |sender, (i, (deqp, tests))| {
                sender.send(deqp.process_caselist(tests, i as u32))
            })
            .unwrap();

        // As we leave this scope, crossbeam will join the results collection
        // thread.  Note that it terminates cleanly because we moved the sender
        // into the rayon iterator.
    })
    .unwrap();

    Ok(run_results)
}

// Parses a deqp-runner regex set list.  We ignore empty lines and lines starting with "#", so you can
// leave notes in your skips/flakes lists about why.
pub fn parse_regex_set<I, S>(exprs: I) -> Result<RegexSet>
where
    S: AsRef<str>,
    I: IntoIterator<Item = S>,
{
    RegexSet::new(
        exprs
            .into_iter()
            .filter(|x| !x.as_ref().is_empty() && !x.as_ref().starts_with('#')),
    )
    .context("Parsing regex set")
}

pub fn read_lines<P, I: IntoIterator<Item = P>>(files: I) -> Result<Vec<String>>
where
    P: AsRef<Path>,
{
    let mut lines: Vec<String> = Vec::new();

    for path in files {
        let path = path.as_ref();
        for line in BufReader::new(
            File::open(&path).with_context(|| format!("opening path: {}", path.display()))?,
        )
        .lines()
        {
            let line = line.with_context(|| format!("reading line from {}", path.display()))?;
            // In newer dEQP, vk-master.txt just contains a list of .txt
            // caselist files relative to its current path, so recursively read
            // thoseand append their contents.
            if line.ends_with(".txt") {
                let sub_path = path.parent().context("Getting path parent dir")?.join(line);

                lines.extend_from_slice(
                    &read_lines(&[sub_path.as_path()])
                        .with_context(|| format!("reading sub-caselist {}", sub_path.display()))?,
                );
            } else {
                lines.push(line)
            }
        }
    }
    Ok(lines)
}

pub fn process_results(
    results: &RunnerResults,
    output_dir: &Path,
    summary_limit: usize,
) -> Result<()> {
    results.write_results(&mut File::create(&output_dir.join("results.csv"))?)?;
    results.write_failures(&mut File::create(&output_dir.join("failures.csv"))?)?;

    results.print_summary(if summary_limit == 0 {
        std::usize::MAX
    } else {
        summary_limit
    });

    if !results.is_success() {
        std::process::exit(1);
    }

    Ok(())
}

pub fn read_baseline(path: Option<&PathBuf>) -> Result<RunnerResults> {
    match path {
        Some(path) => {
            let mut file = File::open(path).context("Reading baseline")?;
            RunnerResults::from_csv(&mut file)
        }
        None => Ok(RunnerResults::new()),
    }
}