use std::borrow::Cow;
use std::iter;
use std::path::PathBuf;

use wesl::{
    syntax::PathOrigin, EscapeMangler, FileResolver, Mangler, ModulePath, ResolveError, Resolver,
    VirtualResolver, Wesl,
};

pub struct ShaderSources {
    virtual_resolver: VirtualResolver<'static>,
    file_resolver: Option<FileResolver>,
    mangler: EscapeMangler,
}

impl Default for ShaderSources {
    fn default() -> Self {
        Self::new(Some("wesl-lib".into()))
    }
}

impl ShaderSources {
    pub fn new(lib_dir: Option<PathBuf>) -> Self {
        Self {
            virtual_resolver: VirtualResolver::new(),
            file_resolver: lib_dir.map(FileResolver::new),
            mangler: Default::default(),
        }
    }

    pub fn set_module(&mut self, module: ModulePath, source: impl Into<String>) {
        self.virtual_resolver
            .add_module(module, Cow::Owned(source.into()));
    }

    pub fn compile(&mut self, root_module: &ModulePath) -> Result<String, wesl::Error> {
        Wesl::new_barebones()
            .set_custom_resolver(&*self)
            .set_custom_mangler(self.mangler.clone())
            .set_options(wesl::CompileOptions {
                mangle_root: true,
                ..Default::default()
            })
            .compile(root_module)
            .map(|compiled| compiled.to_string())
    }

    pub fn mangler(&self) -> &impl Mangler {
        &self.mangler
    }
}

impl Resolver for ShaderSources {
    fn resolve_source<'a>(&'a self, path: &ModulePath) -> Result<Cow<'a, str>, ResolveError> {
        match path.origin {
            PathOrigin::Absolute => self.virtual_resolver.resolve_source(path),
            PathOrigin::Package(ref pname) => {
                if let Some(ref files) = self.file_resolver {
                    let fpath = ModulePath {
                        origin: PathOrigin::Absolute,
                        components: iter::once(pname)
                            .chain(path.components.iter())
                            .cloned()
                            .collect(),
                    };
                    files.resolve_source(&fpath)
                } else {
                    Err(ResolveError::ModuleNotFound(
                        path.clone(),
                        "no library director found".into(),
                    ))
                }
            }
            _ => Err(ResolveError::ModuleNotFound(
                path.clone(),
                "can't resolve relative module".into(),
            )),
        }
    }
}