import "ppp_tokens.ppp";
import "ppp_lexer.ppp";
import "ppp_ast.ppp";
import "ppp_parser.ppp";
import "ppp_types.ppp";
import "ppp_object.ppp";
import "ppp_stdlib.ppp";

enum VariableState {
	Declared(Type, Object),
	Undeclared(Type),
	Constant(Type, Object)
}

func declared_from_obj(obj: Object) -> VariableState return VariableState.Declared(object_get_type(obj), obj);

type Module = dict[str, VariableState];

struct Program {
	modules: dict[str, Module],
	contexts: dict[str, VariableState][]
}

func program_exists(program: Program, name: str) -> bool {
	for context in program.contexts do {
		for name_ in context do {
			if name_ == name return true;
		}
	}
	return false;
}

func program_access_variable(program: Program, name: str) -> Object {
	i: int = len(program.contexts) - 1;
	while i >= 0 do {
		context: dict[str, VariableState] = program.contexts[i];
		for name_ in context do {
			if name_ == name do {
				value: VariableState = context[name];
				match value in {
					case Declared(_, value) return value;
					case Const(_, value) return value;
					case Undeclared(_) assert false, "'%s' is not defined!" % name;
					case _ assert false, "Unimplemented program_access_variable %s" % value;
				}
			}
		}
		i = i - 1;
	}
	assert false, "'%s' is not defined!" % name;
}

func program_declare_variable(program: Program, name: str, type_: Type) {
	for name_ in program.contexts[len(program.contexts)-1] assert name_ != name, "'%s' has already been declared!" % name;
	program.contexts[len(program.contexts)-1][name] = VariableState.Undeclared(type_);
}

func program_assign_variable(program: Program, name: str, value: Object) {
	i: int = len(program.contexts) - 1;
	while i >= 0 do {
		context: dict[str, VariableState] = program.contexts[i];
		for name_ in context do {
			if name_ == name do {
				variable_type: Type;
				match context[name] in {
					case Undeclared(type_) variable_type = type_;
					case _ assert false, "Unimplemented program_assign_variable %s" % context[name];
				}
				assert type_is_subtype_of(object_get_type(value), variable_type), "In the assignment of '%s', expected value of type '%s', but got a value of type '%s'!" % (name, type_represent(variable_type, type_represent(object_get_type(value))));
				context[name] = VariableState.Declared(variable_type, value);
				return none;
			}
		}
		i = i - 1;
	}
	assert false, "'%s' doesn't exist!" % name;
}

func program_declare_and_assign_variable(program: Program, name: str, value: Object) {
	program_declare_variable(program, name, object_get_type(value));
	program_assign_variable(program, name, value);
}

func calculate_expression(expression: Expression, program: Program) -> Object {
	match expression in {
		case String(string) return Object.Str(string);
		case Array(array_) do {
			if len(array_) == 0 return Object.List(Type.List(Type.Variable("")), []);
			elements_type: Type;
			array_elements: Object[] = [];
			for i in range(0, len(array_)) do {
				element: Object = calculate_expression(array_[i], program);
				if i == 0 elements_type = object_get_type(element);
				else assert type_is_subtype_of(object_get_type(element), elements_type), "Array element invalid type";
				array_elements = array_elements + [element];
			}
			assert false, "array %s" % array_elements;
		}
		case Dictionary(dict_) do {
			if len(dict_) == 0 return Object.Dictionary(Type.Dictionary(Type.Variable(""), Type.Variable("")), {});
			key_type: Type;
			value_type: Type;
			dict: dict[Object, Object] = {};
			for i in range(0, len(dict_)) do {
				key: Object = calculate_expression(dict_[i][0], program);
				value: Object = calculate_expression(dict_[i][1], program);
				if i == 0 do {
					key_type = object_get_type(key);
					value_type = object_get_type(value);
				} else {
					assert type_is_subtype_of(object_get_type(key), key_type), "Dict element invalid key type";
					assert type_is_subtype_of(object_get_type(value), value_type), "Dict element invalid value type";
				}
				dict[key] = value;
			}
			assert false, "dict %s" % dict;
		}
		case FieldAccess(expression_, field) do {
			value: Object = calculate_expression(expression_, program);
			match value in {
				case Type(type_) match type_ in {
					case Enum(name, members, _) do {
						for member in members do {
							if member == field do {
								if len(members[field]) == 0 return Object.EnumValue(type_, field, []);
								func return_member(name: str, parameters: (str, Type)[], return_type: Type, statement: Statement, args: Object[]) -> Object {
									match return_type in {
										case Enum(_, _, _) do {}
										case _ assert false, "Unreachable";
									}

									assert len(args) == len(parameters), "%s.%s expected %s arguments but got %s!" % (type_represent(type_), field, int_to_str(len(parameters)), int_to_str(len(args)));
									for i in range(0, len(args)) assert type_is_subtype_of(object_get_type(args[i]), parameters[i][1]);
									return Object.EnumValue(return_type, name, args);
								}
								return Object.Function(Type.Function(members[field], type_), (field, [("", member_type) for member_type in members[field]], type_, Statement.Statements([]), return_member));
							}
						}
						assert false, "g";
					}
					case _ assert false, "Unimplemented calculate_expression field access type %s" % type_represent(type_);
				}
				case _ assert false, "Unimplemented calculate_expression field access %s" % value;
			}
		}
		case Variable(name) return program_access_variable(program, name);
		case _ assert false, "Unimplemented calculate_expression %s" % expression;
	}
}

