//! A module for the syntax nodes representing WGSL expressions.

use std::{fmt, sync::Arc};

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

use crate::{
	common::{ArgumentList, TypeDecl},
	token::{brace, keyword, operator, punct},
	ParseStream, Token, TokenKind,
};

#[derive(Clone, DebugLisp)]
pub enum Expr {
	UnaryPre(UnaryPreExpr),
	UnaryPost(UnaryPostExpr),
	Binary(BinaryExpr),
	Assignment(BinaryExpr),
	FnCall(FnCallExpr),
	TypeCtor(TypeCtorExpr),
	Group(GroupExpr),
	Bitcast(BitcastExpr),
	Literal(Token),
	Ident(IdentExpr),
	Primary(PrimaryExpr),
}

#[derive(Clone, DebugLisp)]
pub struct UnaryPreExpr {
	pub op: Token,
	pub expr: Arc<Expr>,
}

#[derive(Clone, DebugLisp)]
pub struct UnaryPostExpr {
	pub expr: Arc<Expr>,
	pub op: Token,
}

#[derive(Clone, DebugLisp)]
pub struct BinaryExpr {
	pub lhs: Arc<Expr>,
	pub op: Token,
	pub rhs: Arc<Expr>,
}

#[derive(Clone, DebugLisp)]
pub struct FnCallExpr {
	pub ident: IdentExpr,
	pub arguments: ArgumentList,
}

#[derive(Clone, DebugLisp)]
pub struct TypeCtorExpr {
	pub ty: TypeDecl,
	pub arguments: ArgumentList,
}

#[derive(Clone, DebugLisp)]
pub struct GroupExpr {
	pub brace_open: Token,
	pub expr: Arc<Expr>,
	pub brace_close: Token,
}

#[derive(Clone, DebugLisp)]
pub struct BitcastExpr {
	pub keyword: Token,
	pub ty: TypeDecl,
	pub expr: GroupExpr,
}

#[derive(Clone, DebugLisp)]
pub enum IdentExpr {
	Namespaced(NamespacedIdent),
	Leaf(Token),
}

#[derive(Clone, DebugLisp)]
pub struct NamespacedIdent {
	pub namespace: Token,
	pub ident: Arc<IdentExpr>,
}

#[derive(Clone, DebugLisp)]
pub enum IdentExprBuilder {
	Namespaced(NamespacedIdentBuilder),
	Leaf(Token),
}

#[derive(Clone, DebugLisp)]
pub struct NamespacedIdentBuilder {
	pub namespace: Token,
	pub ident: Box<IdentExprBuilder>,
}

#[derive(Clone, DebugLisp)]
pub struct PrimaryExpr {
	pub expr: Arc<Expr>,
	pub postfix: Option<PostfixExpr>,
}

struct PrimaryExprBuilder {
	pub expr: Box<Expr>,
	pub postfix: Option<PostfixExpr>,
}

#[derive(Clone, DebugLisp)]
pub struct PostfixExpr {
	pub accessor: Accessor,
	pub expr: Arc<Expr>,
	pub postfix: Option<Arc<PostfixExpr>>,
}

#[derive(Clone)]
pub enum Accessor {
	Dot(Token),
	Index([Token; 2]),
}

impl gramatika::DebugLisp for Accessor {
	fn fmt(&self, f: &mut fmt::Formatter<'_>, indent: usize) -> fmt::Result {
		match self {
			Accessor::Dot(token) => write!(f, "(Accessor::Dot {:?})", token),
			Accessor::Index(tokens) => {
				writeln!(f, "(Accessor::Index")?;
				for token in tokens {
					write!(f, "{}", debug::INDENT.repeat(indent + 1))?;
					writeln!(f, "{:?},", token)?;
				}
				write!(f, "{})", debug::INDENT.repeat(indent))
			}
		}
	}
}

// -- Recursive descent trait ------------------------------------------------------------

trait RecursiveDescent {
	type Token: gramatika::Token;

	fn assignment(&mut self) -> Result<Expr>;
	fn short_circuit_or(&mut self) -> Result<Expr>;
	fn short_circuit_and(&mut self) -> Result<Expr>;
	fn inclusive_or(&mut self) -> Result<Expr>;
	fn exclusive_or(&mut self) -> Result<Expr>;
	fn and(&mut self) -> Result<Expr>;
	fn equality(&mut self) -> Result<Expr>;
	fn relational(&mut self) -> Result<Expr>;
	fn shift(&mut self) -> Result<Expr>;
	fn additive(&mut self) -> Result<Expr>;
	fn multiplicative(&mut self) -> Result<Expr>;
	fn unary(&mut self) -> Result<Expr>;
	fn binary(
		&mut self,
		operators: &[Self::Token],
		operand_method: fn(&mut Self) -> Result<Expr>,
	) -> Result<Expr>;
}

