error.rs

use std::{borrow::Cow, collections::HashMap, ops::Range};

use codespan_reporting::{
    diagnostic::{Diagnostic, Label},
    files::SimpleFile,
    term,
    term::termcolor::WriteColor,
};
use thiserror::Error;
use tracing::trace;

use super::{preprocess::PreprocessOutput, Composer, ShaderDefValue};
use crate::{compose::SPAN_SHIFT, redirect::RedirectError};

#[derive(Debug)]
pub enum ErrSource {
    Module {
        name: String,
        offset: usize,
        defs: HashMap<String, ShaderDefValue>,
    },
    Constructing {
        path: String,
        source: String,
        offset: usize,
    },
}

impl ErrSource {
    pub fn path<'a>(&'a self, composer: &'a Composer) -> &'a String {
        match self {
            ErrSource::Module { name, .. } => &composer.module_sets.get(name).unwrap().file_path,
            ErrSource::Constructing { path, .. } => path,
        }
    }

    pub fn source<'a>(&'a self, composer: &'a Composer) -> Cow<'a, String> {
        match self {
            ErrSource::Module { name, defs, .. } => {
                let raw_source = &composer.module_sets.get(name).unwrap().sanitized_source;
                let Ok(PreprocessOutput {
                    preprocessed_source: source,
                    ..
                }) = composer.preprocessor.preprocess(raw_source, defs)
                else {
                    return Default::default();
                };

                Cow::Owned(source)
            }
            ErrSource::Constructing { source, .. } => Cow::Borrowed(source),
        }
    }

    pub fn offset(&self) -> usize {
        match self {
            ErrSource::Module { offset, .. } | ErrSource::Constructing { offset, .. } => *offset,
        }
    }
}

#[derive(Debug, Error)]
#[error("Composer error: {inner}")]
pub struct ComposerError {
    #[source]
    pub inner: ComposerErrorInner,
    pub source: ErrSource,
}

#[derive(Debug, Error)]
pub enum ComposerErrorInner {
    #[error("{0}")]
    ImportParseError(String, usize),
    #[error("required import '{0}' not found")]
    ImportNotFound(String, usize),
    #[error("{0}")]
    WgslParseError(naga::front::wgsl::ParseError),
    #[cfg(feature = "glsl")]
    #[error("{0:?}")]
    GlslParseError(naga::front::glsl::ParseErrors),
    #[error("naga_oil bug, please file a report: failed to convert imported module IR back into WGSL for use with WGSL shaders: {0}")]
    WgslBackError(naga::back::wgsl::Error),
    #[cfg(feature = "glsl")]
    #[error("naga_oil bug, please file a report: failed to convert imported module IR back into GLSL for use with GLSL shaders: {0}")]
    GlslBackError(naga::back::glsl::Error),
    #[error("naga_oil bug, please file a report: composer failed to build a valid header: {0}")]
    HeaderValidationError(naga::WithSpan<naga::valid::ValidationError>),
    #[error("failed to build a valid final module: {0}")]
    ShaderValidationError(naga::WithSpan<naga::valid::ValidationError>),
    #[error(
        "Not enough '# endif' lines. Each if statement should be followed by an endif statement."
    )]
    NotEnoughEndIfs(usize),
    #[error("Too many '# endif' lines. Each endif should be preceded by an if statement.")]
    TooManyEndIfs(usize),
    #[error("'#else' without preceding condition.")]
    ElseWithoutCondition(usize),
    #[error("Unknown shader def operator: '{operator}'")]
    UnknownShaderDefOperator { pos: usize, operator: String },
    #[error("Unknown shader def: '{shader_def_name}'")]
    UnknownShaderDef { pos: usize, shader_def_name: String },
    #[error(
        "Invalid shader def comparison for '{shader_def_name}': expected {expected}, got {value}"
    )]
    InvalidShaderDefComparisonValue {
        pos: usize,
        shader_def_name: String,
        expected: String,
        value: String,
    },
    #[error("multiple inconsistent shader def values: '{def}'")]
    InconsistentShaderDefValue { def: String },
    #[error("Attempted to add a module with no #define_import_path")]
    NoModuleName,
    #[error("source contains internal decoration string, results probably won't be what you expect. if you have a legitimate reason to do this please file a report")]
    DecorationInSource(Range<usize>),
    #[error("naga oil only supports glsl 440 and 450")]
    GlslInvalidVersion(usize),
    #[error("invalid override :{0}")]
    RedirectError(#[from] RedirectError),
    #[error(
        "override is invalid as `{name}` is not virtual (this error can be disabled with feature 'override_any')"
    )]
    OverrideNotVirtual { name: String, pos: usize },
    #[error(
        "Composable module identifiers must not require substitution according to naga writeback rules: `{original}`"
    )]
    InvalidIdentifier { original: String, at: naga::Span },
    #[error("Invalid value for `#define`d shader def {name}: {value}")]
    InvalidShaderDefDefinitionValue {
        name: String,
        value: String,
        pos: usize,
    },
    #[error("#define statements are only allowed at the start of the top-level shaders")]
    DefineInModule(usize),
    #[error("Invalid WGSL directive '{directive}' at line {position}: {reason}")]
    InvalidWgslDirective {
        directive: String,
        position: usize,
        reason: String,
    },
}

struct ErrorSources<'a> {
    current: Option<&'a (dyn std::error::Error + 'static)>,
}

impl<'a> ErrorSources<'a> {
    fn of(error: &'a dyn std::error::Error) -> Self {
        Self {
            current: error.source(),
        }
    }
}

impl<'a> Iterator for ErrorSources<'a> {
    type Item = &'a (dyn std::error::Error + 'static);

    fn next(&mut self) -> Option<Self::Item> {
        let current = self.current;
        self.current = self.current.and_then(std::error::Error::source);
        current
    }
}

