#include"ast.h"

#include<stdint.h>
#include<string.h>
#include<stdlib.h>
#include<assert.h>
#include<stdarg.h>
#include"ntc.h"
#include"reporting.h"
#include"utils.h"
#include"x86/arch.h"

const char *AST_KIND_STR[] = { AST_KINDS(GEN_STRI) };

void generic_visitor(AST **nptr, AST *stmt, AST *stmtPrev, AST *chu, AST *tlc, void *ud, GenericVisitorHandler preHandler, GenericVisitorHandler postHandler) {
	if(preHandler) preHandler(nptr, stmt, stmtPrev, chu, tlc, ud);
	
	AST *n = *nptr;
	
	if(n->nodeKind == AST_CHUNK) {
		AST *sPrev = NULL;
		AST **s = &n->chunk.statementFirst;
		while(*s) {
			generic_visitor(s, *s, sPrev, n, tlc, ud, preHandler, postHandler);
			
			sPrev = *s;
			s = &sPrev->statement.next;
		}
	} else if(n->nodeKind == AST_STMT_ASSIGN) {
		generic_visitor(&n->stmtAssign.what, stmt, stmtPrev, chu, tlc, ud, preHandler, postHandler);
		if(n->stmtAssign.to) {
			generic_visitor(&n->stmtAssign.to, stmt, stmtPrev, chu, tlc, ud, preHandler, postHandler);
		}
	} else if(n->nodeKind == AST_STMT_IF) {
		generic_visitor(&n->stmtIf.expression, stmt, stmtPrev, chu, tlc, ud, preHandler, postHandler);
		generic_visitor(&n->stmtIf.then, stmt, stmtPrev, chu, tlc, ud, preHandler, postHandler);
		if(n->stmtIf.elss) {
			generic_visitor(&n->stmtIf.elss, stmt, stmtPrev, chu, tlc, ud, preHandler, postHandler);
		}
	} else if(n->nodeKind == AST_STMT_LOOP) {
		generic_visitor(&n->stmtLoop.body, stmt, stmtPrev, chu, tlc, ud, preHandler, postHandler);
	} else if(n->nodeKind == AST_STMT_BREAK) {
	} else if(n->nodeKind == AST_STMT_CONTINUE) {
	} else if(n->nodeKind == AST_STMT_EXT_ALIGN) {
	} else if(n->nodeKind == AST_STMT_DECL) {
		if(n->stmtDecl.expression) {
			generic_visitor(&n->stmtDecl.expression, stmt, stmtPrev, chu, tlc, ud, preHandler, postHandler);
		}
	} else if(n->nodeKind == AST_STMT_EXPR) {
		generic_visitor(&n->stmtExpr.expr, stmt, stmtPrev, chu, tlc, ud, preHandler, postHandler);
	} else if(n->nodeKind == AST_STMT_EXT_ORG) {
	} else if(n->nodeKind == AST_STMT_EXT_SECTION) {
	} else if(n->nodeKind == AST_STMT_RETURN) {
		if(n->stmtReturn.val) {
			generic_visitor(&n->stmtReturn.val, stmt, stmtPrev, chu, tlc, ud, preHandler, postHandler);
		}
	} else if(n->nodeKind == AST_EXPR_BINARY_OP) {
		generic_visitor(&n->exprBinOp.operands[0], stmt, stmtPrev, chu, tlc, ud, preHandler, postHandler);
		generic_visitor(&n->exprBinOp.operands[1], stmt, stmtPrev, chu, tlc, ud, preHandler, postHandler);
	} else if(n->nodeKind == AST_EXPR_CALL) {
		generic_visitor(&n->exprCall.what, stmt, stmtPrev, chu, tlc, ud, preHandler, postHandler);
		
		for(size_t i = 0; i < n->exprCall.what->expression.type->pointer.of->function.argCount; i++) {
			generic_visitor(&n->exprCall.args[i], stmt, stmtPrev, chu, tlc, ud, preHandler, postHandler);
		}
	} else if(n->nodeKind == AST_EXPR_CAST) {
		generic_visitor(&n->exprCast.what, stmt, stmtPrev, chu, tlc, ud, preHandler, postHandler);
	} else if(n->nodeKind == AST_EXPR_FUNC) {
		generic_visitor(&n->exprFunc.chunk, NULL, NULL, n->exprFunc.chunk, n->exprFunc.chunk, ud, preHandler, postHandler);
	} else if(n->nodeKind == AST_EXPR_UNARY_OP) {
		generic_visitor(&n->exprUnOp.operand, stmt, stmtPrev, chu, tlc, ud, preHandler, postHandler);
	} else if(n->nodeKind == AST_EXPR_VAR) {
	} else if(n->nodeKind == AST_EXPR_STACK_POINTER) {
	} else if(n->nodeKind == AST_EXPR_PRIMITIVE) {
	} else if(n->nodeKind == AST_EXPR_STRING_LITERAL) {
	} else if(n->nodeKind == AST_EXPR_ARRAY) {
		assert(n->expression.type->type == TYPE_TYPE_ARRAY);
		assert(n->expression.type->array.length != 0);
		
		for(size_t i = 0; i < n->expression.type->array.length; i++) {
			generic_visitor(&n->exprArray.items[i], stmt, stmtPrev, chu, tlc, ud, preHandler, postHandler);
		}
	} else if(n->nodeKind == AST_EXPR_EXT_SALLOC) {
	} else if(n->nodeKind == AST_EXPR_DOT) {
		generic_visitor(&n->exprDot.a, stmt, stmtPrev, chu, tlc, ud, preHandler, postHandler);
	} else if(n->nodeKind == AST_EXPR_EXT_SIZEOF) {
		if(n->exprExtSizeOf.ofExpr) {
			generic_visitor(&n->exprExtSizeOf.ofExpr, stmt, stmtPrev, chu, tlc, ud, preHandler, postHandler);
		}
	} else if(n->nodeKind == AST_STMT_JUMP) {
		if(n->stmtJump.condition) {
			generic_visitor(&n->stmtJump.condition, stmt, stmtPrev, chu, tlc, ud, preHandler, postHandler);
		}
	} else if(n->nodeKind == AST_STMT_LABEL) {
	} else if(n->nodeKind == AST_EXPR_NULL) {
	} else {
		stahp_node(n, "generic_visitor: unhandled %s", AST_KIND_STR[n->nodeKind]);
	}
	
	if(postHandler) postHandler(nptr, stmt, stmtPrev, chu, tlc, ud);
}

