//! Helpers for working with [`naga-oil`] modules.
//!
//! [`naga-oil`]: https://github.com/bevyengine/naga_oil

use std::sync::Arc;

use gramatika::{Parse, ParseStreamer, Span, Spanned, SpannedError, Token as _};

use crate::{
	parser::ErrorRecoveringParseStream,
	token::{brace, directive, ident, punct, Token, TokenKind},
	ParseStream,
};

/// Represents a `#define_import_path` statement like the ones supported by
/// [`naga_oil`](https://github.com/bevyengine/naga_oil):
///
/// ```wgsl
/// #define_import_path my::shader::module
/// ```
#[derive(Clone, DebugLisp)]
pub struct ImportPathDecl {
	pub keyword: Token,
	pub path: ImportPath,
}

/// Represents an `#import` statement like the ones supported by
/// [`naga_oil`](https://github.com/bevyengine/naga_oil):
///
/// ```wgsl
/// #import my::shader::{
///     module_a,
///     module_b::{TypeA, TypeB},
///     module_c::function,
///     module_d::function as d_function,
/// }
/// ```
#[derive(Clone, DebugLisp)]
pub struct ImportDecl {
	pub keyword: Token,
	pub path: ImportPath,
}

#[derive(Clone, DebugLisp)]
pub enum ImportPath {
	Namespaced(NamespacedImportPath),
	Block(ImportPathBlock),
	Leaf(ImportPathLeaf),
}

#[derive(Clone, DebugLisp)]
pub struct NamespacedImportPath {
	pub namespace: Token,
	pub path: Arc<ImportPath>,
}

#[derive(Clone, DebugLisp)]
pub struct ImportPathBlock {
	pub brace_open: Token,
	pub paths: Arc<[ImportPath]>,
	pub brace_close: Token,
}

#[derive(Clone, DebugLisp)]
pub struct ImportPathLeaf {
	pub name: Token,
	pub as_binding: Option<Token>,
}

impl ImportPathDecl {
	/// Returns the token at the leaf node of this declaration.
	///
	/// I.e., for the following import path declaration:
	/// ```wgsl
	/// #define_import_path foo::bar::baz
	/// ```
	/// This function will return the `bar` token at the end of the line.
	pub fn name(&self) -> &Token {
		let mut import_path = &self.path;
		loop {
			match import_path {
				ImportPath::Namespaced(NamespacedImportPath { path, .. }) => {
					import_path = path.as_ref();
				}
				ImportPath::Leaf(ImportPathLeaf { name, .. }) => {
					break name;
				}
				ImportPath::Block(_) => {
					unreachable!();
				}
			}
		}
	}
}

impl Parse for ImportPathDecl {
	type Stream = ParseStream;

	fn parse(input: &mut Self::Stream) -> gramatika::Result<Self> {
		let keyword = input.consume(directive!["#define_import_path"])?;
		let path = input.parse()?;

		// Return an error if any portion of the path is a block or has an `as` binding
		let mut parsed_path = &path;
		loop {
			match parsed_path {
				ImportPath::Namespaced(NamespacedImportPath { path, .. }) => {
					parsed_path = path.as_ref();
				}
				ImportPath::Block(block) => {
					return Err(SpannedError {
						message: "Path blocks are not valid in `#define_import_path`".into(),
						source: input.source(),
						span: Some(block.span()),
					});
				}
				ImportPath::Leaf(leaf) => {
					if let Some(as_binding) = leaf.as_binding.as_ref() {
						return Err(SpannedError {
							message: "`as` bindings are not valid in `#define_import_path`".into(),
							source: input.source(),
							span: Some(as_binding.span()),
						});
					} else {
						break;
					}
				}
			}
		}

		Ok(Self { keyword, path })
	}
}

impl Spanned for ImportPathDecl {
	fn span(&self) -> Span {
		self.keyword.span().through(self.path.span())
	}
}

impl Parse for ImportDecl {
	type Stream = ParseStream;

	fn parse(input: &mut Self::Stream) -> gramatika::Result<Self> {
		let keyword = input.consume(directive!["#import"])?;
		let path = input.parse()?;

		Ok(Self { keyword, path })
	}
}

impl Spanned for ImportDecl {
	fn span(&self) -> Span {
		self.keyword.span().through(self.path.span())
	}
}

impl Parse for ImportPath {
	type Stream = ParseStream;

	fn parse(input: &mut Self::Stream) -> gramatika::Result<Self> {
		use TokenKind::*;

		match input.next() {
			Some(mut next) => match next.as_matchable() {
				(Ident | Path, _, _) => {
					next = match next.kind() {
						TokenKind::Ident => input.upgrade_last(TokenKind::Ident, Token::module)?,
						TokenKind::Path => next,
						_ => unreachable!(),
					};

					if input.check(punct![::]) {
						input.discard();

						Ok(ImportPath::Namespaced(NamespacedImportPath {
							namespace: next,
							path: Arc::new(input.parse()?),
						}))
					} else {
						let name = next;
						let as_binding = if input.check(ident![as]) {
							let _ = input.consume_as(TokenKind::Ident, Token::keyword)?;
							Some(input.consume_as(TokenKind::Ident, Token::module)?)
						} else {
							None
						};

						Ok(ImportPath::Leaf(ImportPathLeaf { name, as_binding }))
					}
				}
				(Brace, "{", _) => {
					let brace_open = next;

					let paths =
						input.parse_seq_separated(punct![,], |input| !input.check(brace!("}")))?;

					let brace_close = input.consume(brace!("}"))?;

					Ok(ImportPath::Block(ImportPathBlock {
						brace_open,
						paths: paths.into(),
						brace_close,
					}))
				}
				(_, _, span) => Err(SpannedError {
					message: "Expected identifier or `{`".into(),
					source: input.source(),
					span: Some(span),
				}),
			},
			None => Err(SpannedError {
				message: "Unexpected end of input".into(),
				source: input.source(),
				span: input.prev().map(|token| token.span()),
			}),
		}
	}
}

impl Spanned for ImportPath {
	fn span(&self) -> Span {
		match self {
			ImportPath::Namespaced(NamespacedImportPath { namespace, path }) => {
				namespace.span().through(path.span())
			}
			ImportPath::Block(inner) => inner.span(),
			ImportPath::Leaf(inner) => inner.span(),
		}
	}
}

impl Spanned for ImportPathBlock {
	fn span(&self) -> Span {
		self.brace_open.span().through(self.brace_close.span())
	}
}

impl Spanned for ImportPathLeaf {
	fn span(&self) -> Span {
		if let Some(binding) = self.as_binding.as_ref() {
			self.name.span().through(binding.span())
		} else {
			self.name.span()
		}
	}
}