gpu-trace-perf 1.6.0

use std::{path::Path, process::Command};

use anyhow::{Context, Result, bail};
use log::{debug, error, warn};
use regex::Regex;
use serde::Deserialize;

use crate::{Replay, ReplayOutput, replay_command, snapshot::SnapshotResult};

pub struct ApitraceTrace {
    file: String,
    is_directx: bool,
}

impl ApitraceTrace {
    pub fn new(file: &str) -> ApitraceTrace {
        ApitraceTrace {
            file: file.to_owned(),
            is_directx: apitrace_file_is_directx(file).unwrap_or_else(|e| {
                error!("Failure calling apitrace info, assuming file is GL: {}", e);
                true
            }),
        }
    }
}

pub fn call_apitrace<R: Replay>(trace: &R, command: Command) -> Result<ReplayOutput> {
    let output = trace.run_replay_command(command);

    if !output.status.success() {
        if output.stderr.contains("waffle_context_create failed") {
            warn!(
                "apitrace reported waffle_context_create() failed, likely due to trace requiring too new of a GL version"
            );
        }
        bail!("Failed to start apitrace");
    }

    Ok(output)
}

fn parse_snapshot_line(line: &str) -> Option<String> {
    if line.starts_with("Wrote ") {
        Some(line.trim_start_matches("Wrote ").to_string())
    } else {
        None
    }
}

impl Replay for ApitraceTrace {
    fn replay(&self, wrapper: Option<&str>, envs: &[(String, String)]) -> Result<ReplayOutput> {
        // apitrace replay otherwise assumes x11 (glretrace) for glx traces instead of using waffle
        let apitrace_command = [
            "eglretrace",
            "--pgpu",
            "--headless",
            "--loop=1", // loop the last frame once so we know that caches are hot, shaders are compiled, etc.
            &self.file,
        ];

        call_apitrace(self, replay_command(&apitrace_command, wrapper, envs))
    }

    fn fps(&self, output: &ReplayOutput) -> Result<f64> {
        parse_apitrace_pgpu_output(&output.stdout)
    }

    fn name(&self) -> &str {
        &self.file
    }

    fn can_snapshot(&self) -> bool {
        true
    }

    fn snapshot(&self, output_dir: &str, loops: u32) -> Result<SnapshotResult> {
        let last_frame = apitrace_last_frame(&self.file)?;

        let output_dir = self.output_dir(output_dir)?.unwrap();
        let mut args = Vec::new();
        if self.is_directx {
            args.push("wine".to_string());
            args.push("d3dretrace.exe".to_string());
        } else {
            args.push("apitrace".to_string());
            args.push("replay".to_string());
        }
        args.push("--headless".to_string());

        if loops > 1 {
            args.push(format!("--loop={}", loops - 1));
            args.push("--call-nos=false".to_string());
        } else {
            args.push(format!("--snapshot={last_frame}"));
        }
        args.push(format!(
            "--snapshot-prefix={}/snapshot",
            output_dir.display()
        ));
        args.push(self.file.to_string());

        let start_time = std::time::Instant::now();
        let command = replay_command(&args, None, &[]);

        let cmdline = format!("{:?}", &command);

        let result = call_apitrace(self, command)?;

        let mut files = Vec::new();
        for line in result.stdout.lines() {
            if let Some(path) = parse_snapshot_line(line) {
                files.push(
                    Path::new(&path)
                        .strip_prefix(&output_dir)
                        .context("getting output dir relative to snapshot")?
                        .to_path_buf(),
                );
            }
        }

        // If we're taking looped snapshots, then just capture the loop count of
        // frames (assuming a single-frame trace capture).
        let loops = loops as usize;
        let files = if loops > 1 {
            if files.len() >= loops {
                let first_save = files.len() - loops;

                // Delete .pngs for the frames we're not saving.
                for file in &files[0..first_save] {
                    let full_path = output_dir.join(file);
                    std::fs::remove_file(&full_path).unwrap_or_else(|e| {
                        error!("Removing unneeded snapshot {}: {e}", full_path.display())
                    });
                }

                files[first_save..files.len()].to_vec()
            } else {
                warn!(
                    "Frame looping on {} didn't appear to generate per-frame snapshots, do you have https://github.com/apitrace/apitrace/pull/968",
                    self.name()
                );
                files
            }
        } else {
            files
        };

        Ok(SnapshotResult {
            files,
            cmdline,
            stdout: result.stdout,
            stderr: result.stderr,
            runtime: start_time.elapsed(),
        })
    }
}