int ast_expression_equal(AST *a, AST *b) {
	if(a->nodeKind != b->nodeKind) return 0;
	
	if(a->nodeKind == AST_EXPR_PRIMITIVE) {
		return a->exprPrim.val == b->exprPrim.val;
	} else if(a->nodeKind == AST_EXPR_VAR) {
		return a->exprVar.thing == b->exprVar.thing;
	} else if(a->nodeKind == AST_EXPR_UNARY_OP) {
		return a->exprUnOp.operator == b->exprUnOp.operator && ast_expression_equal(a->exprUnOp.operand, b->exprUnOp.operand);
	} else if(a->nodeKind == AST_EXPR_BINARY_OP) {
		return a->exprBinOp.operator == b->exprBinOp.operator && ast_expression_equal(a->exprBinOp.operands[0], b->exprBinOp.operands[0]) && ast_expression_equal(a->exprBinOp.operands[1], b->exprBinOp.operands[1]);
	} else if(a->nodeKind == AST_EXPR_STACK_POINTER) {
		return 1;
	} else if(a->nodeKind == AST_EXPR_CAST) {
		return ast_expression_equal(a->exprCast.what, b->exprCast.what) && type_equal(a->exprCast.to, b->exprCast.to) && a->exprCast.reinterpretation == b->exprCast.reinterpretation;
	} else {
		stahp_node(a, "ast_expression_equal: unhandled %s", AST_KIND_STR[a->nodeKind]);
	}
}

