#include "interpreter.h"

#include <stdalign.h>

static int CallFrame_Create(CallFrame* self, FlupFunction* self_function)
{
	if (DynamicArray_Create(&self->stack, sizeof(Value), 8, NULL)) {
		return ENOMEM;
	}
	self->memory_pad.memory = NULL;
	self->self_function = self_function;
	self->alternative = self->self_function->alternatives;
	self->instruction_pointer = self->alternative->condition_token_start;
	
	self->variables = NULL;
	self->functions = NULL;

	return EXIT_SUCCESS;
}

static void CallFrame_Destroy(CallFrame* self)
{
	if (self->memory_pad.memory != NULL) {
		Scratchpad_Destroy(&self->memory_pad);
	}

	DynamicArray_Destroy(&self->stack);
}

static void* CallFrame_Reserve(CallFrame* self, size_t amount, size_t alignment)
{
	if (self->memory_pad.memory == NULL) {
		if (Scratchpad_Create(&self->memory_pad, 1024, NULL)) return NULL;
	}

	return Scratchpad_ReserveAligned(&self->memory_pad, amount, alignment);
}

static FlupFunctionAlternative* FlupFunctionAlternative_Malloc(size_t condition_token_start, size_t condition_token_end, size_t body_token_start, size_t body_token_end)
{
	FlupFunctionAlternative* a = malloc(sizeof(FlupFunctionAlternative));
	a->next = NULL;
	a->condition_token_start = condition_token_start;
	a->condition_token_end   = condition_token_end;
	a->body_token_start      = body_token_start;
	a->body_token_end        = body_token_end;

	return a;
}

int Interpreter_Create(Interpreter* self, DynamicArray* tokens)
{
	if (DynamicArray_Create(&self->call_frames, sizeof(CallFrame), 16, NULL)) {
		return ENOMEM;
	}
	self->tokens = tokens;

	return EXIT_SUCCESS;
}

Token* Interpreter_GetToken(Interpreter* self, size_t index)
{
	return DynamicArray_GetPointer(self->tokens, index);
}

Token* Interpreter_ExpectToken(Interpreter* self, CallFrame* top_frame, size_t stop_index)
{
	if (top_frame->instruction_pointer > stop_index) {
		return NULL;
	}

	return Interpreter_GetToken(self, top_frame->instruction_pointer);
}

int Interpreter_ParseFunction(Interpreter* self, CallFrame* top_frame, size_t stop_index)
{
	FlupFunction* f = CallFrame_Reserve(top_frame, sizeof(FlupFunction), alignof(FlupFunction));
	if (f == NULL) return ENOMEM;

	Token* function_name = Interpreter_GetToken(self, top_frame->instruction_pointer);

	f->name = function_name->get.identifier;
	top_frame->instruction_pointer++;

	// one token colon
	top_frame->instruction_pointer++;

	Token* parameter_type = Interpreter_ExpectToken(self, top_frame, stop_index);
	f->parameters = NULL;
	ParameterDefinition** pdef = &f->parameters;
	while (parameter_type->type != TOKENTYPE_ARROW) {
		top_frame->instruction_pointer++;
		Token* parameter_name = Interpreter_ExpectToken(self, top_frame, stop_index);
		top_frame->instruction_pointer++;

		if (parameter_type->type != TOKENTYPE_IDENTIFIER) {
			// TODO: Error message
			return EXIT_FAILURE;
		}
		if (parameter_name->type != TOKENTYPE_IDENTIFIER) {
			// TODO: Error message
			return EXIT_FAILURE;
		}

		*pdef = CallFrame_Reserve(top_frame, sizeof(**pdef), alignof(**pdef));
		if (*pdef == NULL) {
			return ENOMEM;
		}

		(*pdef)->type = parameter_type->get.identifier;
		(*pdef)->name = parameter_name->get.identifier;
		(*pdef)->next = NULL;
		pdef = &((*pdef)->next);

		parameter_type = Interpreter_ExpectToken(self, top_frame, stop_index);
	}

	// one arrow token
	top_frame->instruction_pointer++;

	{
		Token* return_type = Interpreter_ExpectToken(self, top_frame, stop_index);
		if (return_type == NULL) {
			return EXIT_FAILURE;
		}
		if (return_type->type != TOKENTYPE_IDENTIFIER) {
			return EXIT_FAILURE;
		}
		top_frame->instruction_pointer++;

		f->return_type = return_type->get.identifier;
	}

	// pipe token
	Token* alternative_start = Interpreter_ExpectToken(self, top_frame, stop_index);
	if (alternative_start == NULL) {
		return EXIT_FAILURE;
	}

	FlupFunctionAlternative** fdef = &f->alternatives;
	while (alternative_start->type == TOKENTYPE_PIPE) {
		top_frame->instruction_pointer++;
		size_t condition_start = top_frame->instruction_pointer;

		Token* current;
		do {
			current = Interpreter_ExpectToken(self, top_frame, stop_index);
			top_frame->instruction_pointer++;
			if (current == NULL) {
				return EXIT_FAILURE;
			}
			// TODO: Count inner colon and semicolon to allow for nested functions
		} while (current->type != TOKENTYPE_ARROW);
		size_t condition_end = top_frame->instruction_pointer - 2;

		size_t body_start = top_frame->instruction_pointer;
		do {
			current = Interpreter_ExpectToken(self, top_frame, stop_index);
			top_frame->instruction_pointer++;
			if (current == NULL) {
				return EXIT_FAILURE;
			}
			// TODO: Count inner colon and semicolon to allow for nested functions
		} while (current->type != TOKENTYPE_PIPE && current->type != TOKENTYPE_SEMICOLON);
		top_frame->instruction_pointer--;
		size_t body_end = top_frame->instruction_pointer;

		*fdef = CallFrame_Reserve(top_frame, sizeof(**pdef), alignof(**pdef));
		if (*fdef == NULL) {
			return ENOMEM;
		}
		(*fdef)->condition_token_start = condition_start;
		(*fdef)->condition_token_end   = condition_end;
		(*fdef)->body_token_start = body_start;
		(*fdef)->body_token_end   = body_end;

		fdef = &((*fdef)->next);

		alternative_start = Interpreter_ExpectToken(self, top_frame, stop_index);
		if (alternative_start == NULL) {
			return EXIT_FAILURE;
		}
	}
	
	// one token colon
	top_frame->instruction_pointer++;

	f->next = top_frame->functions;
	top_frame->functions = f;

	return EXIT_SUCCESS;
}

