//! A module for the syntax nodes representing top-level WGSL declarations.

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

mod alias;
mod func;
pub mod import;
pub mod legacy;
mod struc;
mod var;

use crate::{common::AttributeList, ParseStream, Token, TokenKind};

pub use self::{
	alias::TypeAliasDecl,
	func::{FunctionDecl, ParamDecl},
	import::{
		ImportDecl, ImportPath, ImportPathBlock, ImportPathDecl, ImportPathLeaf,
		NamespacedImportPath,
	},
	struc::{FieldDecl, StructBody, StructDecl},
	var::VarDecl,
};

#[derive(Clone, DebugLisp)]
pub enum Decl {
	Var(VarDecl),
	Const(VarDecl),
	TypeAlias(TypeAliasDecl),
	Struct(StructDecl),
	Field(FieldDecl),
	Function(FunctionDecl),
	Param(ParamDecl),
	Extension(legacy::ExtensionDecl), // TODO
	ImportPath(ImportPathDecl),
	Import(ImportDecl),
	Module(legacy::ModuleDecl), // TODO
}

lazy_static! {
	// FIXME: This is gross -- is there a better way to handle this?
	static ref ERR_TOKEN: Token = Token::Ident(literal_substr!("ERROR"), Span::default());
}

impl Decl {
	pub fn name(&self) -> &Token {
		match self {
			Decl::Var(decl) | Decl::Const(decl) => &decl.name,
			Decl::TypeAlias(decl) => &decl.name,
			Decl::Struct(decl) => &decl.name,
			Decl::Field(decl) => &decl.name,
			Decl::Function(decl) => &decl.name,
			Decl::Param(decl) => &decl.name,
			Decl::Extension(decl) => &decl.name,
			Decl::ImportPath(decl) => decl.name(),
			Decl::Import(_) => &ERR_TOKEN,
			Decl::Module(decl) => &decl.name,
		}
	}

	pub fn attributes(&self) -> Option<&AttributeList> {
		match self {
			Decl::Var(decl) => decl.attributes.as_ref(),
			Decl::Const(decl) => decl.attributes.as_ref(),
			Decl::TypeAlias(_) => None,
			Decl::Struct(decl) => decl.attributes.as_ref(),
			Decl::Field(decl) => decl.attributes.as_ref(),
			Decl::Function(decl) => decl.attributes.as_ref(),
			Decl::Param(decl) => decl.attributes.as_ref(),
			Decl::Extension(_) => None,
			Decl::ImportPath(_) => None,
			Decl::Import(_) => None,
			Decl::Module(_) => None,
		}
	}
}

impl Parse for Decl {
	type Stream = ParseStream;

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

		match input.peek() {
			Some(token) => match token.as_matchable() {
				(Punct, "@", _) => {
					let attributes = input.parse::<AttributeList>()?;

					match input.parse::<Decl>()? {
						Decl::Var(mut inner) => {
							inner.attributes = Some(attributes);
							Ok(Decl::Var(inner))
						}
						Decl::Const(mut inner) => {
							inner.attributes = Some(attributes);
							Ok(Decl::Const(inner))
						}
						Decl::Struct(mut inner) => {
							inner.attributes = Some(attributes);
							Ok(Decl::Struct(inner))
						}
						Decl::Function(mut inner) => {
							inner.attributes = Some(attributes);
							Ok(Decl::Function(inner))
						}
						_ => Err(SpannedError {
							message: "Attributes are not valid in this position".into(),
							span: Some(attributes.span()),
							source: input.source(),
						}),
					}
				}
				(Keyword, "var", _) => Ok(Decl::Var(input.parse()?)),
				(Keyword, "let" | "const" | "override", _) => Ok(Decl::Const(input.parse()?)),
				(Keyword, "alias", _) => Ok(Decl::TypeAlias(input.parse()?)),
				(Keyword, "struct", _) => Ok(Decl::Struct(input.parse()?)),
				(Keyword, "fn", _) => Ok(Decl::Function(input.parse()?)),
				(Keyword, "enable", _) => Ok(Decl::Extension(input.parse()?)),
				(Directive, "#define_import_path", _) => Ok(Decl::ImportPath(input.parse()?)),
				(Directive, "#import", _) => Ok(Decl::Import(input.parse()?)),
				(Keyword, "import", _) => Ok(Decl::Module(input.parse()?)),
				(_, _, span) => Err(SpannedError {
					message: "Expected `var`, `let`, `const`, `override`, `alias`, \
						`struct`, `fn`, `enable`, or `import`"
						.into(),
					span: Some(span),
					source: input.source(),
				}),
			},
			None => Err(SpannedError {
				message: "Unexpected end of input".into(),
				source: input.source(),
				span: input.prev().map(|token| token.span()),
			}),
		}
	}
}

impl Spanned for Decl {
	fn span(&self) -> Span {
		use Decl::*;

		match self {
			Var(inner) => inner.span(),
			Const(inner) => inner.span(),
			TypeAlias(inner) => inner.span(),
			Struct(inner) => inner.span(),
			Field(inner) => inner.span(),
			Function(inner) => inner.span(),
			Param(inner) => inner.span(),
			Extension(inner) => inner.span(),
			ImportPath(inner) => inner.span(),
			Import(inner) => inner.span(),
			Module(inner) => inner.span(),
		}
	}
}