// px_typechecker.cpp

#define PX_INTERNAL_SOURCE
#include "px_typechecker.h"

#define TYPE_INFER "var"

static bool abject_failure = false;

/*
The typechecker can infer types for
var x;
x = {something with a type}

It can't infer types for
[var x] {
	x();
}

(see original state of test3). Although it could be extended to do this, by
  scanning multiple layers of inference, since you can(?) know the type by the
  way it's used at the call site

*/

typedef lookup_chain<std::string, PXType*> symbol_table;

#define DIE(err, ...) { \
	printf("%%fatal \"TypeError: " err "\"\n", __VA_ARGS__); \
	abject_failure = true; \
	return NULL; \
}

PXType* walk_type(token* t, symbol_table& s, token* current_function) {
	if (abject_failure) return NULL; // abort
	if (t->vartype.is_valid()) return & t->vartype; // Already exists

	switch (t->type) {

		// block := "{" body "}"
		case TOK_BLOCK: {
			// Walk children with new symbol table, don't return any type
			walk_type(t->children[1], symbol_table(&s), current_function);
			return NULL;
		} break;

		// body := statement [ body ]
		case TOK_BODY: {
			// Walk children with same symbol table, don't return any type
			walk_type(t->children[0], s, current_function);
			if (t->children.size() > 1) {
				walk_type(t->children[1], s, current_function);
			}
			return NULL;
		} break;

		// dimvar := type identifier [ = rvalue ]
		case TOK_DIM_VAR: {
			if (! PXType::unserialize(t->children[0], t->vartype)) {
				DIE("failed to unserialize type in t-dim-var");
			}

			if (t->children.size() > 2) {
				// An initialiser was present, use it for type information.
				walk_type(t->children[3], s, current_function);

				// var this = {something strongly typed}
				if (t->vartype.is_inferable()) {
					t->vartype = t->children[3]->vartype;

				} else if (!t->vartype.equals(t->children[3]->vartype)) {
					DIE("Incompatible type for variable initialiser");
				}

			} else {
				t->children[1]->vartype = t->vartype; 
			}

			
			s.insert(t->children[1]->data, & t->vartype);
			return & t->vartype;
		} break;
			
		// import-clause := "import" typedef identifier
		case TOK_IMPORT_CLAUSE: {
			if (! PXType::unserialize(t->children[1], t->vartype)) {
				DIE("failed to unserialize type in t-dim-var");
			}
			//s.insert(t->children[1]->data, & t->vartype);
			return & t->vartype;
		} break;

		// fndef-arglist := [ type identifier [ "," fndef-arglist ] ]
		case TOK_FNDEF_ARGLIST: {
			// Register the definition in the fndef variable scope
			if (t->children.size()) {
				if (! PXType::unserialize(t->children[0], t->vartype)) {
					DIE("failed to unserialize type in t-fndef-arglist");
				}
				//t->vartype = t->children[0]->data; // not really semantic, but OK
				s.insert(t->children[1]->data, & t->vartype);

				if (t->children.size() > 2) {
					walk_type(t->children[3], s, current_function);
				}
			}
		} break;

		// fndef := "[" fndef-arglist "]" block
		case TOK_FNDEF: {
			// Function parameters have a scope inbetween the definee and the
			//    function body
			symbol_table param_env(& s);
			walk_type(t->children[1], param_env, current_function); // the current_function won't ever be used, no way for fndef-arglist to get to return-stmt without going through fndef and making a new current_function

			// Argument types are now in param_env
			symbol_table fnbody_env(& param_env);
			
			// Pass the current token as current_function, so that we can
			//    get the return type if it's inferred
			// Eventually if we specify fixed return types, you could set it
			//    here

			walk_type(t->children[3], fnbody_env, t);

			// Walking the function has set this token's type to be the return type.
			// Now explicitally wrap this as Function<ReturnType, arg, arg, ...>
			if (! t->vartype.is_valid() ||
				t->vartype.is_inferable()) {
				t->vartype = PXType::make_scalar("NoReturnType"); // whatever is in eax
				// Use NoReturnType as the return type, to prevent accidentally
				//   assigning eax junk to a variable.
			}

			PXType pxt = PXType::make_scalar("function");
			pxt.children.push_back( t->vartype ); // return type
			token* arglist = t->children[1];
			for(;;) {
				if (! arglist->children.size()) break;

				pxt.children.push_back( t->children[1]->vartype );

				if ( arglist->children.size() > 2) {
					arglist = arglist->children[3];
				} else {
					break;
				}
			}

			t->vartype = pxt; // operator= copy

			return & t->vartype; //t->vartype.is_valid() ? & t->vartype : NULL; // void if no return type
		} break;

		// fncall := lvalue "(" fncall-arglist ")"
		case TOK_FNCALL: {
			// The type of a fncall is simply the function's return type
			// FIXME: Once again duplicating some lvalue stuff here... this will have to stop!

			//                           lvalue    := identifier ...
			PXType *functype = s.find(t->children[0]->children[0]->data);
			if (! functype) {
				DIE("unknown function, or tried to assign the result of a void function");
			}

			bool already_checked_children = false;

			// functype (should) have the form function<returntype, args...>
			if (functype->type != "function" ||
				! functype->children.size() ) {

				// if we're calling a var, then see if var's type can be inferred
				//  from the local scope...

				if (functype->is_inferable()) {
					// the correct type for the function we are about to call is
					//  function< ret_type [, argtypes...] >

					// begin infer {{
					//std::string correct_type = "function<var";
					functype->type = "function";
					functype->children.clear();
					// first add the return type
					functype->children.push_back( PXType::make_scalar("var") );

					// 	syntax[TOK_FNCALL_ARGLIST]	= vadder<tok_types>(TOK_RULE_OPTIONAL_REMAINING)(TOK_RVALUE)(TOK_RULE_OPTIONAL_REMAINING)(TOK_OP_COMMA)(TOK_FNCALL_ARGLIST);
					token* arglist = t->children[2];
					while( arglist->children.size() ) {
						PXType *rvtype = walk_type(arglist->children[0], s, current_function);

						//correct_type.append( "," );
						//correct_type.append( rvtype->serialize() );
						functype->children.push_back( *rvtype );

						if (arglist->children.size() == 1) break;
						arglist = arglist->children[2];
					};
					//correct_type.append( ">" );

					already_checked_children = true;

					// infer
					//functype->type = correct_type;
					// }} endinfer

				} else {
					DIE("tried to call something that didn't have the type function<ret,...>");
				}
			}

			if (functype->children[0].is_inferable()) {
				DIE("return type of function must be inferred before called");
			}
			
			// Don't forget to infer types for all the arguments to this function
			if (! already_checked_children) {
				token* arglist = t->children[2];
				while( arglist->children.size() ) {
					PXType *rvtype = walk_type(arglist->children[0], s, current_function);
					if (arglist->children.size() == 1) break;
					arglist = arglist->children[2];
				};
			}

			t->vartype = functype->children[0]; // copy ; //*functype;
			return & t->vartype;
		} break;

		// Things where the type is the type of the first child:
		// statement := statement-inner ";"
		case TOK_STATEMENT: 
		// statement-inner := block | return-stmt | if-stmt | for-stmt | dim-var | fncall | assignment
		case TOK_STATEMENT_INNER:
		// rvalue-inner2 := fn-call | numeric-literal | lvalue | quoted-string | fndef
		case TOK_RVALUE_INNER2:
		{
			PXType *subtype = walk_type(t->children[0], s, current_function);
			if (! subtype) {
				return NULL; // failed to resolve type of RHS
			}

			t->vartype = *subtype;
			return & t->vartype;
		} break;
		
		case TOK_RVALUE_INNER: {
			// rvalue-inner: rvalue-inner2 [ "." fncall ]

			PXType *subtype = walk_type(t->children[0], s, current_function);
			if (!subtype) {
				return NULL;
			}
			
			// Syntax sugar: var.x(y) ==> var's_type_x(var, y);
			if (t->children.size() > 1) {
			
				// Hoist function up-level
				token *tmp = new token(*t); // copy
				// FIXME: free *t WITHOUT calling its dtor!
				*t = *(t->children[2]);
				t->parent = tmp->parent;
				tmp->parent = t;

				// Patch function name
				token *identifier = t->children[0]->children[0];
				identifier->data = subtype->serialize() + "_" + identifier->data;

				// Remove the .call() from the rvalue...
				tmp->children.resize(1);

				// ...and add it as an extra function argument
				token *arglist = new token();
				arglist->type = TOK_FNCALL_ARGLIST;
				arglist->parent = t;

				arglist->children.push_back(tmp);
				arglist->children[0]->parent = arglist;

				if (t->children[2]->children.size()) {
					// Add comma
					token *comma = new token();
					comma->type = TOK_OP_COMMA;
					comma->parent = arglist;
					arglist->children.push_back(comma);

					// Copy remaining args
					t->children[2]->parent = arglist;
					arglist->children.push_back(t->children[2]);
				}

				// Replace existing arglist
				t->children[2] = arglist;
			

				// FIXME: There are still some problems with incorrect token
				//  parents here. But that doesn't seem to cause any codegen
				//  problems yet
				

				// Typecheck the generated funcall
				*subtype = *walk_type(t, s, current_function);
				if (!subtype) {
					return NULL;
				}
			}

			t->vartype = *subtype;
			return &t->vartype;

		} break;

		// rvalue := rvalue-inner [ binary-operator rvalue-inner ]
		case TOK_RVALUE: {
			PXType *lhstype = walk_type(t->children[0], s, current_function);
			if (! lhstype) {
				return NULL; // failed to resolve type of LHS
			}
			
			if (t->children.size() > 1) {
				PXType *rhstype = walk_type(t->children[2], s, current_function);
				if (! rhstype) return NULL; // failed to resolve type of RHS

				if (! lhstype->equals(*rhstype)) {

					// {{{ some primitive attempts at better error messages
					if (lhstype->is_inferable() && t->children[0]->children[0]->type == TOK_LVALUE) {
						fprintf(stderr,
							"%%fatal \"TypeError: Tried to use binary operator on variable '%s' with unknown type\"\n",
							t->children[0]->children[0]->children[0]->data.c_str()
						);
					}
					if (rhstype->is_inferable() && t->children[0]->children[0]->type == TOK_LVALUE) {
						fprintf(stderr,
							"%%fatal \"TypeError: Tried to use binary operator on variable '%s' with unknown type\"\n",
							t->children[2]->children[0]->children[0]->data.c_str()
						);
					}
					fprintf(stderr,
						"%%fatal \"TypeError: types don't match for binary operator ('%s' != '%s')\n",
						lhstype->serialize().c_str(),
						rhstype->serialize().c_str()
					);
					// }}} 

					DIE("types don't match for binary operator");
				}

				/*
				if (! lhstype->is_scalar("int")) {
					DIE("binary operators are only defined for type int");
				}
				*/
			}

			t->vartype = *lhstype;
			return & t->vartype;
		} break;

		// lvalue := identifier [ "[" rvalue "]" ]
		case TOK_LVALUE: {
			// FIXME: For now, let's skip the brackets and assert that the type is
			//   always the type of the identifier
			// Once we add generics to the parser, this can be fixed
			
			//PXType *identtype = walk_type(t->children[0], s, current_function);
			// Type of identifier
			PXType *identtype = s.find(t->children[0]->data);
			if (! identtype) {
				return NULL;
			}
			t->children[0]->vartype = *identtype;

			if (t->children.size() > 2) {

				if (identtype->type != "string" && identtype->type != "vector") {
					DIE("Offset indexing only supported on string|vector, attempted %s", identtype->serialize().c_str());
				}

				PXType *childtype = walk_type(t->children[2], s, current_function);
				if (childtype == NULL) {
					return NULL; // throw
				}
				if (childtype->type != "int") {
					DIE("Offset index should be type int, got %s", childtype->serialize().c_str());
				}

				// Dereferencing a vector by an int produces the child-type of the vector
				if (identtype->type == "vector") {
					identtype = &identtype->children[0]; // FIXME memory leak
				}
			}

			t->vartype = *identtype;
			return & t->vartype;
		} break;

		// assignment := lvalue "=" rvalue
		case TOK_ASSIGNMENT: {
			// FIXME: Same as in the compiler proper, we have to duplicate some work here
			//   with lvalues
			// Ref. above comment re generics

			//                                lvalue :=    identifier ...
			PXType *assignee_type = s.find(t->children[0]->children[0]->data);
			if (! assignee_type) {
				DIE("unknown variable");
			}

			// Look up assignee rvalue's type

			PXType *value_type = walk_type(t->children[2], s, current_function);
			if (! value_type) {
				return NULL; // failed to resolve type of RHS
			}

			if (assignee_type->is_inferable()) {
				*assignee_type = PXType(*value_type); // copy constructor
			}

			if (! assignee_type->equals(*value_type)) {
				DIE("incompatible types for assignment");
			}

			t->vartype = *value_type;
			return & t->vartype;
		} break;

		// return := "return" rvalue
		case TOK_RETURN_STATEMENT: {
			PXType *value_type = walk_type(t->children[1], s, current_function);
			if (! value_type) {
				return NULL; // failed to resolve type of RHS
			}

			if (current_function /* can be null if we're returning from the main script */) {
				if (
					!current_function->vartype.is_valid() || // no type assigned - current syntax doesn't support declaring return types
					current_function->vartype.is_inferable() // type explicitally set to inferable
				) {
					current_function->vartype = *value_type;
				}

				if (! current_function->vartype.equals(*value_type)) {
					DIE("incompatible return type");
				}

			} else {
				// Only allow returning int from the main script
				if (value_type->type != "int") {
					DIE("Returning %s from main script, expected int", value_type->serialize().c_str());
				}
			}

			t->vartype = *value_type;
			return & t->vartype;
		} break;

		// if-stmt := "if" "(" rvalue ")" block [ "else" block ]
		case TOK_IF_STATEMENT: {
			walk_type(t->children[2], s, current_function);
			walk_type(t->children[4], s, current_function);

			bool has_else_clause = (t->children.size() > 5);

			if (has_else_clause) {
				walk_type(t->children[6], s, current_function);
			}
			return NULL;
		} break;

		// for-stmt := "for" "(" stmt rvalue ";" stmt-inner ")" block
		case TOK_FOR_STATEMENT: {
			walk_type(t->children[2], s, current_function); // often a variable declaration(?)
			PXType* condtype = walk_type(t->children[3], s, current_function);
			if (! condtype->is_scalar("int")) {
				DIE("condition type of for loop is not coercible to bool");
			}
			walk_type(t->children[5], s, current_function); // statement-inner

			walk_type(t->children[7], s, current_function);
			return NULL;
		} break;

		case TOK_NUMERIC_LITERAL: {
			t->vartype = PXType::make_scalar("int");
			return & t->vartype;
		} break;

		case TOK_QUOTED_STRING: {
			t->vartype = PXType::make_scalar("string");
			return & t->vartype;
		} break;

		case TOK_SLICE: {
			walk_type(t->children[0], s, current_function);
			t->vartype = t->children[0]->vartype;

			return & t->vartype;
		} break;

		case TOK_SLICE_INNER: {
			PXType* lchild = walk_type(t->children[0], s, current_function);
			PXType* rchild = walk_type(t->children[2], s, current_function);

			if (lchild->type != "int" || rchild->type != "int") {
				DIE(
					"Slice indexes must have type [int:int], got [%s:%s]",
					lchild->serialize().c_str(), rchild->serialize().c_str()
				);
			}

			t->vartype = *lchild;
			return &t->vartype;
		} break;

		case TOK_VECTOR_LITERAL: {
			PXType* childtype = walk_type(t->children[1], s, current_function);

			if (childtype == NULL) {
				// FIXME: we can do better than this
				DIE("Please specify concrete type for vector<var>");
			}

			t->vartype = PXType::make_scalar("vector");
			t->vartype.children.push_back(*childtype);

			return &t->vartype;

		} break;

		case TOK_VECTOR_ELEMENT: {
			if (!t->children.size()) {
				return NULL; // empty vector literal
			}

			t->vartype = *walk_type(t->children[0], s, current_function);

			if (t->children.size() > 1) {
				// Walk next vector element

				walk_type(t->children[2], s, current_function);

				// Assert types are the same for consecutive (i.e. all) vector
				//   elements

				if (!t->vartype.equals(t->children[2]->vartype)) {
					DIE(
						"Vectors must have one type only (found %s, expected %s)",
						t->children[2]->vartype.serialize().c_str(),
						t->vartype.serialize().c_str()
					);
				}
			}

			return &t->vartype;
		} break;

		case TOK_IDENTIFIER:
		case TOK_TYPE_IDENTIFIER:
		case TOK_FNCALL_ARGLIST:
		{
			DIE("tried to look up an unexpected type"); // Shouldn't ever get here
			abject_failure = true;
		} break;

		case TOK_RWORD_BREAK:
		case TOK_RWORD_CONTINUE:
		{
			// Nothing to do
		} break;

		default: {
			DIE("unexpected type to look up");
			abject_failure = true;
		} break;
	};

	return NULL; // type wasn't determined
}

//
// PUBLIC INTERFACE
//

bool typecheck(token* ast) {
	if (ast->type == TOK_INVALID) return false;

	symbol_table env(NULL);
	abject_failure = false;
	
	token* current_function = NULL;
	if (ast->type == TOK_FNDEF) current_function = ast;

	walk_type(ast, env, current_function);

	return ! abject_failure;
};