FlupFunction* Interpreter_FindFunction(Interpreter* self, StringView name)
{
	DYNAMICARRAY_REVERSEITERATE(self->call_frames, CallFrame, frame) {
		for (FlupFunction* function = frame->functions; function != NULL; function = function->next) {
			if (StringView_Equal(function->name, name)) {
				return function;
			}
		}
	}

	return NULL;
}

bool _are_types_compatible(enum ValueType type, StringView type_name)
{
	const enum ValueType type_mapping[] = {VALUETYPE_INT64, VALUETYPE_DOUBLE};
	const char* name_mapping[] = {"int", "double"};

	for (size_t i = 0; i < sizeof(name_mapping) / sizeof(name_mapping[0]); i++) {
		if (1==1
			&& type == type_mapping[i] 
			&& StringView_Equal(type_name, StringView_FromString(name_mapping[i]))
			) {
			return true;
		}
	}

	return false;
}

int CallFrame_DefineVariable(CallFrame* self, StringView name, Value value)
{
	FlupVariable* v = CallFrame_Reserve(self, sizeof(FlupVariable), alignof(FlupVariable));
	if (v == NULL) {
		return ENOMEM;
	}

	v->name = name;
	v->value = value;

	v->next = self->variables;
	self->variables = v;

	return EXIT_SUCCESS;
}

int Interpreter_CallFunction(Interpreter* self, CallFrame* parent_frame, Token* name)
{
	FlupFunction* function = Interpreter_FindFunction(self, name->get.identifier);
	if (function == NULL) {
		// TODO: Error message
		return EXIT_FAILURE;
	}

	CallFrame* new_frame;
	DynamicArray_AppendEmpty(&self->call_frames, (void**) &new_frame);
	parent_frame = DynamicArray_GetPointer(&self->call_frames, DynamicArray_GetLength(&self->call_frames) - 2);
	if (new_frame == NULL) {
		// TODO: Error message
		return ENOMEM;
	}
	CallFrame_Create(new_frame, function);

	size_t argument_count = 0;
	size_t stack_size = DynamicArray_GetLength(&parent_frame->stack);
	for (
			ParameterDefinition* d = function->parameters;
			d != NULL;
			d = d->next,
			argument_count++);
	if (stack_size < argument_count) {
		// TODO: Error message
		return EXIT_FAILURE;
	}

	ParameterDefinition* parameter_def = function->parameters;
	for (size_t stack_index = stack_size - argument_count; stack_index != stack_size; stack_index++, parameter_def = parameter_def->next) {
		Value* parameter_value = DynamicArray_GetPointer(&parent_frame->stack, stack_index);
		if (! _are_types_compatible(parameter_value->type, parameter_def->type)) {
			// TODO: Error message
			return EXIT_FAILURE;
		}

		int define_code = CallFrame_DefineVariable(new_frame, parameter_def->name, *parameter_value);
		if (define_code != EXIT_SUCCESS) {
			return define_code;
		}
	}

	for (size_t i = 0; i < argument_count; i++) {
		DynamicArray_Remove(&parent_frame->stack, DynamicArray_GetLength(&parent_frame->stack) - 1);
	}

	return EXIT_SUCCESS;
}

