########################################################
##	IMPORTS
########################################################
from strings_with_arrows import *
import string
import os
import math
import constants
from error import *
from position import *
from lexer import *
from nodes import *


########################################################
##	BUILDING GRAMMAR TREE
########################################################



########################################################
##	PARSE RESULT
########################################################

class ParseResult:
	def __init__(self):
		self.error = None
		self.node = None
		self.advance_count = 0 ##to check if the error needs to be overwritten, how many times we advacned
		self.last_registered_advance_count = 0
		self.advance_count = 0
		self.to_reverse_count = 0

	def register_advancements(self):
	    self.last_registered_advance_count = 1
	    self.advance_count += 1

	def register(self,res):##take in  a parse result or a node 
	    self.last_registered_advance_count = res.advance_count
	    self.advance_count += res.advance_count
	    if res.error: self.error = res.error
	    return res.node ##extract the node and return it 

	def try_register(self, res):
		if res.error:
		  self.to_reverse_count = res.advance_count
		  return None
		return self.register(res)

	def success(self, node):
		self.node = node
		return self

	def failure(self,error):
		if not self.error or self.advance_count == 0 :
			self.error = error
		return self

########################################################
##	PARSER
########################################################

class Parser:
	def __init__(self,tokens): ##initialise token and token index 
		self.tokens = tokens
		self.tok_idx = -1
		self.advance()

	def advance(self):
		self.tok_idx += 1
		self.update_current_tok()
		return self.current_tok

	def reverse(self, amount=1):
		self.tok_idx -= amount## we take amount we have to reverse and subtract it from the token index
		self.update_current_tok()
		return self.current_tok

	def update_current_tok(self):
		if self.tok_idx >= 0 and self.tok_idx < len(self.tokens):
		  self.current_tok = self.tokens[self.tok_idx]

	#########################################################

	def parse(self):	
		res = self.statements() ##result after parsing expression
		if not res.error and self.current_tok.type != TT_EOF: ##still havennt reached the EOF even though we have finished it means we have a syntax error 
			return res.failure(InvalidSyntaxError(self.current_tok.pos_start, self.current_tok.pos_end, "Expected '+', '-' , '*' '/' or '^' "))
		return res
	#########################################################

	def atom(self):
		res = ParseResult()
		tok = self.current_tok

		if tok.type in (TT_INT,TT_FLOAT):
			res.register_advancements()
			self.advance()
			return res.success(NumberNode(tok)) 

		elif tok.type == TT_STRING:
			res.register_advancements()
			self.advance()
			return res.success(StringNode(tok))

		elif tok.type == TT_IDENTIFIER:
			res.register_advancements()
			self.advance()
			return res.success(VarAccessNode(tok))

		elif tok.type == TT_LPAREN:
			res.register_advancements()
			self.advance()
			expr = res.register(self.expr())
			if res.error:
				return res
			if self.current_tok.type == TT_RPAREN:
				res.register_advancements()
				self.advance()
				return res.success(expr)
			else:
				return res.failure(InvalidSyntaxError(self.current_tok.pos_start, self.current_tok.pos_end, "Expected ')' "))

		elif tok.type == TT_LSQUARE:
		  list_expr = res.register(self.list_expr())
		  if res.error: 
		  	return res
		  return res.success(list_expr)

		elif tok.matches(TT_KEYWORD, 'IF'):
			if_expr = res.register(self.if_expr())
			if res.error:
				return res
			return res.success(if_expr)

		elif tok.matches(TT_KEYWORD, 'FOR'):
			for_expr = res.register(self.for_expr())
			if res.error:
				return res
			return res.success(for_expr)

		elif tok.matches(TT_KEYWORD, 'WHILE'):
			while_expr = res.register(self.while_expr())
			if res.error:
				return res
			return res.success(while_expr)

		elif tok.matches(TT_KEYWORD, 'FUN'):
			func_def = res.register(self.func_def())
			if res.error:
				return res
			return res.success(func_def)

		return res.failure(InvalidSyntaxError(tok.pos_start, tok.pos_end, "Expected int, float, Identifier or '+', '-', or '(' ,'[', 'IF', 'FOR', 'WHILE', 'FUN'")) ## if the token type is not INT or FLOAT

	###################################

	def statements(self):

		res = ParseResult()
		statements = [] ##for the list of expressions 
		pos_start = self.current_tok.pos_start.copy()

		while self.current_tok.type == TT_NEWLINE: ##if current token is new line 
			res.register_advancements()
			self.advance()

		statement = res.register(self.statement()) ##add expression to statement 
		if res.error: return res
		statements.append(statement) ##append to statements 

		more_statements = True

		while True:
			newline_count = 0
			while self.current_tok.type == TT_NEWLINE: ##if more that one lien 
				res.register_advancements()
				self.advance()
				newline_count += 1 ##increment new line 
			if newline_count == 0:
				more_statements = False

			if not more_statements: 
				break

			statement = res.try_register(self.statement()) ##try and look for rules 

			if not statement:
				self.reverse(res.to_reverse_count) ##expression count would have moved forward a bunch of times so we need to reverse. 
				more_statements = False
				continue
			statements.append(statement)

		return res.success(ListNode(statements,pos_start,self.current_tok.pos_end.copy()))


	def statement(self):
		res = ParseResult()
		pos_start = self.current_tok.pos_start.copy()

		if self.current_tok.matches(TT_KEYWORD, 'RETURN'): ##we check for the return statement
			res.register_advancement()
			self.advance()

			expr = res.try_register(self.expr()) ##optional expression after return
			if not expr:
				self.reverse(res.to_reverse_count)
			return res.success(ReturnNode(expr, pos_start, self.current_tok.pos_start.copy()))

		if self.current_tok.matches(TT_KEYWORD, 'CONTINUE'): ##check for continue
			res.register_advancement()
			self.advance()
			return res.success(ContinueNode(pos_start, self.current_tok.pos_start.copy()))

		if self.current_tok.matches(TT_KEYWORD, 'BREAK'):
			res.register_advancement()
			self.advance()
			return res.success(BreakNode(pos_start, self.current_tok.pos_start.copy()))

		expr = res.register(self.expr()) ##if none of the keyword we're looking for a normal expression
		if res.error:
			return res.failure(InvalidSyntaxError(self.current_tok.pos_start, self.current_tok.pos_end,"Expected 'RETURN', 'CONTINUE', 'BREAK', 'VAR', 'IF', 'FOR', 'WHILE', 'FUN', int, float, identifier, '+', '-', '(', '[' or 'NOT'"))
		return res.success(expr)


	def list_expr(self):
		res = ParseResult()
		element_nodes = []
		pos_start = self.current_tok.pos_start.copy()

		if self.current_tok.type != TT_LSQUARE:
			return res.failure(InvalidSyntaxError(self.current_tok.pos_start, self.current_tok.pos_end,f"Expected '['"))

		res.register_advancements()
		self.advance()

		if self.current_tok.type == TT_RSQUARE:
			res.register_advancements()
			self.advance()
		else:
			element_nodes.append(res.register(self.expr()))
			if res.error:
				return res.failure(InvalidSyntaxError(self.current_tok.pos_start, self.current_tok.pos_end,"Expected ']', 'VAR', 'IF', 'FOR', 'WHILE', 'FUN', int, float, identifier, '+', '-', '(', '[' or 'NOT'"))

			while self.current_tok.type == TT_COMMA:
				res.register_advancements()
				self.advance()

			element_nodes.append(res.register(self.expr()))
			if res.error: 
				return res

			if self.current_tok.type != TT_RSQUARE:
				return res.failure(InvalidSyntaxError(self.current_tok.pos_start, self.current_tok.pos_end,f"Expected ',' or ']'"))
			res.register_advancements()
			self.advance()

		return res.success(ListNode(element_nodes,pos_start,self.current_tok.pos_end.copy()))

	def if_expr(self):
		res = ParseResult()
		all_cases = res.register(self.if_expr_cases('IF'))
		if res.error: 
			return res
		cases, else_case = all_cases
		return res.success(IfNode(cases, else_case))

	def if_expr_b(self):
		return self.if_expr_cases('ELIF')

	def if_expr_c(self):
		res = ParseResult()
		else_case = None

		if self.current_tok.matches(TT_KEYWORD, 'ELSE'):
			res.register_advancements()
			self.advance()

			if self.current_tok.type == TT_NEWLINE:
				res.register_advancements()
				self.advance()

				statements = res.register(self.statements())
				if res.error: return res
				else_case = (statements, True)

				if self.current_tok.matches(TT_KEYWORD, 'END'):
				  res.register_advancements()
				  self.advance()
				else:
				  return res.failure(InvalidSyntaxError(
				    self.current_tok.pos_start, self.current_tok.pos_end,
				    "Expected 'END'"
				  ))
			else:
				expr = res.register(self.statement())
				if res.error: return res
				else_case = (expr, False)

		return res.success(else_case)

	def if_expr_b_or_c(self):
		res = ParseResult()
		cases, else_case = [], None

		if self.current_tok.matches(TT_KEYWORD, 'ELIF'):
		  all_cases = res.register(self.if_expr_b())
		  if res.error: return res
		  cases, else_case = all_cases
		else:
		  else_case = res.register(self.if_expr_c())
		  if res.error: return res

		return res.success((cases, else_case))

	def if_expr_cases(self, case_keyword):
		res = ParseResult()
		cases = []
		else_case = None

		if not self.current_tok.matches(TT_KEYWORD, case_keyword):
		  return res.failure(InvalidSyntaxError(self.current_tok.pos_start, self.current_tok.pos_end,f"Expected '{case_keyword}'"))

		res.register_advancements()
		self.advance()

		condition = res.register(self.expr())
		if res.error: 
			return res

		if not self.current_tok.matches(TT_KEYWORD, 'THEN'):
			return res.failure(InvalidSyntaxError(self.current_tok.pos_start, self.current_tok.pos_end,f"Expected 'THEN'"))

		res.register_advancements()
		self.advance()

		if self.current_tok.type == TT_NEWLINE:
		  res.register_advancements()
		  self.advance()

		  statements = res.register(self.statements())
		  if res.error: 
		  	return res
		  cases.append((condition, statements, True))

		  if self.current_tok.matches(TT_KEYWORD, 'END'):
		    res.register_advancements()
		    self.advance()
		  else:
		    all_cases = res.register(self.if_expr_b_or_c())
		    if res.error: 
		    	return res
		    new_cases, else_case = all_cases
		    cases.extend(new_cases)
		else:
		  expr = res.register(self.statement())
		  if res.error: 
		  	return res
		  cases.append((condition, expr, False))

		  all_cases = res.register(self.if_expr_b_or_c())
		  if res.error: return res
		  new_cases, else_case = all_cases
		  cases.extend(new_cases)

		return res.success((cases, else_case))

	def for_expr(self):
		res = ParseResult()

		if not self.current_tok.matches(TT_KEYWORD, 'FOR'):
			return res.failure(InvalidSyntaxError(self.current_tok.pos_start, self.current_tok.pos_end,f"Expected 'FOR'"))

		res.register_advancement()
		self.advance()

		if self.current_tok.type != TT_IDENTIFIER:
			return res.failure(InvalidSyntaxError(self.current_tok.pos_start, self.current_tok.pos_end,f"Expected identifier"))

		var_name = self.current_tok
		res.register_advancement()
		self.advance()

		if self.current_tok.type != TT_EQ:
		  return res.failure(InvalidSyntaxError(self.current_tok.pos_start, self.current_tok.pos_end,f"Expected '='"))

		res.register_advancement()
		self.advance()

		start_value = res.register(self.expr())
		if res.error: return res

		if not self.current_tok.matches(TT_KEYWORD, 'TO'):
			return res.failure(InvalidSyntaxError(self.current_tok.pos_start, self.current_tok.pos_end,f"Expected 'TO'"))

		res.register_advancement()
		self.advance()

		end_value = res.register(self.expr())
		if res.error: return res

		if self.current_tok.matches(TT_KEYWORD, 'STEP'):
		  res.register_advancement()
		  self.advance()

		  step_value = res.register(self.expr())
		  if res.error: return res
		else:
		  step_value = None

		if not self.current_tok.matches(TT_KEYWORD, 'THEN'):
			return res.failure(InvalidSyntaxError(self.current_tok.pos_start, self.current_tok.pos_end,f"Expected 'THEN'"))

		res.register_advancement()
		self.advance()

		if self.current_tok.type == TT_NEWLINE:
			res.register_advancement()
			self.advance()

			body = res.register(self.statements())
			if res.error: return res

			if not self.current_tok.matches(TT_KEYWORD, 'END'):
				return res.failure(InvalidSyntaxError(self.current_tok.pos_start, self.current_tok.pos_end,f"Expected 'END'"))

			res.register_advancement()
			self.advance()

			return res.success(ForNode(var_name, start_value, end_value, step_value, body, True))

		body = res.register(self.statement())
		if res.error: return res

		return res.success(ForNode(var_name, start_value, end_value, step_value, body, False))

	def while_expr(self):
		res = ParseResult()

		if not self.current_tok.matches(TT_KEYWORD, 'WHILE'):
		  return res.failure(InvalidSyntaxError(self.current_tok.pos_start, self.current_tok.pos_end,f"Expected 'WHILE'"))

		res.register_advancement()
		self.advance()

		condition = res.register(self.expr())
		if res.error: return res

		if not self.current_tok.matches(TT_KEYWORD, 'THEN'):
			return res.failure(InvalidSyntaxError(self.current_tok.pos_start, self.current_tok.pos_end,f"Expected 'THEN'"))

		res.register_advancement()
		self.advance()

		if self.current_tok.type == TT_NEWLINE:
			res.register_advancement()
			self.advance()

			body = res.register(self.statements())
			if res.error: return res

			if not self.current_tok.matches(TT_KEYWORD, 'END'):
				return res.failure(InvalidSyntaxError(self.current_tok.pos_start, self.current_tok.pos_end,f"Expected 'END'"))

			res.register_advancement()
			self.advance()

			return res.success(WhileNode(condition, body, True))

		body = res.register(self.statement())
		if res.error: return res

		return res.success(WhileNode(condition, body, False))
	
	def power(self):
		return self.bin_op(self.call,(TT_POW,),self.factor)

	def call(self):
		res = ParseResult()
		atom = res.register(self.atom())
		if res.error: 
			return res

		if self.current_tok.type == TT_LPAREN:
			res.register_advancements()
			self.advance()
			arg_nodes = []

			if self.current_tok.type == TT_RPAREN:
				res.register_advancements()
				self.advance()
			else:
				arg_nodes.append(res.register(self.expr()))
				if res.error:
					return res.failure(InvalidSyntaxError(self.current_tok.pos_start, self.current_tok.pos_end,"Expected ')', 'VAR', 'IF', 'FOR', 'WHILE', 'FUN', int, float, identifier, '+', '-', '(', '[' or 'NOT'"))

				while self.current_tok.type == TT_COMMA:
					res.register_advancements()
					self.advance()

					arg_nodes.append(res.register(self.expr()))
					if res.error: 
						return res

				if self.current_tok.type != TT_RPAREN:
					return res.failure(InvalidSyntaxError(self.current_tok.pos_start, self.current_tok.pos_end,f"Expected ',' or ')'"))

				res.register_advancements()
				self.advance()

			return res.success(CallNode(atom, arg_nodes))
		return res.success(atom)

	def factor(self):

		res = ParseResult()
		tok = self.current_tok

		if tok.type in (TT_PLUS, TT_MINUS):

			res.register_advancements()
			self.advance()

			factor = res.register(self.factor())

			if res.error:
				return res
			
			else:
				return res.success(UnaryOpNode(tok,factor))

		return self.power()
		

	def term(self):
		return self.bin_op(self.factor, (TT_MUL,TT_DIV))

	def arith_expr(self):
		return self.bin_op(self.term, (TT_PLUS,TT_MINUS))

	def comp_expr(self):
		res = ParseResult()

		if self.current_tok.matches(TT_NE,'NE'):
			op_tok = self.current_tok
			res.register_advancements()
			self.advance()

			node = res.register(self.comp_expr())

			if res.error:
				return res

			return res.success(UnaryOpNode(op_tok,node))

		node = res.register(self.bin_op(self.arith_expr,(TT_EE,TT_NE,TT_LT,TT_GT,TT_LTE,TT_GTE)))

		if res.error:
			return res.failure(InvalidSyntaxError(self.current_tok.pos_start, self.current_tok.pos_end, "Expected int, float, identifier, '+', '-', '(', '[', 'IF', 'FOR', 'WHILE', 'FUN' or 'NOT'"))

		return res.success(node)

	def expr(self):
		res = ParseResult()
		if self.current_tok.matches(TT_KEYWORD,'VAR'):
			res.register_advancements()
			self.advance()

			if self.current_tok.type != TT_IDENTIFIER:
				return res.failure(InvalidSyntaxError(self.current_tok.pos_start,self.current_tok.pos_end,"Expected identifier"))

			var_name = self.current_tok

			res.register_advancements()
			self.advance()

			if self.current_tok.type != TT_EQ:
				return res.failure(InvalidSyntaxError(self.current_tok.pos_start,self.current_tok.pos_end,"Expected '=' "))

			res.register_advancements()
			self.advance()
			expr = res.register(self.expr())
			if res.error:
				return res

			return res.success(VarAssignNode(var_name,expr))


		node = res.register(self.bin_op(self.comp_expr, ((TT_KEYWORD,"AND"),(TT_KEYWORD,"OR"))))

		if res.error:
			return res.failure(InvalidSyntaxError(self.current_tok.pos_start,self.current_tok.pos_end,"Expected 'VAR', 'IF', 'FOR', 'WHILE', 'FUN', int, float, identifier, '+', '-', '(', '[' or 'NOT'")) ##overwrite error with term 
		else:
			return res.success(node)

	def bin_op(self, func_a, ops, func_b=None):
		if func_b == None:
			func_b = func_a

		res = ParseResult()
		left = res.register(func_a()) ## if there is an error in the func it will be registered to the res obj 
		if res.error:
			return res
		while self.current_tok.type in ops or (self.current_tok.type, self.current_tok.value) in ops:
			op_tok = self.current_tok
			res.register_advancements()
			self.advance()
			right = res.register(func_b())
			if res.error:
				return res
			left = BinOpNode(left,op_tok,right)
		return res.success(left)

	def func_def(self):

		res = ParseResult()

		if not self.current_tok.matches(TT_KEYWORD, 'FUN'):
			return res.failure(InvalidSyntaxError(self.current_tok.pos_start, self.current_tok.pos_end,f"Expected 'FUN'"))

		res.register_advancement()
		self.advance()

		if self.current_tok.type == TT_IDENTIFIER:
			var_name_tok = self.current_tok
			res.register_advancement()
			self.advance()
			if self.current_tok.type != TT_LPAREN:
				return res.failure(InvalidSyntaxError(self.current_tok.pos_start, self.current_tok.pos_end,f"Expected '('"))
		else:
			var_name_tok = None
			if self.current_tok.type != TT_LPAREN:
				return res.failure(InvalidSyntaxError(self.current_tok.pos_start, self.current_tok.pos_end,f"Expected identifier or '('"))

		res.register_advancement()
		self.advance()
		arg_name_toks = []

		if self.current_tok.type == TT_IDENTIFIER:
			arg_name_toks.append(self.current_tok)
			res.register_advancement()
			self.advance()

			while self.current_tok.type == TT_COMMA:
				res.register_advancement()
				self.advance()

				if self.current_tok.type != TT_IDENTIFIER:
					return res.failure(InvalidSyntaxError(self.current_tok.pos_start, self.current_tok.pos_end,f"Expected identifier"))

				arg_name_toks.append(self.current_tok)
				res.register_advancement()
				self.advance()

				if self.current_tok.type != TT_RPAREN:
					return res.failure(InvalidSyntaxError(self.current_tok.pos_start, self.current_tok.pos_end,f"Expected ',' or ')'"))

		else:
			if self.current_tok.type != TT_RPAREN:
				return res.failure(InvalidSyntaxError(self.current_tok.pos_start, self.current_tok.pos_end,f"Expected identifier or ')'"))

		res.register_advancement()
		self.advance()

		if self.current_tok.type == TT_ARROW:
			res.register_advancement()
			self.advance()

			body = res.register(self.expr())
			if res.error: return res

			return res.success(FuncDefNode(var_name_tok,arg_name_toks,body,True))

		if self.current_tok.type != TT_NEWLINE:
			return res.failure(InvalidSyntaxError(self.current_tok.pos_start, self.current_tok.pos_end,f"Expected '->' or NEWLINE"))

		res.register_advancement()
		self.advance()

		body = res.register(self.statements())
		if res.error: return res

		if not self.current_tok.matches(TT_KEYWORD, 'END'):
			return res.failure(InvalidSyntaxError(self.current_tok.pos_start, self.current_tok.pos_end,f"Expected 'END'"))

		res.register_advancement()
		self.advance()

		return res.success(FuncDefNode(var_name_tok,arg_name_toks,body, False))##no arrow
