xbasic 0.3.2

use crate::error_handler::ErrorHandler;
use crate::expr::{Expr, ExprValue};
use crate::function::{Function, FunctionDefinition};
use crate::stmt::Stmt;
use crate::tokens::{Token, TokenType};

pub(crate) struct Parser<'a> {
	tokens: Vec<Token>,
	functions: Vec<FunctionDefinition>,
	current: usize,
	error_handler: &'a mut ErrorHandler,
}

impl<'a> Parser<'a> {
	pub(crate) fn new(
		tokens: Vec<Token>,
		functions: Vec<FunctionDefinition>,
		error_handler: &'a mut ErrorHandler,
	) -> Self {
		Self {
			tokens,
			functions,
			current: 0,
			error_handler,
		}
	}

	pub(crate) fn parse(&mut self) -> Option<(Vec<Stmt>, Vec<Function>)> {
		let mut v: Vec<Stmt> = Vec::new();
		let mut f: Vec<Function> = Vec::new();
		while !self.is_at_end() {
			while self.match_type(&TokenType::Newline).is_some() {}

			if self.match_type(&TokenType::Function).is_some() {
				f.push(self.function()?);
			} else {
				match self.statement() {
					Some(stmt) => v.push(stmt),
					None => {
						self.synchronize(); // TODO this is useless
						return None;
					}
				}
			}
		}
		Some((v, f))
	}

	fn function(&mut self) -> Option<Function> {
		// Read in identifier
		let name = self.consume(TokenType::Identifier, "Expected identifier after FUNCTION.")?;
		// Read in (
		self.consume(TokenType::LeftParen, "Expected '(' after function name.")?;

		// Get parameters
		let mut params = Vec::new();
		if self.peek().token_type != TokenType::RightParen {
			let mut another = true;
			while !self.is_at_end() && another {
				// Should never fail since we already checked this in the resolver
				params.push(self.consume(TokenType::Identifier, "Expected parameter name.")?);
				another = self.match_type(&TokenType::Comma).is_some();
			}
		}

		// Read in )
		self.consume(
			TokenType::RightParen,
			"Expected ')' after function parameters.",
		)?;

		// Newline
		self.consume(
			TokenType::Newline,
			"Expected newline after function parameters.",
		);

		// Keep going until we read in "end function"
		let mut statements = Vec::new();
		while !self.is_at_end()
			&& (self.peek().token_type != TokenType::End
				|| self.peek_next() != &TokenType::Function)
		{
			statements.push(self.statement()?)
		}

		self.consume(
			TokenType::End,
			"Expected END FUNCTION after function declaration.",
		)?;

		self.consume(
			TokenType::Function,
			"Expected END FUNCTION after function declaration.",
		)?;
		self.consume(TokenType::Newline, "Expected newline after END FUNCTION")?;

		Some(Function::new(name.lexeme, params, statements))
	}

	fn statement(&mut self) -> Option<Stmt> {
		if self.match_type(&TokenType::Print).is_some() {
			return self.print_statement();
		}

		if self.match_type(&TokenType::Input).is_some() {
			return self.input_statement();
		}

		if self.match_type(&TokenType::Return).is_some() {
			return self.return_statement();
		}

		if self.match_type(&TokenType::If).is_some() {
			return self.if_statement("IF");
		}

		if self.match_type(&TokenType::While).is_some() {
			return self.while_statement();
		}

		if self.match_type(&TokenType::For).is_some() {
			return self.for_statement();
		}

		if self.peek_next() == &TokenType::Equal {
			let identifier = self.advance();

			// Consume equal
			let equals = self.advance();

			let val = self.expression()?;

			// Consume newline
			self.advance();

			return Some(Stmt::new_assign(identifier, val, equals.line));
		}

		if [
			TokenType::PlusEqual,
			TokenType::MinusEqual,
			TokenType::StarEqual,
			TokenType::SlashEqual,
			TokenType::CaretEqual,
			TokenType::PercentEqual,
		]
		.contains(self.peek_next())
		{
			let identifier = self.advance();
			let operator = self.advance();
			let val = self.expression()?;

			self.consume(TokenType::Newline, "Expected newline after assignment.")?;

			return Some(Stmt::new_binary_assign(
				identifier,
				operator.token_type,
				val,
				operator.line,
			));
		}

		self.expression_statement()
	}

