use build_print::{info, warn};
use naga::{
    ShaderStage,
    back::spv::{self, WriterFlags},
    front::glsl::{self, Options as GlslOptions, ParseErrors},
    valid::{ValidationError, ValidationFlags, Validator},
};
use std::{
    env, fs,
    path::{Path, PathBuf},
};
use thiserror::Error;

#[derive(Debug, Error)]
pub enum BuildError {
    #[error("I/O error: {0}")]
    Io(#[from] std::io::Error),
    #[error("environment variable error: {0}")]
    EnvVar(#[from] env::VarError),
    #[error("unsupported shader extension: {0}")]
    UnsupportedExt(String),
    #[error("GLSL parse errors in `{0}`:\n{1}")]
    ParseErrors(String, ParseErrors),
    #[error("validation errors in `{0}`: {1}")]
    ValidateErrors(String, ValidationError),
    #[error("SPIR-V write error for `{0}`: {1}")]
    Spv(String, spv::Error),
}

impl From<(String, ParseErrors)> for BuildError {
    fn from((s, e): (String, ParseErrors)) -> Self {
        BuildError::ParseErrors(s, e)
    }
}

impl From<(String, ValidationError)> for BuildError {
    fn from((s, e): (String, ValidationError)) -> Self {
        BuildError::ValidateErrors(s, e)
    }
}

impl From<(String, spv::Error)> for BuildError {
    fn from((s, e): (String, spv::Error)) -> Self {
        BuildError::Spv(s, e)
    }
}

fn compile_shader(path: &Path, out_dir: &Path) -> Result<(), BuildError> {
    let ext = path
        .extension()
        .and_then(|e| e.to_str())
        .map(str::to_string)
        .ok_or_else(|| BuildError::UnsupportedExt(path.display().to_string()))?;
    let stage = match ext.as_str() {
        "vert" => ShaderStage::Vertex,
        "frag" => ShaderStage::Fragment,
        "comp" => ShaderStage::Compute,
        _ => return Err(BuildError::UnsupportedExt(ext)),
    };
    let src = fs::read_to_string(path)?;

    let module = glsl::Frontend::default()
        .parse(&GlslOptions::from(stage), &src)
        .map_err(|e| (ext.clone(), e))?;
    let info = Validator::new(ValidationFlags::all(), Default::default())
        .validate(&module)
        .map_err(|e| (ext.clone(), e.into_inner()))?;
    let mut writer = spv::Writer::new(&spv::Options {
        flags: WriterFlags::empty(),
        ..Default::default()
    })
    .map_err(|e| (ext.clone(), e))?;
    let mut spirv = Vec::new();
    writer
        .write(&module, &info, None, &None, &mut spirv)
        .map_err(|e| (ext.clone(), e))?;
    let out_path = out_dir.join(format!("{}.spv", ext));
    fs::write(&out_path, bytemuck::cast_slice(&spirv))?;
    info!("Compiled {} → {}", path.display(), out_path.display());
    Ok(())
}

fn main() -> Result<(), BuildError> {
    println!("cargo:rerun-if-changed=shaders");
    let out_dir = PathBuf::from(env::var("OUT_DIR")?);
    for entry in fs::read_dir("shaders")? {
        let path = entry?.path();
        if !path.is_file() {
            continue;
        }
        if let Err(e) = compile_shader(&path, &out_dir) {
            if matches!(e, BuildError::UnsupportedExt(_)) {
                warn!("{}", e);
                continue;
            }
            return Err(e);
        }
    }
    Ok(())
}