// Returns the fps from the last frame of an apitrace replay --pgpu output
fn parse_apitrace_pgpu_output(output: &str) -> Result<f64> {
    lazy_static! {
        static ref CALL_RE: Regex = Regex::new("^call [0-9]+ -?[0-9]+ ([0-9]+)").unwrap();
    }

    let mut total = 0;
    let mut start_of_frame = true;
    for line in output.lines() {
        if line == "frame_end" {
            start_of_frame = true;
        } else {
            let cap = CALL_RE.captures(line);
            if let Some(cap) = cap {
                if start_of_frame {
                    total = 0;
                    start_of_frame = false;
                }
                match cap[1].parse::<i64>() {
                    Ok(gpu) => {
                        if gpu >= 0 {
                            total += gpu;
                        } else {
                            anyhow::bail!(
                                "apitrace produced GL_TIME_ELAPSED < 0, skipping(gpu hang?)"
                            );
                        }
                    }
                    Err(_) => {
                        anyhow::bail!("failed to parse apitrace's GL_TIME_ELAPSED");
                    }
                }
            }
        }
    }

    if total == 0 {
        anyhow::bail!("No times parsed");
    }

    Ok(1_000_000_000.0 / (total as f64))
}

/// If we're measuring times with u_trace, we don't need timestamps, but we do
/// need to loop the frame because there's a bunch of setup at the start that
/// we're not trying to measure.
pub struct ApitraceUtraceTrace {
    file: String,
    is_directx: bool,
}

impl ApitraceUtraceTrace {
    pub fn new(file: &str) -> ApitraceUtraceTrace {
        ApitraceUtraceTrace {
            file: file.to_string(),
            is_directx: apitrace_file_is_directx(file).unwrap_or_else(|e| {
                error!("Failure calling apitrace info, assuming file is GL: {}", e);
                true
            }),
        }
    }
}

impl Replay for ApitraceUtraceTrace {
    fn replay(&self, wrapper: Option<&str>, envs: &[(String, String)]) -> Result<ReplayOutput> {
        // apitrace replay otherwise assumes x11 (glretrace) for glx traces
        // instead of using waffle
        //
        // loops the last frame twice so we know that caches are hot, shaders
        // are compiled, etc., and the middle frame utrace results will be one
        // of the hot ones.
        let apitrace_command = if self.is_directx {
            vec![
                "wine",
                "d3dretrace.exe",
                "--headless",
                "--loop=2",
                &self.file,
            ]
        } else {
            vec!["eglretrace", "--headless", "--loop=2", &self.file]
        };

        call_apitrace(self, replay_command(&apitrace_command, wrapper, envs))
    }

    fn fps(&self, _: &ReplayOutput) -> Result<f64> {
        unreachable!("shouldn't be called");
    }

    fn name(&self) -> &str {
        &self.file
    }
}

#[derive(Deserialize)]
struct ApitraceInfo {
    #[serde(rename = "API")]
    api: String,
}

pub fn apitrace_info_is_directx(input: &str) -> Result<bool> {
    let info = serde_json::from_str::<ApitraceInfo>(input)
        .with_context(|| format!("Parsing apitrace info output:\n{input}"))?;
    Ok(info.api == "DirectX")
}

pub fn apitrace_file_is_directx_native(filename: &str) -> Result<bool> {
    let file = std::fs::File::open(filename).with_context(|| format!("opening {filename}"))?;
    let reader = apitrace::TraceReader::new(file)
        .with_context(|| format!("opening {filename} as apitrace trace"))?;
    match reader.guess_api()? {
        apitrace::call_flags::CallAPI::GL => Ok(false),
        apitrace::call_flags::CallAPI::D3D => Ok(true),
        apitrace::call_flags::CallAPI::Unknown => bail!("Failed to detect API for {filename}"),
    }
}

pub fn apitrace_file_is_directx(file: &str) -> Result<bool> {
    debug!("checking directx on {}", file);

    match apitrace_file_is_directx_native(file) {
        Ok(dx) => return Ok(dx),
        Err(e) => warn!("Failed to check apitrace file for directx-ness: {e}"),
    }

    let mut command = replay_command(&["apitrace", "info", file], None, &[]);
    let output = command.output().context("Calling apitrae info")?;
    apitrace_info_is_directx(&String::from_utf8_lossy(&output.stdout))
}

pub fn apitrace_last_call_number(input: &str) -> Result<u64> {
    let mut last_line = None;

    for line in input.lines() {
        if !line.trim_end().is_empty() {
            last_line = Some(line);
        }
    }

    let last_line = last_line.context("finding a non-empty line in apitrace dump output")?;
    let result = str::parse::<u64>(last_line.split_once(' ').context("finding space")?.0)
        .context("parsing u64 from last line")?;
    Ok(result)
}