	fn print_statement(&mut self) -> Option<Stmt> {
		let mut values: Vec<Expr> = Vec::new();
		while !self.is_at_end() && self.peek().token_type != TokenType::Newline {
			values.push(self.expression()?);
		}

		let line = self.previous().line;
		self.consume(TokenType::Newline, "Expected end of line after values.")?;

		if values.len() > 255 {
			self.error_handler.error(
				line,
				"Cannot pass more than 255 expressions to a PRINT statement.",
			);
			return None;
		}

		Some(Stmt::new_print(values, line))
	}

	fn input_statement(&mut self) -> Option<Stmt> {
		let variable = self.consume(TokenType::Identifier, "Expected identifier after INPUT.")?;
		self.consume(
			TokenType::Newline,
			"Expected newline after INPUT statement.",
		);
		Some(Stmt::new_input(variable))
	}

	fn return_statement(&mut self) -> Option<Stmt> {
		let return_value = if self.match_type(&TokenType::Newline).is_none() {
			let return_value = self.expression()?;
			self.consume(
				TokenType::Newline,
				"Expected newline after RETURN statement.",
			)?;
			return_value
		} else {
			// TODO this should be nil if we add nil
			Expr::new_literal(ExprValue::Integer(0))
		};
		Some(Stmt::new_return(return_value, self.previous().line))
	}

	fn if_statement(&mut self, type_str: &str) -> Option<Stmt> {
		let line = self.previous().line;
		let condition = self.expression()?;

		self.consume(
			TokenType::Then,
			&format!("Expected THEN after {} statement condition.", type_str),
		)?;

		// TODO single line if statements
		let _single_line = self.match_type(&TokenType::Newline).is_none();

		let mut then_stmts = Vec::new();
		while !self.is_at_end_if() && !self.is_at_else() && !self.is_at_end() {
			then_stmts.push(self.statement()?);
		}

		let mut elseif = false;
		let else_stmts = if self.match_type(&TokenType::ElseIf).is_some() {
			//self.consume(TokenType::Newline, "Expected newline after elseif.");
			elseif = true;
			vec![self.if_statement("ELSEIF")?]
		} else if self.match_type(&TokenType::Else).is_some() {
			self.consume(TokenType::Newline, "Expected newline after else.");
			let mut else_stmts = Vec::new();
			while !self.is_at_end_if() && !self.is_at_end() {
				else_stmts.push(self.statement()?);
			}
			else_stmts
		} else {
			Vec::new()
		};

		if !elseif {
			self.consume(TokenType::End, "Expected END IF after IF statement.")?;
			self.consume(TokenType::If, "Expected END IF after IF statement.")?;
			self.consume(TokenType::Newline, "Expected END IF after IF statement.")?;
		}

		Some(Stmt::new_if(condition, then_stmts, else_stmts, line))
	}

	fn while_statement(&mut self) -> Option<Stmt> {
		let line = self.peek().line;
		let condition = self.expression()?;

		self.consume(
			TokenType::Newline,
			"Expected newline after while condition.",
		)?;

		let mut body: Vec<Stmt> = Vec::new();
		while self.peek().token_type != TokenType::Wend && !self.is_at_end() {
			body.push(self.statement()?);
		}

		self.consume(TokenType::Wend, "Expected WEND after while body.")?;
		let _ = self.match_type(&TokenType::Newline);

		Some(Stmt::new_while(condition, body, line))
	}