FlupVariable* CallFrame_FindVariable(CallFrame* self, StringView name)
{
	for (FlupVariable* variable = self->variables; variable != NULL; variable = variable->next) {
		if (StringView_Equal(variable->name, name)) {
			return variable;
		}
	}

	return NULL;
}

int CallFrame_StackPop(CallFrame* self, Value* dest)
{
	size_t stack_size = DynamicArray_GetLength(&self->stack);
	Value* popped = DynamicArray_GetPointer(&self->stack, stack_size - 1);
	if (popped == NULL) {
		return ENOTFOUND;
	}

	*dest = *popped;
	DynamicArray_Remove(&self->stack, stack_size - 1);

	return EXIT_SUCCESS;
}

int CallFrame_StackPush(CallFrame* self, Value* value)
{
	return DynamicArray_Append(&self->stack, value);
}

bool Value_Equal(Value* v1, Value* v2)
{
	if (v1->type != v2->type) return false;
	switch (v1->type) {
		case VALUETYPE_INT64:
			return v1->get.i64 == v2->get.i64;
		case VALUETYPE_DOUBLE:
			return v1->get.f64 == v2->get.f64;
		case VALUETYPE_BOOLEAN:
			return v1->get.boolean = v2->get.boolean;
	}

	return false;
}

int BuiltinFunction_Equality(CallFrame* top_frame)
{
	Value v1;
	Value v2;

	CallFrame_StackPop(top_frame, &v1);
	if (CallFrame_StackPop(top_frame, &v2) != EXIT_SUCCESS) {
		// TODO: Error message
		return EXIT_FAILURE;
	}

	Value result;
	result.type = VALUETYPE_BOOLEAN;
	result.get.boolean = Value_Equal(&v1, &v2);

	top_frame->instruction_pointer++;
	return CallFrame_StackPush(top_frame, &result);

}

int BuiltinFunction_Plus(CallFrame* top_frame)
{
	Value v1;
	Value v2;

	CallFrame_StackPop(top_frame, &v1);
	if (CallFrame_StackPop(top_frame, &v2) != EXIT_SUCCESS) {
		// TODO: Error message
		return EXIT_FAILURE;
	}
	if (v1.type != v2.type) {
		// TODO: Error message
		return EXIT_FAILURE;
	}

	Value result;
	result.type = v1.type;
	switch (v1.type) {
		case VALUETYPE_INT64:
			result.get.i64 = v2.get.i64 + v1.get.i64;
			break;
		case VALUETYPE_DOUBLE:
			result.get.f64 = v2.get.f64 + v1.get.f64;
			break;
		case VALUETYPE_BOOLEAN:
			result.get.boolean = v2.get.boolean + v1.get.boolean;
			break;
        }

	top_frame->instruction_pointer++;
	return CallFrame_StackPush(top_frame, &result);
}

int BuiltinFunction_Minus(CallFrame* top_frame)
{
	Value v1;
	Value v2;

	CallFrame_StackPop(top_frame, &v1);
	if (CallFrame_StackPop(top_frame, &v2) != EXIT_SUCCESS) {
		// TODO: Error message
		return EXIT_FAILURE;
	}
	if (v1.type != v2.type) {
		// TODO: Error message
		return EXIT_FAILURE;
	}

	Value result;
	result.type = v1.type;
	switch (v1.type) {
		case VALUETYPE_INT64:
			result.get.i64 = v2.get.i64 - v1.get.i64;
			break;
		case VALUETYPE_DOUBLE:
			result.get.f64 = v2.get.f64 - v1.get.f64;
			break;
		case VALUETYPE_BOOLEAN:
			result.get.boolean = v2.get.boolean - v1.get.boolean;
			break;
        }

	top_frame->instruction_pointer++;
	return CallFrame_StackPush(top_frame, &result);
}