// This function may return three values: YES (1), NO (0) or UNKNOWN (-1).
// ... Ew
int ast_stmt_is_after(const AST *chunk, const AST *s1, const AST *s2) {
	const AST *s = chunk->chunk.statementFirst;
	
	while(1) {
		if(s && s->nodeKind == AST_STMT_LOOP) {
			int i = ast_stmt_is_after(s->stmtLoop.body, s1, s2);
			if(i == 1 || (i == 0 && s1 != NULL)) {
				return i;
			}
		}
		
		if(s == s1) {
			return 0;
		}
		if(s == s2) {
			return 1;
		}
		
		if(!s) break;
		
		if(s->nodeKind == AST_STMT_IF) {
			int i = ast_stmt_is_after(s->stmtIf.then, s1, s2);
			if(i == 1 || (i == 0 && s1 != NULL)) {
				return i;
			}
			
			if(s->stmtIf.elss) {
				i = ast_stmt_is_after(s->stmtIf.elss, s1, s2);
				if(i == 1 || (i == 0 && s1 != NULL)) {
					return i;
				}
			}
		}
		
		s = s->statement.next;
	}
	
	return -1;
}

AST *ast_get_label_by_name(AST *tlc, const char *name) {
	for(AST *s = tlc->chunk.statementFirst; s; s = s->statement.next) {
		if(s->nodeKind == AST_STMT_LABEL && !strcmp(s->stmtLabel.name, name)) {
			return s;
		}
	}
	
	stahp(0, 0, "Label %s not found", name);
}

static void *memdup(void *a, size_t len) {
	void *r = malloc(len);
	memcpy(r, a, len);
	return r;
}

/*
 * WARNING: Just because you deep copy an AST node, does not mean
 * ast_expression_equal will return true! This matters for example with
 * function calls (a function call is not equal to itself).
 */
AST *ast_deep_copy(AST *src) {
	if(src->nodeKind == AST_EXPR_VAR) {
		return memdup(src, sizeof(ASTExprVar));
	} else if(src->nodeKind == AST_EXPR_PRIMITIVE) {
		return memdup(src, sizeof(ASTExprPrimitive));
	} else if(src->nodeKind == AST_EXPR_VAR) {
		return memdup(src, sizeof(ASTExprVar));
	} else if(src->nodeKind == AST_EXPR_UNARY_OP) {
		ASTExprUnaryOp *n = memdup(src, sizeof(ASTExprUnaryOp));
		n->operand = ast_deep_copy(n->operand);
		return n;
	}
	
	abort();
}

