use anyhow::{Context, Result};
use deqp_runner::JunitGeneratorOptions;
use deqp_runner::{parallel_deqp, DeqpCommand, RunnerResults, TestCase};
use std::fs::{self, File};
use std::io::{self, BufRead, BufReader};
use std::path::{Path, PathBuf};
use std::time::Duration;
use structopt::StructOpt;

#[derive(Debug, StructOpt)]
#[structopt(
    author = "Eric Anholt <eric@anholt.net>",
    about = "Runs dEQP in parallel"
)]
struct Opts {
    #[structopt(subcommand)]
    subcmd: SubCommand,
}

#[derive(Debug, StructOpt)]
enum SubCommand {
    #[structopt(name = "run")]
    Run(Run),

    #[structopt(name = "junit")]
    Junit(Junit),
}

#[derive(Debug, StructOpt)]
pub struct Run {
    #[structopt(long, help = "path to deqp binary")]
    deqp: PathBuf,

    #[structopt(long, help = "path to deqp caselist (such as *-mustpass.txt)")]
    caselist: PathBuf,

    #[structopt(long = "output", help = "path to output directory")]
    output_dir: PathBuf,

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

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

    #[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)"
    )]
    flakes: Option<PathBuf>,

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

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

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

    #[structopt(long, default_value = "1", help = "Runs 1 out of every N tests.")]
    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)."
    )]
    fraction_start: usize,

    #[structopt(
        long,
        default_value = "500",
        help = "Starting number of tests to include in each deqp invocation (mitigates deqp startup overhead)"
    )]
    tests_per_group: usize,

    #[structopt(
        long,
        help = "Minimum number of tests to scale down to in each deqp invocation (defaults to tests_per_group)"
    )]
    min_tests_per_group: Option<usize>,

    #[structopt(
        long,
        help = "Optional path to executor/testlog-to-xml, for converting QPA files to usable XML"
    )]
    testlog_to_xml: Option<PathBuf>,

    #[structopt(help = "arguments to deqp binary")]
    deqp_args: Vec<String>,
}

#[derive(Debug, StructOpt)]
pub struct Junit {
    #[structopt(long, help = "Path to source results.csv or failures.csv")]
    results: PathBuf,

    #[structopt(long, short = "o", help = "Path to write junit XML to")]
    output: PathBuf,

    #[structopt(flatten)]
    junit_generator_options: JunitGeneratorOptions,
}

fn read_lines<P>(filename: P) -> Result<io::Lines<io::BufReader<File>>>
where
    P: AsRef<Path>,
{
    let file = File::open(&filename)
        .with_context(|| format!("opening filename: {}", filename.as_ref().display()))?;
    Ok(BufReader::new(file).lines())
}

fn main() -> Result<()> {
    stderrlog::new().module(module_path!()).init().unwrap();

    let opts = Opts::from_args();

    match opts.subcmd {
        SubCommand::Run(run) => {
            if run.jobs > 0 {
                rayon::ThreadPoolBuilder::new()
                    .num_threads(run.jobs)
                    .build_global()
                    .unwrap();
            }

            // Load all the test names from the caselist file
            if run.fraction < 1 {
                eprintln!("--fraction must be >= 1.");
                std::process::exit(1);
            }
            if run.fraction_start < 1 {
                eprintln!("--fraction_start must be >= 1.");
                std::process::exit(1);
            }
            let test_names: Vec<TestCase> = read_lines(&run.caselist)
                .context("reading caselist")?
                .map(|x| x.context("reading line"))
                .map(|x| TestCase::Deqp(x.unwrap()))
                .skip(run.fraction_start - 1)
                .step_by(run.fraction)
                .collect::<Vec<TestCase>>();

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

            let skips = match run.skips {
                Some(skips) => {
                    let lines = read_lines(skips)?
                        .map(|x| x.context("reading skips list"))
                        .collect::<Result<Vec<String>>>()?;
                    Some(deqp_runner::parse_regex_set(lines)?)
                }
                None => None,
            };

            let flakes = match run.flakes {
                Some(flakes) => {
                    let lines = read_lines(flakes)?
                        .map(|x| x.context("reading flakes list"))
                        .collect::<Result<Vec<String>>>()?;
                    Some(deqp_runner::parse_regex_set(lines)?)
                }
                None => None,
            };

            let baseline = match run.baseline {
                Some(path) => {
                    let mut file = File::open(path).context("Reading baseline")?;
                    RunnerResults::from_csv(&mut file)?
                }
                None => RunnerResults::new(),
            };

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

            let rayon_threads = rayon::current_num_threads();
            let tests_per_group = usize::max(
                1,
                usize::min(
                    (test_names.len() + rayon_threads - 1) / rayon_threads,
                    run.tests_per_group,
                ),
            );
            println!(
                "Running dEQP on {} threads in {}-test groups",
                rayon_threads, tests_per_group
            );

            let deqp = DeqpCommand {
                deqp: run.deqp,
                args: run.deqp_args,
                output_dir: run.output_dir.clone(),
                skips,
                flakes,
                baseline,
                timeout: Duration::from_secs_f32(run.timeout),
                qpa_to_xml: run.testlog_to_xml,
                tests_per_group,
                min_tests_per_group: if let Some(min_tests_per_group) = run.min_tests_per_group {
                    if min_tests_per_group == 0 {
                        eprintln!("--min_tests_per_group must be >= 1.");
                        std::process::exit(1);
                    }
                    min_tests_per_group
                } else {
                    // This runner has clever code I inherited from the parallel-deqp-runner
                    // project to, near the end of the list of test of tests to shard out, make
                    // smaller groups of tests so that you don't end up with a long test list
                    // left in one deqp process limiting your total runtime.
                    //
                    // The problem that cleverness faces is this table of deqp startup time:
                    //
                    //             freedreno  radeon
                    // deqp-gles2  0.2s       0.08s
                    // deqp-vk     2.0s       0.6s
                    //
                    // Even if all the slowest tests I have (a few deqp-vks on radeon at 6-7s)
                    // land in the same test list near the end of the run, you're still going to
                    // end up spawning a ton of extra deqp processes trying to be clever, costing
                    // more than your tests.
                    //
                    // So, disable it by default by using our normal group size the whole time
                    // unless someone sets the option to something else.
                    run.tests_per_group
                },
            };
            let results = parallel_deqp(&deqp, test_names)?;
            results.write_results(&mut File::create(&run.output_dir.join("results.csv"))?)?;
            results.write_failures(&mut File::create(&run.output_dir.join("failures.csv"))?)?;

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

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

        SubCommand::Junit(junit) => {
            let results = RunnerResults::from_csv(&mut File::open(&junit.results)?)
                .context("Reading in results csv")?;

            results.write_junit_failures(
                &mut File::create(&junit.output)?,
                &junit.junit_generator_options,
            )?;
        }
    }

    Ok(())
}