int Interpreter_ExecuteNext(Interpreter* self, size_t stop_token)
{
	size_t frame_count = DynamicArray_GetLength(&self->call_frames);
	CallFrame* top_frame = DynamicArray_GetPointer(&self->call_frames, frame_count - 1);
	Token* t = Interpreter_GetToken(self, top_frame->instruction_pointer);
	switch (t->type) {
		case TOKENTYPE_INTEGER:
		{
			Value* v;
			DynamicArray_AppendEmpty(&top_frame->stack, (void**) &v);
			if (v == NULL) return ENOMEM;
			v->type = VALUETYPE_INT64;
			v->get.i64 = t->get.integer;
			top_frame->instruction_pointer++;
			break;
		}
		case TOKENTYPE_DOUBLE:
		{
			Value* v;
			DynamicArray_AppendEmpty(&top_frame->stack, (void**) &v);
			if (v == NULL) return ENOMEM;
			v->type = VALUETYPE_DOUBLE;
			v->get.f64 = t->get.decimal;
			top_frame->instruction_pointer++;
			break;
		}
		case TOKENTYPE_IDENTIFIER:
		{
			Token* next = Interpreter_GetToken(self, top_frame->instruction_pointer + 1);
			FlupVariable* var;
			if (next != NULL && next->type == TOKENTYPE_COLON) {
				// function definition
				return Interpreter_ParseFunction(self, top_frame, stop_token);
			} else if ((var = CallFrame_FindVariable(top_frame, t->get.identifier)) != NULL) {
				Value* stack_value;
				DynamicArray_AppendEmpty(&top_frame->stack, (void**) &stack_value);
				if (stack_value == NULL) {
					return ENOMEM;
				}
				*stack_value = var->value;
				top_frame->instruction_pointer++;
			} else {
				top_frame->instruction_pointer++;
				return Interpreter_CallFunction(self, top_frame, t);
				// regular call
			}

			break;
		}

		case TOKENTYPE_EQUALITY:
		{
			return BuiltinFunction_Equality(top_frame);
		}

		case TOKENTYPE_PLUS:
		{
			BuiltinFunction_Plus(top_frame);
			break;
		}
		case TOKENTYPE_MINUS:
		{
			BuiltinFunction_Minus(top_frame);
			break;
		}

		case TOKENTYPE_PIPE:
		case TOKENTYPE_SEMICOLON:
		case TOKENTYPE_LEFT_BRACE:
		case TOKENTYPE_RIGHT_BRACE:
		case TOKENTYPE_AMPERSAND:
		case TOKENTYPE_MULTIPLY:
		case TOKENTYPE_DIVIDE:
		case TOKENTYPE_ARROW:
		case TOKENTYPE_COLON:
		case TOKENTYPE_INEQUALITY:
		case TOKENTYPE_LESSTHAN:
		case TOKENTYPE_LESSEQUAL:
		case TOKENTYPE_GREATERTHAN:
		case TOKENTYPE_GREATEREQUAL:
		case TOKENTYPE_COMMA:
			break;

		case TOKENTYPE_NONE:
		case TOKENTYPE_ERROR:
			return EXIT_FAILURE;
        }

	return EXIT_SUCCESS;
}

bool Value_IsTruthy(Value* v)
{
	switch (v->type) {
		case VALUETYPE_INT64:
			return v->get.i64 != 0;
		case VALUETYPE_DOUBLE:
			return v->get.f64 != 0.0;
		case VALUETYPE_BOOLEAN:
			return v->get.boolean;
	}

	return false;
}

static bool CallFrame_IsExecutingCondition(CallFrame* frame)
{
	return frame->instruction_pointer >= frame->alternative->condition_token_start 
		&& frame->instruction_pointer <= frame->alternative->condition_token_end;
}

static bool CallFrame_IsSelectingCondition(CallFrame* frame)
{
	return frame->instruction_pointer == frame->alternative->condition_token_end + 1;
}

static bool CallFrame_IsExecutingBody(CallFrame* frame)
{
	return frame->instruction_pointer < frame->alternative->body_token_end;
}

static bool CallFrame_IsReturningBody(CallFrame* frame)
{
	return frame->instruction_pointer == frame->alternative->body_token_end;
}