	fn for_statement(&mut self) -> Option<Stmt> {
		let line = self.peek().line;
		let variable = self.consume(TokenType::Identifier, "Expected variable name.")?;
		self.consume(TokenType::Equal, "Expected '='.")?;
		let min_value = self.expression()?;
		self.consume(TokenType::To, "Expected TO.")?;
		let max_value = self.expression()?;
		self.consume(TokenType::Newline, "Expected newline.")?;
		let mut stmts = Vec::new();
		while self.peek().token_type != TokenType::Next && !self.is_at_end() {
			stmts.push(self.statement()?);
		}

		self.consume(TokenType::Next, "Expected NEXT after FOR body.")?;

		let var = self.consume(TokenType::Identifier, "Expected identifier after NEXT.")?;
		if var.lexeme != variable.lexeme {
			self.error_handler
				.error_token(&var, "Incorrect variable after NEXT.");
			return None;
		}

		let _ = self.match_type(&TokenType::Newline);

		Some(Stmt::new_for(variable, min_value, max_value, stmts, line))
	}

	fn expression_statement(&mut self) -> Option<Stmt> {
		let expr = self.expression()?;
		let line = self.previous().line;
		self.consume(
			TokenType::Newline,
			"Expected newline after expression statement.",
		)?;

		Some(Stmt::new_expression(expr, line))
	}

	fn expression(&mut self) -> Option<Expr> {
		self.or()
	}

	fn or(&mut self) -> Option<Expr> {
		let mut left = self.and()?;
		while let Some(_operator) = self.match_type(&TokenType::Or) {
			let right = self.and()?;
			left = Expr::new_logical(left, TokenType::Or, right);
		}

		Some(left)
	}

	fn and(&mut self) -> Option<Expr> {
		let mut left = self.equality()?;
		while let Some(_operator) = self.match_type(&TokenType::And) {
			let right = self.and()?;
			left = Expr::new_logical(left, TokenType::And, right);
		}

		Some(left)
	}

	fn equality(&mut self) -> Option<Expr> {
		let mut expr = self.comparison()?;

		while let Some(op) = self.match_type(&TokenType::Equal) {
			match self.comparison() {
				Some(right) => expr = Expr::new_binary(expr, op.token_type.clone(), right),
				None => return None,
			}
		}

		Some(expr)
	}

	fn comparison(&mut self) -> Option<Expr> {
		let mut expr = self.addition()?;

		while let Some(operator) = self.match_types(&[
			TokenType::Greater,
			TokenType::GreaterEqual,
			TokenType::Less,
			TokenType::LessEqual,
		]) {
			match self.addition() {
				Some(right) => expr = Expr::new_binary(expr, operator.token_type.clone(), right),
				None => return None,
			}
		}

		Some(expr)
	}

	fn addition(&mut self) -> Option<Expr> {
		let mut expr = self.multiplication()?;

		while let Some(operator) = self.match_types(&[TokenType::Minus, TokenType::Plus]) {
			match self.multiplication() {
				Some(right) => expr = Expr::new_binary(expr, operator.token_type.clone(), right),
				None => return None,
			}
		}

		Some(expr)
	}

	fn multiplication(&mut self) -> Option<Expr> {
		let mut expr = self.exponentiation()?;

		while let Some(operator) =
			self.match_types(&[TokenType::Star, TokenType::Slash, TokenType::Percent])
		{
			match self.exponentiation() {
				Some(right) => expr = Expr::new_binary(expr, operator.token_type.clone(), right),
				None => return None,
			}
		}

		Some(expr)
	}

	fn exponentiation(&mut self) -> Option<Expr> {
		let mut expr = self.unary()?;
		if let Some(operator) = self.match_type(&TokenType::Caret) {
			match self.exponentiation() {
				Some(right) => expr = Expr::new_binary(expr, operator.token_type, right),
				None => return None,
			}
		}

		Some(expr)
	}

	fn unary(&mut self) -> Option<Expr> {
		if let Some(operator) = self.match_types(&[TokenType::Not, TokenType::Minus]) {
			return match self.unary() {
				Some(right) => Some(Expr::new_unary(operator.token_type, right)),
				None => None,
			};
		}

		self.primary()
	}