// --- gramatika impls -------------------------------------------------------------------

impl Spanned for Expr {
	fn span(&self) -> Span {
		match self {
			Expr::UnaryPre(inner) => inner.span(),
			Expr::UnaryPost(inner) => inner.span(),
			Expr::Binary(inner) => inner.span(),
			Expr::Assignment(inner) => inner.span(),
			Expr::FnCall(inner) => inner.span(),
			Expr::TypeCtor(inner) => inner.span(),
			Expr::Group(inner) => inner.span(),
			Expr::Bitcast(inner) => inner.span(),
			Expr::Literal(inner) => inner.span(),
			Expr::Ident(inner) => inner.span(),
			Expr::Primary(inner) => inner.span(),
		}
	}
}

impl Spanned for PrimaryExpr {
	fn span(&self) -> Span {
		match self.postfix {
			Some(ref postfix) => self.expr.span().through(postfix.span()),
			None => self.expr.span(),
		}
	}
}

impl Spanned for PostfixExpr {
	fn span(&self) -> Span {
		match self.postfix {
			Some(ref postfix) => self.expr.span().through(postfix.span()),
			None => self.expr.span(),
		}
	}
}

impl Spanned for IdentExpr {
	fn span(&self) -> Span {
		match self {
			IdentExpr::Leaf(name) => name.span(),
			IdentExpr::Namespaced(NamespacedIdent { namespace, ident }) => {
				namespace.span().through(ident.span())
			}
		}
	}
}

impl Spanned for FnCallExpr {
	fn span(&self) -> Span {
		self.ident.span().through(self.arguments.brace_close.span())
	}
}

impl Spanned for BinaryExpr {
	fn span(&self) -> Span {
		self.lhs.span().through(self.rhs.span())
	}
}

impl Spanned for UnaryPreExpr {
	fn span(&self) -> Span {
		self.op.span().through(self.expr.span())
	}
}

impl Spanned for UnaryPostExpr {
	fn span(&self) -> Span {
		self.expr.span().through(self.op.span())
	}
}

impl Spanned for TypeCtorExpr {
	fn span(&self) -> Span {
		self.ty.span().through(self.arguments.brace_close.span())
	}
}

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

impl Spanned for BitcastExpr {
	fn span(&self) -> Span {
		self.keyword.span().through(self.expr.brace_close.span())
	}
}

impl Parse for Expr {
	type Stream = ParseStream;

	fn parse(input: &mut Self::Stream) -> Result<Self> {
		input.assignment()
	}
}

impl Parse for PrimaryExprBuilder {
	type Stream = ParseStream;

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

		let expr = match input.peek() {
			Some(token) => match token.as_matchable() {
				(Type, _, _) => Ok(Expr::TypeCtor(input.parse::<TypeCtorExpr>()?)),
				(Brace, "(", _) => Ok(Expr::Group(input.parse::<GroupExpr>()?)),
				(Keyword, "bitcast", _) => Ok(Expr::Bitcast(input.parse::<BitcastExpr>()?)),
				(IntLiteral, _, _) | (FloatLiteral, _, _) | (Keyword, "true" | "false", _) => {
					let token = input.next().unwrap();
					Ok(Expr::Literal(token))
				}
				(Ident, _, _) => {
					let mut ident = input.parse::<IdentExprBuilder>()?;

					if input.check(brace!["("]) {
						let mut expr = &mut ident;
						while let IdentExprBuilder::Namespaced(NamespacedIdentBuilder {
							ident,
							..
						}) = expr
						{
							expr = ident.as_mut();
						}
						let IdentExprBuilder::Leaf(name) = expr else {
							unreachable!();
						};
						*name = input.upgrade_last(TokenKind::Ident, Token::function)?;
						let arguments = input.parse::<ArgumentList>()?;

						Ok(Expr::FnCall(FnCallExpr {
							ident: ident.build(),
							arguments,
						}))
					} else {
						Ok(Expr::Ident(ident.build()))
					}
				}
				(_, _, span) => Err(SpannedError {
					message: "Expected expression".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()),
			}),
		}?;

		let postfix = match input.peek().map(|t| t.as_matchable()) {
			Some((Punct, ".", _) | (Brace, "[", _)) => Some(input.parse::<PostfixExpr>()?),
			_ => None,
		};

		Ok(PrimaryExprBuilder {
			expr: Box::new(expr),
			postfix,
		})
	}
}

impl PrimaryExprBuilder {
	fn build(self) -> PrimaryExpr {
		PrimaryExpr {
			expr: self.expr.into(),
			postfix: self.postfix,
		}
	}
}

impl Parse for FnCallExpr {
	type Stream = ParseStream;