#define RPN(a,b,c) a c b
int Interpreter_RunFrame(Interpreter* self)
{
	size_t frame_count = DynamicArray_GetLength(&self->call_frames);
	CallFrame* top_frame = DynamicArray_GetPointer(&self->call_frames, frame_count - 1);
	if (CallFrame_IsExecutingCondition(top_frame)) {
		// alternative executing
		while (top_frame->instruction_pointer <= top_frame->alternative->condition_token_end) {
			int execute_code = Interpreter_ExecuteNext(self, top_frame->alternative->condition_token_end);
			if (execute_code != EXIT_SUCCESS) {
				return execute_code;
			}

			if (RPN(frame_count, DynamicArray_GetLength(&self->call_frames), !=)) {
				return EXIT_SUCCESS;
			}
		}
	} else if (CallFrame_IsSelectingCondition(top_frame)) {
		// alternative evaluation done
		Value alternative_success;
		if (CallFrame_StackPop(top_frame, &alternative_success) != EXIT_SUCCESS) {
			// TODO: Error message
			return EXIT_FAILURE;
		}

		if (Value_IsTruthy(&alternative_success)) {
			// into body
			top_frame->instruction_pointer = top_frame->alternative->body_token_start;
		} else  {
			if (top_frame->alternative->next == NULL) {
				// TODO: Error message
				return EXIT_FAILURE;
			}

			top_frame->alternative = top_frame->alternative->next;
			top_frame->instruction_pointer = top_frame->alternative->condition_token_start;
			// next alternative
		}

		return EXIT_SUCCESS;
	} else if (CallFrame_IsExecutingBody(top_frame)) {
		// function executing
		do {
			int execute_code = Interpreter_ExecuteNext(self, top_frame->alternative->body_token_end);
			if (execute_code != EXIT_SUCCESS) {
				return execute_code;
			}

			if (RPN(frame_count, DynamicArray_GetLength(&self->call_frames), !=)) {
				return EXIT_SUCCESS;
			}
		} while (top_frame->instruction_pointer < top_frame->alternative->body_token_end);
	} else if (CallFrame_IsReturningBody(top_frame)) {
		// Function return
		if (StringView_Equal(top_frame->self_function->return_type, StringView_FromString("void"))) {
			CallFrame_Destroy(top_frame);
			DynamicArray_Remove(&self->call_frames, frame_count - 1);
			return EXIT_SUCCESS;
		}

		CallFrame* parent_frame = DynamicArray_GetPointer(&self->call_frames, frame_count - 2);
		Value return_value;
		if (CallFrame_StackPop(top_frame, &return_value)) {
			// TODO: Error message
			return EXIT_FAILURE;
		}

		// char function_name[top_frame->self_function->name.length + 1];
		// memset(function_name, 0, sizeof(function_name));
		// StringView_Paste(function_name, top_frame->self_function->name);

		if (! _are_types_compatible(return_value.type, top_frame->self_function->return_type)) {
			// TODO: Error message
			return EXIT_FAILURE;
		}

		CallFrame_Destroy(top_frame);
		DynamicArray_Remove(&self->call_frames, frame_count - 1);
		if (CallFrame_StackPush(parent_frame, &return_value)) {
			return ENOMEM;
		}
	} else {
		// Invalid instruction pointer
		// TODO: Error message
		return EXIT_FAILURE;
	}

	return EXIT_SUCCESS;
}


int Interpreter_Run(Interpreter* self)
{
	size_t frame_count = 1;
	do  {
		if (Interpreter_RunFrame(self)) {
			// TODO: Return the error code
			return EXIT_FAILURE;
		}

		frame_count = DynamicArray_GetLength(&self->call_frames);
	} while (frame_count != 0);
	return EXIT_SUCCESS;
}

int Interpreter_Interpret(Interpreter* self)
{
	CallFrame* first_frame;
	FlupFunctionAlternative* main_alternative = FlupFunctionAlternative_Malloc(SIZE_MAX-1, SIZE_MAX-1, 0, DynamicArray_GetLength(self->tokens));
	FlupFunction main_function =  {
		.name = StringView_FromString("<main>"),
		.alternatives = main_alternative,
		.return_type = StringView_FromString("void"),
		.next = NULL,
		.parameters = NULL,
	};
	if (main_alternative == NULL) {
		return ENOMEM;
	}

	DynamicArray_AppendEmpty(&self->call_frames, (void**) &first_frame);
	if (CallFrame_Create(first_frame, &main_function)) {
		return ENOMEM;
	}
	first_frame->instruction_pointer = 0;

	Interpreter_Run(self);

	free(main_alternative);

	return EXIT_SUCCESS;
}

void Interpreter_Destroy(Interpreter* self)
{
	DynamicArray_Destroy(&self->call_frames);
}