AST *ast_cast_expr(AST *what, Type *to) {
	if(what == NULL) {
		stahp(0, 0, "NULL what passed to ast_cast_expr");
		goto fail;
	}
	
	/* Only exists at parse-time, hence not part of type system and is handled separately */
	if(what->nodeKind == AST_EXPR_STRING_LITERAL) {
		if(to->type == TYPE_TYPE_ARRAY && type_equal(primitive_parse("u8"), to->array.of) && (to->array.length == what->exprStrLit.length || to->array.length == -1)) {
			ASTExprArray *ret = calloc(1, sizeof(*ret));
			ret->nodeKind = AST_EXPR_ARRAY;
			ret->items = calloc(what->exprStrLit.length, sizeof(*ret->items));
			ret->type = to;
			
			// If this declaration is of an unknown-length array type (eg u8[?]) then update the length
			if(to->array.length == -1) {
				to->array.length = what->exprStrLit.length;
			}
			
			for(int i = 0; i < what->exprStrLit.length; i++) {
				uint8_t bajt = what->exprStrLit.data[i];
				
				ASTExprPrimitive *item = calloc(1, sizeof(*item));
				item->nodeKind = AST_EXPR_PRIMITIVE;
				item->type = to->array.of;
				item->val = bajt;
				
				ret->items[i] = (AST*) item;
			}
			
			return (AST*) ret;
		} else if(to->type == TYPE_TYPE_PRIMITIVE) {
			if(to->primitive.width != what->exprStrLit.length * 8) {
				stahp_node(what, "Size mismatch between string literal and target type");
			}
			
			ASTExprPrimitive *ret = calloc(1, sizeof(*ret));
			ret->nodeKind = AST_EXPR_PRIMITIVE;
			ret->type = to;
			ret->val = 0;
			memcpy(&ret->val, what->exprStrLit.data, what->exprStrLit.length);
			return (AST*) ret;
		} else abort();
	}
	
	bool copy = false;
	
	// Make sure an unparametrized generic int parameter hasn't sneaked its way in
	while(what->nodeKind == AST_EXPR_VAR && what->exprVar.thing->kind == SCOPEITEM_CEXPR && what->exprVar.thing->data.cexpr.concrete) {
		what = what->exprVar.thing->data.cexpr.concrete;
		copy = true;
	}
	assert(!(what->nodeKind == AST_EXPR_VAR && what->exprVar.thing->kind == SCOPEITEM_CEXPR));
	
	if(copy) {
		what = ast_deep_copy(what);
	}
	
	if(type_equal(what->expression.type, to)) return what;
	
	if(!type_is_castable(what->expression.type, to)) {
		goto fail;
	}
	
	if(what->nodeKind == AST_EXPR_PRIMITIVE && (to->type == TYPE_TYPE_PRIMITIVE || to->type == TYPE_TYPE_POINTER)) {
		ASTExprPrimitive *ret = calloc(1, sizeof(*ret));
		ret->nodeKind = AST_EXPR_PRIMITIVE;
		ret->type = to;
		
		if(to->type == TYPE_TYPE_PRIMITIVE) {
			ret->val = what->exprPrim.val & ((1UL << to->primitive.width) - 1);
		} else {
			ret->val = what->exprPrim.val & ((1UL << (8 * type_size(to))) - 1);
		}
		
		return (AST*) ret;
	} else {
		ASTExprCast *ret = calloc(1, sizeof(*ret));
		ret->nodeKind = AST_EXPR_CAST;
		ret->type = to;
		ret->what = what;
		ret->to = to;
		return (AST*) ret;
	}
	
fail:
	stahp_node(what, "Cannot cast type %s into %s", type_to_string(what->expression.type), type_to_string(to));
}

struct ReferencesStackState {
	bool yes;
};
static void references_stack_visitor(AST **nptr, AST *stmt, AST *stmtPrev, AST *chunk, AST *tlc, void *ud) {
	AST *n = *nptr;
	if(n->nodeKind == AST_EXPR_STACK_POINTER) {
		((struct ReferencesStackState*) ud)->yes = true;
	}
}
bool ast_references_stack(AST *a) {
	struct ReferencesStackState state = {};
	generic_visitor(&a, NULL, NULL, NULL, NULL, &state, references_stack_visitor, NULL);
	return state.yes;
}

struct IsScopeItemReferenced {
	bool yes;
	ScopeItem *si;
};
static void is_scopeitem_referenced_visitor(AST **aptr, AST *stmt, AST *stmtPrev, AST *chunk, AST *tlc, void *ud) {
	struct IsScopeItemReferenced *state = ud;
	
	AST *n = *aptr;
	if(n->nodeKind == AST_EXPR_UNARY_OP) {
		if(n->exprUnOp.operand->nodeKind == AST_EXPR_VAR && n->exprUnOp.operator == UNOP_REF && n->exprUnOp.operand->exprVar.thing == state->si) {
			state->yes = true;
		}
	}
}
bool ast_is_scopeitem_referenced(AST *tlc, ScopeItem *si) {
	struct IsScopeItemReferenced state = {.si = si};
	generic_visitor(&tlc, NULL, NULL, tlc, tlc, &state, NULL, is_scopeitem_referenced_visitor);
	return state.yes;
}

ScopeItem *ast_tlc_new_var(AST *tlc, char *name, Type *itstype) {
	ScopeItem *vte = calloc(1, sizeof(*vte));
	vte->kind = SCOPEITEM_VAR;
	vte->type = itstype;
	vte->data.var.color = -1;
	vte->data.var.precolored = false;
	vte->data.var.degree = 0;
	vte->data.var.priority = 0;
	vte->data.var.name = name;
	
	// Add to var array
	tlc->chunk.vars = realloc(tlc->chunk.vars, sizeof(*tlc->chunk.vars) * (++tlc->chunk.varCount));
	tlc->chunk.vars[tlc->chunk.varCount - 1] = vte;
	
	return vte;
}