	fn parse(input: &mut Self::Stream) -> Result<Self> {
		let mut ident = input.parse::<IdentExprBuilder>()?;

		// TODO: Is this needed? We're already upgrading in `PrimaryExprBuilder::Parse`
		let mut expr = &mut ident;
		while let IdentExprBuilder::Namespaced(NamespacedIdentBuilder { ident, .. }) = expr {
			expr = ident.as_mut();
		}
		let IdentExprBuilder::Leaf(name) = expr else {
			unreachable!();
		};
		*name = input.upgrade_last(TokenKind::Ident, Token::function)?;

		let arguments = input.parse::<ArgumentList>()?;

		Ok(Self {
			ident: ident.build(),
			arguments,
		})
	}
}

impl Parse for TypeCtorExpr {
	type Stream = ParseStream;

	fn parse(input: &mut Self::Stream) -> Result<Self> {
		let ty = input.parse::<TypeDecl>()?;
		let arguments = input.parse::<ArgumentList>()?;

		Ok(Self { ty, arguments })
	}
}

impl Parse for GroupExpr {
	type Stream = ParseStream;

	fn parse(input: &mut Self::Stream) -> Result<Self> {
		let brace_open = input.consume(brace!["("])?;
		let expr = input.parse::<Expr>()?;
		let brace_close = input.consume(brace![")"])?;

		Ok(Self {
			brace_open,
			expr: Arc::new(expr),
			brace_close,
		})
	}
}

impl Parse for BitcastExpr {
	type Stream = ParseStream;

	fn parse(input: &mut Self::Stream) -> Result<Self> {
		let keyword = input.consume(keyword![bitcast])?;
		input.consume(operator![<])?;
		let ty = input.parse::<TypeDecl>()?;
		input.consume(operator![>])?;
		let expr = input.parse::<GroupExpr>()?;

		Ok(Self { keyword, ty, expr })
	}
}

impl Parse for IdentExprBuilder {
	type Stream = ParseStream;

	fn parse(input: &mut Self::Stream) -> Result<Self> {
		let mut ident = input.consume_kind(TokenKind::Ident)?;

		if input.check(punct![::]) {
			ident = input.upgrade_last(TokenKind::Ident, Token::module)?;
			input.consume(punct![::])?;

			Ok(Self::Namespaced(NamespacedIdentBuilder {
				namespace: ident,
				ident: Box::new(input.parse()?),
			}))
		} else {
			Ok(Self::Leaf(ident))
		}
	}
}

impl IdentExprBuilder {
	pub fn build(&self) -> IdentExpr {
		match self {
			IdentExprBuilder::Leaf(name) => IdentExpr::Leaf(name.clone()),
			IdentExprBuilder::Namespaced(NamespacedIdentBuilder { namespace, ident }) => {
				IdentExpr::Namespaced(NamespacedIdent {
					namespace: namespace.clone(),
					ident: Arc::new(ident.build()),
				})
			}
		}
	}
}

impl From<&IdentExpr> for IdentExprBuilder {
	fn from(value: &IdentExpr) -> Self {
		match value {
			IdentExpr::Leaf(name) => IdentExprBuilder::Leaf(name.clone()),
			IdentExpr::Namespaced(NamespacedIdent { namespace, ident }) => {
				IdentExprBuilder::Namespaced(NamespacedIdentBuilder {
					namespace: namespace.clone(),
					ident: Box::new(ident.as_ref().into()),
				})
			}
		}
	}
}

impl Parse for IdentExpr {
	type Stream = ParseStream;

	fn parse(input: &mut Self::Stream) -> Result<Self> {
		Ok(input.parse::<IdentExprBuilder>()?.build())
	}
}

impl Parse for PostfixExpr {
	type Stream = ParseStream;

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

		if input.check(punct![.]) {
			let dot = input.next().unwrap();
			let accessor = Accessor::Dot(dot);

			let mut expr = input.parse::<PrimaryExprBuilder>()?;
			if let Expr::Ident(ident) = expr.expr.as_mut() {
				if let IdentExpr::Leaf(name) = ident {
					*name = input.upgrade(name.clone(), Token::field)?;
				} else {
					// FIXME: This is pretty wildly inefficient
					let mut builder = IdentExprBuilder::from(&*ident);
					let mut ident_expr = &mut builder;
					while let IdentExprBuilder::Namespaced(NamespacedIdentBuilder {
						ident, ..
					}) = ident_expr
					{
						ident_expr = ident.as_mut();
					}
					let IdentExprBuilder::Leaf(name) = ident_expr else {
						unreachable!();
					};
					*name = input.upgrade(name.clone(), Token::field)?;
					*ident = builder.build();
				}
			}

			let postfix = match input.peek().map(|t| t.as_matchable()) {
				Some((Punct, ".", _) | (Brace, "[", _)) => {
					Some(Arc::new(input.parse::<PostfixExpr>()?))
				}
				_ => None,
			};

			Ok(Self {
				accessor,
				expr: Arc::new(Expr::Primary(expr.build())),
				postfix,
			})
		} else {
			let brace_open = input.consume(brace!["["])?;
			let expr = input.parse::<Expr>()?;
			let brace_close = input.consume(brace!["]"])?;

			let postfix = match input.peek().map(|t| t.as_matchable()) {
				Some((Punct, ".", _) | (Brace, "[", _)) => {
					Some(Arc::new(input.parse::<PostfixExpr>()?))
				}
				_ => None,
			};

			Ok(Self {
				accessor: Accessor::Index([brace_open, brace_close]),
				expr: Arc::new(expr),
				postfix,
			})
		}
	}
}

