from emis_funky_funktions import *
from typing import Collection, Sequence, TypeAlias

from ir import Expression, ReplHole, subst_all
from ir.genir import json_to_ir, PatternParseProblem, BranchTypesDiffer, UndefinedVariable
from ir.types_ import BUILTINS_CONTEXT, UnificationError

import json
from dataclasses import dataclass
from operator import add

def evaluate(expr: Expression) -> Expression:
	"""
	>>> funktion = Function((
	...    (SPattern(NamePattern('n')), (
	...        Application((
	...            Variable('+'),
	...            Variable('n'),
	...            Variable('n'),
	...        ))
	...    )),
	...    (NamePattern('0'), (
	...        Int(1312)
	...    )),
	... ))
	>>> evaluate(Application((funktion, Int(3))))
	4
	>>> evaluate(Application((funktion, Int(0))))
	1312
	"""
	while not expr.is_value():
		match expr.step():
			case Some(next):
				expr = next
				# Loop
			case None:
				raise AssertionError('Evaluate called on a value which cannot step:', expr)
	return expr

def repl_expr(expr: Expression, bindings: ReplHole = ReplHole(BUILTINS_CONTEXT)):
	expr_subst = subst_all(bindings.val_bindings, expr)
	result = evaluate(expr_subst)
	if isinstance(result, ReplHole):
		new_bindings = [(name, type) for (name, type) in result.typ_bindings.variable_types.items() if name not in bindings.typ_bindings]
		print('\n'.join(f'\033[39m{var_name}\033[39m : \033[94m{var_type}' for (var_name, var_type) in new_bindings))
		print('\033[39m')
		repl(result)
	else:
		print(result, end='\n\n')
		repl(bindings)

def repl(bindings: ReplHole = ReplHole(BUILTINS_CONTEXT)):
	print('Enter a JSON expression:')
	try:
		expr = input('\033[95m~> \033[96m')
	except EOFError:
		print('\033[91m<exit>\033[39m')
		return
	try:
		ast = json.loads(expr)
	except json.decoder.JSONDecodeError as e:
		print('\n\033[91m  -----[ ERROR ]-----\033[39m')
		print(f'Bad json: ', e.args[0], end='\n\n')
		print('\033[39m')
		return repl(bindings)
	# TODO handle this
	match json_to_ir(ast, bindings.typ_bindings):
		case Ok((new_expr, new_ty, substs)):
			print('\033[39m')
			repl_expr(new_expr, bindings)
		case Err(err):
			print('\n\033[91m  -----[ ERROR ]-----\033[39m')
			match err:
				case PatternParseProblem(pattern):
					print(f"Couldn't parse the pattern {repr(pattern)}")
				case UnificationError(ty1, ty2):
					print(f"Couldn't unify the types {ty1} and {ty2}")
				case BranchTypesDiffer(b1_pattern, b2_pattern, b1_type, b2_type):
					print(f"Two or more branches had different types!")
					print(f"")
					print(f"The branch that matches the pattern\033[96m")
					print(b1_pattern)
					print(f"\033[39mhas the type\033[94m")
					print(b1_type)
					print(f"\033[39m")
					print(f"The branch that matches the pattern\033[96m")
					print(b2_pattern)
					print(f"\033[39mhas the type\033[94m")
					print(b2_type)
					print(f"\033[39m")
				case UndefinedVariable(var_name):
					print(f'The variable {repr(var_name)} was referenced but not defined')
			print('\033[39m')
			repl(bindings)

if __name__ == '__main__':
	import doctest
	doctest.testmod()