use std::{collections::HashMap, convert::TryFrom};

use anyhow::Result;

use crate::shader_parser_ir3;

/// Parsed information about a shader from Mesa driver debug output.
#[derive(Debug, PartialEq, Clone)]
pub struct Shader {
    /// Hash that identifies the shader, which will be what's captured from the
    /// utrace events for the draw calls.
    pub sha1: String,
    /// Shader stage (arbitrary string, but useful for understanding results).
    pub stage: String,
    /// Original dumped shader code.
    pub code: String,
    /// Any useful information parsed out of the shader dump that might get used
    /// in a heuristic.
    pub stats: HashMap<String, u64>,
}

impl std::convert::TryFrom<&str> for Shader {
    type Error = anyhow::Error;

    fn try_from(code: &str) -> std::result::Result<Self, Self::Error> {
        // Introducing other shader parsers will probably involve a
        // .or_else(||) here.
        shader_parser_ir3::parse(code)
    }
}

/// Walks the stderr output, parsing shader code dumps as they're found.
pub fn find_shaders(output: &str) -> Result<Vec<Shader>> {
    let mut shaders = Vec::new();

    let mut code = String::new();
    for line in output.lines() {
        if !code.is_empty() && line == "MESA: info: " {
            shaders.push(Shader::try_from(code.as_str())?);
            code.clear();
        } else if !code.is_empty() || line.starts_with("MESA: info: Native code") {
            // New shader parsers will likely need to add to the starts_with()
            // check above for delimiters between shaders.
            code.push_str(line);
            code.push('\n');
        }
    }

    if !code.is_empty() {
        shaders.push(Shader::try_from(code.as_str())?);
    }

    Ok(shaders)
}

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

    #[test]
    fn test_full_output() {
        static STDERR: &str = include_str!("test_data/tu-shader-dump-stderr.txt");

        let shaders = find_shaders(STDERR).unwrap();

        let shas: Vec<_> = STDERR
            .lines()
            .filter(|x| x.contains("Native code for"))
            .map(|x| {
                let x = x.split(" with sha1 ").nth(1).unwrap().trim_end_matches(':');
                x.to_string()
            })
            .collect();

        for (i, shader) in shaders.iter().enumerate() {
            assert_eq!((i, &shader.sha1), (i, &shas[i]));

            if shader
                .code
                .lines()
                .filter(|x| x.contains("Native code for"))
                .count()
                != 1
            {
                println!("Parse fail, too many native codes for {}:", shader.sha1);
                println!("{}", shader.code);
            }
            println!("found {}", &shader.sha1);
        }

        assert_eq!(shaders.len(), 182);
    }

    static SINGLE_IR3_SHADER: &str = include_str!("test_data/single-ir3-shader.txt");

    #[test]
    fn test_single() {
        let shaders = find_shaders(SINGLE_IR3_SHADER).unwrap();
        assert_eq!(shaders.len(), 1);
        let shader = shaders.into_iter().next().unwrap();

        println!("stats:");
        for stat in &shader.stats {
            println!("{}: {}", stat.0, stat.1);
        }

        assert_eq!(shader.stats["full"], 6);
    }
}