// impl<'a> FusedIterator for ErrorSources<'a> {}

impl ComposerError {
    /// format a Composer error
    pub fn emit_to_string(&self, composer: &Composer) -> String {
        composer.undecorate(&self.emit_to_string_internal(composer))
    }

    fn emit_to_string_internal(&self, composer: &Composer) -> String {
        let path = self.source.path(composer);
        let source = self.source.source(composer);
        let source_offset = self.source.offset();

        trace!("source:\n~{}~", source);
        trace!("source offset: {}", source_offset);

        let map_span = |rng: Range<usize>| -> Range<usize> {
            ((rng.start & ((1 << SPAN_SHIFT) - 1)).saturating_sub(source_offset))
                ..((rng.end & ((1 << SPAN_SHIFT) - 1)).saturating_sub(source_offset))
        };

        let files = SimpleFile::new(path, source.as_str());
        let config = term::Config::default();
        let (labels, notes) = match &self.inner {
            ComposerErrorInner::DecorationInSource(range) => {
                (vec![Label::primary((), range.clone())], vec![])
            }
            ComposerErrorInner::HeaderValidationError(v)
            | ComposerErrorInner::ShaderValidationError(v) => (
                v.spans()
                    .map(|(span, desc)| {
                        trace!(
                            "mapping span {:?} -> {:?}",
                            span.to_range().unwrap_or(0..0),
                            map_span(span.to_range().unwrap_or(0..0))
                        );
                        Label::primary((), map_span(span.to_range().unwrap_or(0..0)))
                            .with_message(desc.to_owned())
                    })
                    .collect(),
                ErrorSources::of(&v)
                    .map(|source| source.to_string())
                    .collect(),
            ),
            ComposerErrorInner::ImportNotFound(msg, pos) => (
                vec![Label::primary((), *pos..*pos)],
                vec![format!("missing import '{msg}'")],
            ),
            ComposerErrorInner::ImportParseError(msg, pos) => (
                vec![Label::primary((), *pos..*pos)],
                vec![format!("invalid import spec: '{msg}'")],
            ),
            ComposerErrorInner::WgslParseError(e) => (
                e.labels()
                    .map(|(range, msg)| {
                        Label::primary((), map_span(range.to_range().unwrap_or(0..0)))
                            .with_message(msg)
                    })
                    .collect(),
                vec![e.message().to_owned()],
            ),
            #[cfg(feature = "glsl")]
            ComposerErrorInner::GlslParseError(e) => (
                e.errors
                    .iter()
                    .map(|naga::front::glsl::Error { kind, meta }| {
                        Label::primary((), map_span(meta.to_range().unwrap_or(0..0)))
                            .with_message(kind.to_string())
                    })
                    .collect(),
                vec![],
            ),
            ComposerErrorInner::NotEnoughEndIfs(pos)
            | ComposerErrorInner::TooManyEndIfs(pos)
            | ComposerErrorInner::ElseWithoutCondition(pos)
            | ComposerErrorInner::UnknownShaderDef { pos, .. }
            | ComposerErrorInner::UnknownShaderDefOperator { pos, .. }
            | ComposerErrorInner::InvalidShaderDefComparisonValue { pos, .. }
            | ComposerErrorInner::OverrideNotVirtual { pos, .. }
            | ComposerErrorInner::GlslInvalidVersion(pos)
            | ComposerErrorInner::DefineInModule(pos)
            | ComposerErrorInner::InvalidShaderDefDefinitionValue { pos, .. }
            | ComposerErrorInner::InvalidWgslDirective { position: pos, .. } => {
                (vec![Label::primary((), *pos..*pos)], vec![])
            }
            ComposerErrorInner::WgslBackError(e) => {
                return format!("{path}: wgsl back error: {e}");
            }
            #[cfg(feature = "glsl")]
            ComposerErrorInner::GlslBackError(e) => {
                return format!("{path}: glsl back error: {e}");
            }
            ComposerErrorInner::InconsistentShaderDefValue { def } => {
                return format!("{path}: multiple inconsistent shader def values: '{def}'");
            }
            ComposerErrorInner::RedirectError(..) => (
                vec![Label::primary((), 0..0)],
                vec![format!("override error")],
            ),
            ComposerErrorInner::NoModuleName => {
                return format!(
                    "{path}: no #define_import_path declaration found in composable module"
                );
            }
            ComposerErrorInner::InvalidIdentifier { at, .. } => (
                vec![Label::primary((), map_span(at.to_range().unwrap_or(0..0)))
                    .with_message(self.inner.to_string())],
                vec![],
            ),
        };

        let diagnostic = Diagnostic::error()
            .with_message(self.inner.to_string())
            .with_labels(labels)
            .with_notes(notes);

        let mut msg = Vec::with_capacity(256);

        let mut color_writer;
        let mut no_color_writer;
        let writer: &mut dyn WriteColor = if supports_color() {
            color_writer = term::termcolor::Ansi::new(&mut msg);
            &mut color_writer
        } else {
            no_color_writer = term::termcolor::NoColor::new(&mut msg);
            &mut no_color_writer
        };

        term::emit(writer, &config, &files, &diagnostic).expect("cannot write error");

        String::from_utf8_lossy(&msg).into_owned()
    }
}

#[cfg(any(test, target_arch = "wasm32"))]
fn supports_color() -> bool {
    false
}

// termcolor doesn't expose this logic when using custom buffers
#[cfg(not(any(test, target_arch = "wasm32")))]
fn supports_color() -> bool {
    // TODO: https://github.com/tokio-rs/tracing/issues/3378
    false
}