func calculate_type_expression(expression: TypeExpression, program: Program, must_resolve: bool) -> Type {
	match expression in {
		case Name(name) do {
			if !program_exists(program, name) && !must_resolve return Type.Variable(name);
			type_obj: Object = program_access_variable(program, name);
			match type_obj in {
				case Type(type_) return type_;
				case _ assert false, "Unimplemented %s" % type_obj;
			}
		}
		case List(type_) return Type.List(calculate_type_expression(type_, program, must_resolve));
		case Tuple(types) return Type.Tuple([calculate_type_expression(type_, program, must_resolve) for type_ in types]);
		case _ assert false, "Unimplemented calculate_type_expression %s" % expression;
	}
}

func update_types(type_: Type, program: Program) {
	type_name: str;
	match type_ in {
		case Enum(name, _, _) type_name = name;
		case Struct(name, _, _) type_name = name;

		case _ assert false, "Unimplemented update_types %s" % type_represent(type_);
	}
	for context in program.contexts do {
		for variable_ in context do {
			match context[variable_] in {
				case Declared(variable_type, value) match value in {
					case Type(variable_type_) do {
						assert variable_type == Type.Type;
						type_fill(variable_type_, {type_name: type_}, []);
					}
				}
				case Undeclared(_) do {}
				case _ assert false, "Unimplemented update_types %s" % context[variable_];
			}
		}
	}
}

enum Result {
	Return(Object),
	Continue,
	Break,
	Nothing
}

func interpret_statements(statements: Statement[], program: Program) -> Result {
	for statement in statements do {
		match statement in {
			case Import(file_) do {
				file_path: Object = calculate_expression(file_, program);
				match file_path in {
					case Str(file_path_str) do {
						found: bool = false;
						module: Module;
						for module_ in program.modules do {
							if module_ == file_path_str do {
								found = true;
								module = program.modules[module_];
							}
						}
						if !found do {
							print("Importing %s\n" % file_path_str);
							module = interpret_file(file_path_str, program.modules);
						}

						for variable in module do program.contexts[0][variable] = module[variable];

						if !found program.modules[file_path_str] = module;
					}
					case _ assert false, "Unimplemented interpret_statement import %s" % file_path;
				}
			}
			case EnumDefinition(name, entries) do {
				enum_type: Type = Type.Enum(name, {entry.name: [calculate_type_expression(type_, program, false) for type_ in entry.types] for entry in entries}, []);
				program_declare_and_assign_variable(program, name, Object.Type(enum_type));
				update_types(enum_type, program);
			}
			case FunctionDefinition(name, arguments_, return_type_, body) do {
				func run_function(name: str, arguments: (str, Type)[], return_type: Type, body: Statement, args: Object[]) -> Object {
					assert false, "run_function";
				}

				arguments: (str, Type)[] = [(argument.name, calculate_type_expression(argument.type_, program, true)) for argument in arguments_];
				return_type: Type;
				match return_type_ in {
					case Some(type_) return_type = calculate_type_expression(type_, program, true);
					case None return_type = Type.Void;
				}
				function_type: Type = Type.Function([argument[1] for argument in arguments], return_type);
				object: Object = Object.Function(function_type, (name, arguments, return_type, body, run_function));
				program_declare_and_assign_variable(program, name, object);
			}
			case StructDefinition(name, entries) do {
				struct_type: Type = Type.Struct(name, {entry.name: calculate_type_expression(entry.type_, program, false) for entry in entries}, []);
				program_declare_and_assign_variable(program, name, Object.Type(struct_type));
				update_types(struct_type, program);
			}
			case Assignment(lhs, rhs, type_) do {
				assert is_valid_target(lhs);
				match lhs in {
					case Variable(name) do {
						value: Object = calculate_expression(rhs, program);
						match type_ in {
							case _ assert false, "uuuu %s" % type_;
						}
					}
					case _ assert false, "Unimplemented interpret_statement assignment %s" % lhs;
				}
			}
			case _ assert false, "Unimplemented interpret_statement %s" % statement;
		}
	}
	return Result.Nothing;
}


func interpret_file(file_path: str, modules: dict[str, Module]) -> Module {
	print("Parsing %s...\n" % file_path);
	lexer: Lexer = lexer_from_file(file_path);
	statements: Statement[] = [];
	while !is_some_token(lexer_take_token(lexer, TokenContents.Eof)) statements = statements + [parse_statement(lexer)];
	new_variables: dict[str, VariableState] = {};
	for variable in variables new_variables[variable] = declared_from_obj(variables[variable]);
	program: Program = Program{modules=modules, contexts=[new_variables, {}]};
	print("Interpreting %s...\n" % file_path);
	return_value: Result = interpret_statements(statements, program);
	assert len(program.contexts) == 2;
	match return_value in {
		case Nothing do {}
		case Return(_) assert false, "Cannot return from outside a function!";
		case Continue assert false,  "Cannot continue from outside a loop!";
		case Break assert false, "Cannot break from outside a loop!";
		case _ assert false, "Unimplemented interpret_file return_value";
	}
	return program.contexts[1];
}