	fn primary(&mut self) -> Option<Expr> {
		Some(match &self.advance().token_type {
			TokenType::False => Expr::new_literal(ExprValue::Boolean(false)),
			TokenType::True => Expr::new_literal(ExprValue::Boolean(true)),

			TokenType::Integer(x) => Expr::new_literal(ExprValue::Integer(*x)),
			TokenType::Decimal(x) => Expr::new_literal(ExprValue::Decimal(*x)),
			TokenType::String(x) => Expr::new_literal(ExprValue::String(x.to_string())),

			TokenType::LeftParen => {
				if self.is_at_end() {
					self.error_handler
						.error_token(&self.previous(), "Expected ')' after expression.");
					return None;
				}
				match self.expression() {
					Some(expr) => {
						self.consume(TokenType::RightParen, "Expected ')' after expression.")?;
						expr
					}
					None => return None,
				}
			}
			TokenType::Identifier => {
				// Could be a variable or a function
				let name = self.previous();
				if self.match_type(&TokenType::LeftParen).is_some() {
					self.function_call(name)?
				} else {
					Expr::new_variable(name)
				}
			}
			_ => {
				self.error_handler
					.error_token(&self.peek(), "Expected expression.");
				return None;
			}
		})
	}

	fn function_call(&mut self, identifier: Token) -> Option<Expr> {
		// Verify function exists
		let mut definition = None;
		for def in self.functions.clone() {
			if def.name == identifier.lexeme {
				definition = Some(def);
				break;
			}
		}

		let definition = match definition {
			Some(x) => x,
			None => {
				self.error_handler
					.error_token(&identifier, "Not a function.");
				return None;
			}
		};

		// Collect args
		let mut args = Vec::new();
		if self.peek().token_type != TokenType::RightParen {
			let mut again = true;
			while !self.is_at_end() && again {
				args.push(self.expression()?);
				again = self.match_type(&TokenType::Comma).is_some();
			}
		}

		// Consume )
		self.consume(
			TokenType::RightParen,
			"Expected ')' after function call arguments.",
		)?;

		if args.len() > 255 {
			self.error_handler.error_token(
				&identifier,
				"Cannot call function with more than 255 arguments.",
			);
			return None;
		}

		// Check arity
		if (definition.arity) != args.len() as u8 {
			self.error_handler.error_token(
				&identifier,
				&format!(
					"Expected {} arguments, got {}.",
					definition.arity,
					args.len()
				),
			);
			return None;
		}

		// Build function call expression
		Some(Expr::new_call(identifier.lexeme, args))
	}

	fn synchronize(&mut self) {
		self.advance();
		while !self.is_at_end() {
			// TODO more here?
			if self.peek().token_type == TokenType::Newline {
				// jump over newline
				self.advance();
				return;
			}

			self.advance();
		}
	}

	fn consume(&mut self, token_type: TokenType, message: &str) -> Option<Token> {
		if self.check(&token_type) {
			return Some(self.advance());
		}

		self.error_handler.error_token(&self.peek(), message);
		None
	}

	fn match_type(&mut self, token_type: &TokenType) -> Option<Token> {
		if self.check(token_type) {
			return Some(self.advance());
		}

		None
	}

	fn match_types(&mut self, types: &[TokenType]) -> Option<Token> {
		for t in types {
			if let Some(x) = self.match_type(t) {
				return Some(x);
			}
		}

		None
	}

	fn is_at_end(&self) -> bool {
		self.peek().token_type == TokenType::Eof
	}

	fn is_at_end_if(&self) -> bool {
		self.peek().token_type == TokenType::End && self.peek_next() == &TokenType::If
	}

	fn is_at_else(&self) -> bool {
		self.peek().token_type == TokenType::Else || self.peek().token_type == TokenType::ElseIf
	}

	fn check(&self, token_type: &TokenType) -> bool {
		if self.is_at_end() {
			return false;
		}
		&self.peek().token_type == token_type
	}

	fn advance(&mut self) -> Token {
		let v = self.peek();
		if !self.is_at_end() {
			self.current += 1;
		}
		v
	}

	fn peek(&self) -> Token {
		self.tokens[self.current].clone()
	}

	fn peek_next(&self) -> &TokenType {
		if self.current + 1 == self.tokens.len() {
			&TokenType::NoToken
		} else {
			&self.tokens[self.current + 1].token_type
		}
	}

	fn previous(&self) -> Token {
		self.tokens[self.current - 1].clone()
	}
}