pub fn apitrace_last_frame(file: &str) -> Result<u64> {
    // TODO: Convert this over to using apitrace-rs (once we flag all the frame-end call sigs)
    debug!("apitrace dumping {}", file);
    let args = ["apitrace", "dump", "--calls=frame", file];
    let mut command = Command::new(args[0]);
    for arg in &args[1..] {
        command.arg(arg);
    }

    let output = ReplayOutput::from(command.output().context("calling apitrace dump")?);

    apitrace_last_call_number(&output.stdout)
        .with_context(|| format!("Getting last frame in {}", file))
}

#[cfg(test)]
mod tests {
    use super::*;
    use assert_approx_eq::assert_approx_eq;

    #[test]
    fn test_apitrace_parsing() {
        // Actual apitrace output, trimmed down for the testcase.
        let apitrace_input = "
# call no gpu_start gpu_dura cpu_start cpu_dura vsize_start vsize_dura rss_start rss_dura pixels program name
call 44 0 0 0 0 0 0 0 0 0 0 glViewport
call 56 25082334 50166 0 0 0 0 0 0 0 0 glClear
call 81 41719667 0 0 0 0 0 0 0 0 0 glClear
call 176 42206667 472583 0 0 0 0 0 0 0 7 glDrawArrays
frame_end
call 222 0 0 0 0 0 0 0 0 0 4 glClearColor
call 224 45001334 21666 0 0 0 0 0 0 0 4 glClear
call 231 45023750 38000 0 0 0 0 0 0 0 7 glClear
call 239 45062584 519333 0 0 0 0 0 0 0 7 glDrawArrays
frame_end
call 222 0 0 0 0 0 0 0 0 0 4 glClearColor
call 224 47438000 13666 0 0 0 0 0 0 0 4 glClear
call 231 47452417 59583 0 0 0 0 0 0 0 7 glClear
call 239 47512917 579083 0 0 0 0 0 0 0 7 glDrawArrays
frame_end
Rendered 3 frames in 0.0539452 secs, average of 55.612 fps
";
        assert_approx_eq!(
            parse_apitrace_pgpu_output(apitrace_input).unwrap(),
            1.0 / ((13_666 + 59_583 + 579_083) as f64 / 1_000_000_000.0)
        )
    }

    #[test]
    fn test_apitrace_parsing_negatve_start() {
        let apitrace_input = "call 318 -8883437858 156 0 0 0 0 0 0 0 0 glBlitFramebuffer";
        assert_approx_eq!(
            parse_apitrace_pgpu_output(apitrace_input).unwrap(),
            1.0 / (156.0 / 1_000_000_000.0)
        );
    }

    #[test]
    fn test_apitrace_parsing_empty() {
        let apitrace_input = "
# call no gpu_start gpu_dura cpu_start cpu_dura vsize_start vsize_dura rss_start rss_dura pixels program name
call 44 0 0 0 0 0 0 0 0 0 0 glViewport
frame_end
";
        assert!(parse_apitrace_pgpu_output(apitrace_input).is_err());
    }

    #[test]
    fn test_apitrace_directx_info() -> Result<()> {
        assert!(apitrace_info_is_directx(
            r#"
{
  "FileName": "/home/anholt/src/traces-db/unigine/heaven-scene1-low-d3d11.trace-dxgi",
  "ContainerVersion": 6,
  "ContainerType": "Brotli",
  "API": "DirectX",
  "FramesCount": 104,
  "ActualDataSize": 130120914,
  "ContainerSize": 63387386
}
"#
        )?);
        assert!(!apitrace_info_is_directx(
            r#"
{
  "FileName": "/home/anholt/src/traces-db/neverball/neverball-v2.trace",
  "ContainerVersion": 6,
  "ContainerType": "Brotli",
  "API": "OpenGL + GLX/WGL/CGL",
  "FramesCount": 147,
  "ActualDataSize": 21503984,
  "ContainerSize": 1554696
}
"#
        )?);
        Ok(())
    }

    #[test]
    fn test_apitrace_last_call_number() {
        let input = r#"
// process.name = "/usr/bin/glxgears"
1384 glXSwapBuffers(dpy = 0x56060e921f80, drawable = 31457282)

1413 glXSwapBuffers(dpy = 0x56060e921f80, drawable = 31457282)
"#;

        assert_eq!(apitrace_last_call_number(input).unwrap(), 1413);
    }

    #[test]
    fn test_parse_snapshot_line() {
        assert_eq!(
            parse_snapshot_line("Wrote /path/to/snapshot0001.png"),
            Some("/path/to/snapshot0001.png".to_string())
        );
        assert_eq!(
            parse_snapshot_line("Wrote output/dir/frame-123.png"),
            Some("output/dir/frame-123.png".to_string())
        );
        assert_eq!(parse_snapshot_line("Some other output"), None);
        assert_eq!(parse_snapshot_line(""), None);
    }
}