// --- RecursiveDescent impl -------------------------------------------------------------

impl RecursiveDescent for ParseStream {
	type Token = Token;

	fn assignment(&mut self) -> Result<Expr> {
		let lhs = self.short_circuit_or()?;

		// if self.check(operator![=]) {
		if matches!(
			self.peek().map(|t| t.as_matchable()),
			Some((
				TokenKind::Operator,
				"+=" | "-=" | "*=" | "/=" | "%=" | "&=" | "|=" | "^=" | "<<=" | ">>=" | "=",
				_
			))
		) {
			let op = self.next().unwrap();
			let value = self.assignment()?;

			match &lhs {
				Expr::Primary(PrimaryExpr { expr, .. })
					if matches!(expr.as_ref(), Expr::Ident(_)) =>
				{
					Ok(Expr::Assignment(BinaryExpr {
						lhs: Arc::new(lhs),
						op,
						rhs: Arc::new(value),
					}))
				}
				other => Err(SpannedError {
					message: "Invalid assignment target".into(),
					source: self.source(),
					span: Some(other.span()),
				}),
			}
		} else {
			Ok(lhs)
		}
	}

	fn short_circuit_or(&mut self) -> Result<Expr> {
		self.binary(&[operator![||]], Self::short_circuit_and)
	}

	fn short_circuit_and(&mut self) -> Result<Expr> {
		self.binary(&[operator![&&]], Self::inclusive_or)
	}

	fn inclusive_or(&mut self) -> Result<Expr> {
		self.binary(&[operator![|]], Self::exclusive_or)
	}

	fn exclusive_or(&mut self) -> Result<Expr> {
		self.binary(&[operator![^]], Self::and)
	}

	fn and(&mut self) -> Result<Expr> {
		self.binary(&[operator![&]], Self::equality)
	}

	fn equality(&mut self) -> Result<Expr> {
		self.binary(&[operator![==], operator![!=]], Self::relational)
	}

	fn relational(&mut self) -> Result<Expr> {
		self.binary(
			&[operator![<=], operator![>=], operator![<], operator![>]],
			Self::shift,
		)
	}

	fn shift(&mut self) -> Result<Expr> {
		self.binary(&[operator![<<], operator![>>]], Self::additive)
	}

	fn additive(&mut self) -> Result<Expr> {
		self.binary(&[operator!["+"], operator!["-"]], Self::multiplicative)
	}

	fn multiplicative(&mut self) -> Result<Expr> {
		self.binary(&[operator![*], operator![/], operator![%]], Self::unary)
	}

	fn unary(&mut self) -> Result<Expr> {
		if matches!(
			self.peek().map(|t| t.as_matchable()),
			Some((TokenKind::Operator, "-" | "!" | "~" | "*" | "&", _))
		) {
			let op = self.next().unwrap();
			let expr = self.unary()?;

			Ok(Expr::UnaryPre(UnaryPreExpr {
				op,
				expr: Arc::new(expr),
			}))
		} else {
			let primary = Expr::Primary(self.parse::<PrimaryExprBuilder>()?.build());

			if let Some((TokenKind::Operator, "++" | "--", _)) =
				self.peek().map(|t| t.as_matchable())
			{
				let op = self.next().unwrap();

				Ok(Expr::UnaryPost(UnaryPostExpr {
					expr: Arc::new(primary),
					op,
				}))
			} else {
				Ok(primary)
			}
		}
	}

	fn binary(
		&mut self,
		operators: &[Self::Token],
		operand_method: fn(&mut Self) -> Result<Expr>,
	) -> Result<Expr> {
		let mut expr = operand_method(self)?;

		while operators.iter().any(|op| self.check(op.clone())) {
			let op = self.next().unwrap();
			let rhs = operand_method(self)?;

			expr = Expr::Binary(BinaryExpr {
				lhs: Arc::new(expr),
				op,
				rhs: Arc::new(rhs),
			});
		}

		Ok(expr)
	}
}