package main

import "core:fmt"

Parser :: struct {
	lexer:           ^Lexer,
	tok, next:       Token,
	can_be_function: bool,
}

parser_create :: proc(lexer: ^Lexer) -> (ret: Parser) {
	ret = {
		lexer           = lexer,
		can_be_function = true,
	}
	parser_next(&ret)
	parser_next(&ret)
	return
}

@(private = "file")
parser_next :: proc(parser: ^Parser) {
	parser.tok = parser.next
	parser.next = lexer_next(parser.lexer)
}

@(private = "file")
accept :: proc(parser: ^Parser, tok: TokenKind) -> bool {
	if parser.tok.kind == tok {
		parser_next(parser)
		return true
	}
	return false
}

@(private = "file")
expect :: proc(parser: ^Parser, tok: TokenKind) -> bool {
	if !accept(parser, tok) {
		append(
			&g_message_list,
			message_create(
				.Error,
				fmt.aprintf("Expected {}, got {} at {}", tok, parser.tok.kind, "TODO"),
				parser.tok.range,
			),
		)
		return false
	}
	return true
}

parser_parse :: proc(parser: ^Parser) -> (ret: ^Node) {
	ret = parser_parse_block(parser, .EOF)
	return
}

@(private = "file")
parser_parse_block :: proc(parser: ^Parser, end: TokenKind) -> (ret: ^Node) {
	range := parser.tok.range
	statements: [dynamic]^Node
	for parser.tok.kind != end && parser.tok.kind != .EOF {
		if accept(parser, .Let) {
			ret := parser_parse_definitions(parser)
			expect(parser, .Semicolon)
			if ret != nil {
				for stmt in ret {
					if stmt != nil {
						append(&statements, stmt)
					}
				}
			}
		} else {
			stmt := parser_parse_statement(parser)
			if stmt != nil {
				append(&statements, stmt)
			}
		}
	}
	expect(parser, end)
	return node_create_block(range, statements)
}

@(private = "file")
parser_parse_statement :: proc(parser: ^Parser) -> (ret: ^Node) {
	range_beg := parser.tok.range
	if accept(parser, .Semicolon) {
		ret = nil
		return
	} else if parser.tok.kind == .If {
		expect(parser, .If)
		ret = parser_parse_if_statement(parser)
	} else if parser.tok.kind == .Use {
		ret = parser_parse_use_statement(parser)
		expect(parser, .Semicolon)
	} else if parser.tok.kind == .OpenBrace {
		parser_next(parser)
		ret = parser_parse_block(parser, .CloseBrace)
		expect(parser, .CloseBrace)
	} else if parser.tok.kind == .For {
		ret = parser_parse_for_statement(parser)
	} else if parser.tok.kind == .Function {
		ret = parser_parse_function_definition(parser)
	} else if parser.tok.kind == .Struct {
		ret = parser_parse_struct_definition(parser)
	} else if parser.tok.kind == .Enum {
		ret = parser_parse_enum_definition(parser)
	} else if parser.tok.kind == .Union {
		ret = parser_parse_union_definition(parser)
	} else if parser.tok.kind == .Ret {
		ret = parser_parse_ret(parser)
	} else {
		ret = parser_parse_expression(parser)
	}
	if ret != nil {
		ret.range.range.start = range_beg.range.start
	}
	expect(parser, .Semicolon)
	return
}

@(private = "file")
parser_parse_ret :: proc(parser: ^Parser) -> ^Node {
	range := parser.tok.range
	expect(parser, .Ret)
	if parser.tok.kind == .Semicolon {
		return node_create_ret(range, nil)
	}
	ret := parser_parse_expression(parser)
	return node_create_ret(range, ret)
}

@(private = "file")
parser_parse_struct_definition :: proc(parser: ^Parser) -> ^Node {
	range := parser.tok.range
	expect(parser, .Struct)
	name: [dynamic]u8
	if parser.tok.kind == .Identifier {
		name = parser.tok.value.([dynamic]u8)
		parser_next(parser)
	} else {
		expect(parser, .Identifier)
		return nil
	}
	is_packed := true
	if parser.tok.kind == .String {
		value := parser.tok.value.([dynamic]u8)
		if compare_dyn_arr_string(&value, "unpacked") {
			is_packed = false
		} else {
			panic("TODO, unknown struct attribute")
		}
		parser_next(parser)
	}
	expect(parser, .OpenBrace)
	fields := parser_parse_definitions(parser, .CloseBrace)
	expect(parser, .CloseBrace)
	n := node_create_struct_or_union(range, .Struct, name, fields)
	n.value = is_packed
	return n
}

@(private = "file")
enum_resolve_value_items :: proc(values: ^[dynamic]EnumValueItem) {
	// IMPORTANT: Although this should be obvious, THIS DOES NOT DO ANY TYPE
	//            CHECKING! The `type_check` procedure is for that!

	// What this algorithm basically does is find the first resolved item, then
	// it resolved backwards, afterwards, it goes and resolves stuff forwards.
	// This resolves everything as signed integers which may be an issue but i
	// frankly don't care, if you have an enum which has more 2^63 negative
	// values that's on you.

	// 1. Find the first element that is resolved, default to the beginning if
	//    none is found
	idx_first_resolved := 0
	for &value, idx in values {
		if value.resolved {
			idx_first_resolved = idx
		}
	}

	// 2. Resolve backwards from first resolved element
	for i := idx_first_resolved; i >= 1; i -= 1 {
		values[i - 1].value = values[i].value - 1
		values[i].resolved = true
		values[i - 1].resolved = true
	}

	// 3. Resolve forwards until the end of the values vector
	for i in idx_first_resolved ..< len(values) - 1 {
		values[i + 1].value = values[i].value + 1
		values[i].resolved = true
		values[i + 1].resolved = true
	}
}

@(private = "file")
parser_parse_enum_definition :: proc(parser: ^Parser) -> ^Node {
	range := parser.tok.range
	expect(parser, .Enum)

	enum_value := EnumValue {
		values = [dynamic]EnumValueItem{},
		type   = nil,
	}

	name: [dynamic]u8
	if parser.tok.kind == .Identifier {
		name = parser.tok.value.([dynamic]u8)
	}
	if !expect(parser, .Identifier) {
		if accept(parser, .Colon) {
			for parser.tok.kind != .OpenBrace && parser.tok.kind != .EOF && parser.tok.kind != .Semicolon {
				parser_next(parser)
			}
		}
		if accept(parser, .OpenBrace) {
			for parser.tok.kind != .CloseBrace && parser.tok.kind != .EOF {
				parser_next(parser)
			}
		}
		return nil
	}

	type_node: ^Node = nil
	if accept(parser, .Colon) {
		type_node = parser_parse_type(parser)
	} else {
		type_node = node_create_value(.Identifier, range, [dynamic]u8{'i', '3', '2'})
	}

	expect(parser, .OpenBrace)

	for parser.tok.kind != .CloseBrace && parser.tok.kind != .EOF {
		item_range := parser.tok.range
		item := EnumValueItem{}
		item.range = item_range
		item.resolved = false
		if parser.tok.kind != .Identifier {
			append(&g_message_list, message_create(.Error, "Expected identifier in enum value list", parser.tok.range))
			for parser.tok.kind != .CloseBrace && parser.tok.kind != .EOF {
				parser_next(parser)
			}
			return nil
		}
		name := parser.tok.value.([dynamic]u8)
		item.name = name
		parser_next(parser)

		value := transmute(u64)(i64(-1))
		if accept(parser, .Assign) {
			if parser.tok.kind != .Integer {
				append(
					&g_message_list,
					message_create(.Error, "Only integers can be values for enums", parser.tok.range),
				)
				continue
			}
			item.resolved = true
			value = parser.tok.value.(u64)
			item_range.range.end = parser.tok.range.range.end
			parser_next(parser)
		}
		item.value = value

		append(&enum_value.values, item)

		if parser.tok.kind == .CloseBrace {
			continue
		}
		expect(parser, .Comma)
	}
	range.range.end = parser.tok.range.range.end
	parser_next(parser)

	enum_resolve_value_items(&enum_value.values)

	enum_value.type = type_node

	return node_create_enum(range, enum_value, name)
}

@(private = "file")
parser_parse_union_definition :: proc(parser: ^Parser) -> ^Node {
	range := parser.tok.range
	expect(parser, .Union)
	name: [dynamic]u8
	if parser.tok.kind == .Identifier {
		name = parser.tok.value.([dynamic]u8)
		parser_next(parser)
	} else {
		expect(parser, .Identifier)
	}
	expect(parser, .OpenBrace)
	fields := parser_parse_definitions(parser, .CloseBrace)
	expect(parser, .CloseBrace)
	return node_create_struct_or_union(range, .Union, name, fields)
}

@(private = "file")
parser_parse_function_definition :: proc(parser: ^Parser) -> ^Node {
	expect(parser, .Function)
	name: [dynamic]u8
	if parser.tok.kind == .Identifier {
		name = parser.tok.value.([dynamic]u8)
		parser_next(parser)
	} else {
		expect(parser, .Identifier)
	}

	params: [dynamic]^Node
	if accept(parser, .OpenParen) {
		params = parser_parse_definitions(parser, .CloseParen)
		expect(parser, .CloseParen)
	} else {
		params = {}
	}

	type: ^Node = nil
	if parser.tok.kind != .OpenBrace && parser.tok.kind != .Semicolon {
		type = parser_parse_type(parser)
	}
	if parser.tok.kind == .Semicolon {
		return node_create_extern_function(parser.tok.range, name, type, params)
	}
	expect(parser, .OpenBrace)
	body := parser_parse_block(parser, .CloseBrace)
	return node_create_function(parser.tok.range, name, type, body, params)
}

@(private = "file")
parser_parse_for_statement :: proc(parser: ^Parser) -> ^Node {
	range := parser.tok.range
	expect(parser, .For)
	if accept(parser, .OpenBrace) {
		body := parser_parse_block(parser, .CloseBrace)
		return node_create_for(range, nil, nil, nil, body)
	}
	if parser.tok.kind == .Let {
		panic("TODO, let in for not implemented yet")
	}
	init := parser_parse_expression(parser)
	if accept(parser, .OpenBrace) {
		body := parser_parse_block(parser, .CloseBrace)
		return node_create_for(range, nil, init, nil, body)
	}
	expect(parser, .Semicolon)
	condition: ^Node = nil
	if parser.tok.kind != .Semicolon {
		condition = parser_parse_expression(parser)
	}
	expect(parser, .Semicolon)
	if accept(parser, .OpenBrace) {
		body := parser_parse_block(parser, .CloseBrace)
		return node_create_for(range, init, condition, nil, body)
	}
	after := parser_parse_expression(parser)
	expect(parser, .OpenBrace)
	body := parser_parse_block(parser, .CloseBrace)
	return node_create_for(range, init, condition, after, body)
}

@(private = "file")
parser_parse_type :: proc(parser: ^Parser) -> (ret: ^Node) {
	range := parser.tok.range
	if parser.tok.kind == .Identifier {
		ret = node_create_value(.Identifier, range, parser.tok.value.([dynamic]u8))
		parser_next(parser)
	} else if accept(parser, .BitwiseXOR) {
		ret = node_create_pointer(range, parser_parse_type(parser))
	} else if accept(parser, .OpenBracket) {
		size: u64 = 0
		if parser.tok.kind == .Integer {
			size = parser.tok.value.(u64)
			parser_next(parser)
		}

		expect(parser, .CloseBracket)
		ret = node_create_array(range, size, parser_parse_type(parser))
	} else {
		append(
			&g_message_list,
			message_create(
				.Error,
				fmt.aprintf("Expected type, got {} at {}", parser.tok.kind, "TODO"),
				parser.tok.range,
			),
		)
		ret = nil
	}
	return
}

@(private = "file")
parser_parse_definitions :: proc(parser: ^Parser, end := TokenKind.Semicolon) -> [dynamic]^Node {
	range := parser.tok.range
	vars: [dynamic]^Node
	type: ^Node = nil
	are_constants := false
	uninitialized := false
	if parser.tok.kind != .Identifier {
		append(
			&g_message_list,
			message_create(.Error, "Built-in keywords cannot be used for variable names", parser.tok.range),
		)
		for parser.tok.kind != end && parser.tok.kind != .EOF {
			parser_next(parser)
		}
		return nil
	}

	for parser.tok.kind != end && parser.tok.kind != .EOF {
		names := [dynamic][dynamic]u8{}

		for parser.tok.kind == .Identifier {
			tok := parser.tok
			if !expect(parser, .Identifier) {
				break
			}

			append(&names, tok.value.([dynamic]u8))
		}

		expect(parser, .Colon)

		if parser.tok.kind != .Colon && parser.tok.kind != .Assign {
			type = parser_parse_type(parser)
		}

		if accept(parser, .Assign) {
			are_constants = false
		} else if accept(parser, .Colon) {
			are_constants = true
		} else {
			uninitialized = true
		}

		for i in 0 ..< len(names) {
			value: ^Node = nil
			if uninitialized == false {
				value = parser_parse_expression(parser)
			}
			name_node := node_create_value(.Identifier, range, names[i])
			append(&vars, node_create_variable(range, name_node, type, value, are_constants))
		}

		if !accept(parser, .Comma) {
			break
		}

		type = nil
		are_constants = false
		uninitialized = false
	}
	return vars
}

@(private = "file")
parser_parse_use_statement :: proc(parser: ^Parser) -> ^Node {
	range := parser.tok.range
	expect(parser, .Use)
	alias: [dynamic]u8
	if parser.tok.kind == .Identifier {
		alias = parser.tok.value.([dynamic]u8)
		parser_next(parser)
	}
	tok := parser.tok
	if !expect(parser, .String) {
		return nil
	}
	return node_create_use(range, tok.value.([dynamic]u8), alias)
}

@(private = "file")
parser_parse_if_statement :: proc(parser: ^Parser) -> ^Node {
	range := parser.tok.range
	condition := parser_parse_expression(parser)
	expect(parser, .OpenBrace)
	body := parser_parse_block(parser, .CloseBrace)
	if accept(parser, .Elif) {
		falsebody := parser_parse_if_statement(parser)
		return node_create_if(range, condition, body, falsebody)
	}
	if accept(parser, .Else) {
		expect(parser, .OpenBrace)
		falsebody := parser_parse_block(parser, .CloseBrace)
		return node_create_if(range, condition, body, falsebody)
	}
	return node_create_if(range, condition, body, nil)
}

@(private = "file")
parser_parse_expression :: proc(parser: ^Parser) -> ^Node {
	return parser_parse_assignment(parser)
}

@(private = "file")
parser_parse_binary_expression :: proc(
	parser: ^Parser,
	kinds: []TokenKind,
	next: proc(parser: ^Parser) -> ^Node,
	right_assoc := false,
) -> ^Node {
	lhs := next(parser)
	i := 0
	for i < len(kinds) {
		kind := kinds[i]
		for accept(parser, kind) {
			i = 0
			prev_can_be_function := parser.can_be_function
			parser.can_be_function = false
			rhs: ^Node
			if !right_assoc {
				rhs = next(parser)
			} else {
				rhs = parser_parse_binary_expression(parser, kinds, next, right_assoc)
			}

			lhs = node_create_binary(kind, lhs.range, lhs, rhs)
			lhs^.range.range.end = rhs.range.range.end
			parser.can_be_function = prev_can_be_function
		}
		i += 1
	}

	return lhs
}

@(private = "file")
parser_parse_assignment :: proc(parser: ^Parser) -> ^Node {
	return parser_parse_binary_expression(parser, {.Assign}, parser_parse_arrow, true)
}

@(private = "file")
parser_parse_arrow :: proc(parser: ^Parser) -> ^Node {
	// Basically, a -> b is the same as function_call(b, {a})
	lhs := parser_parse_equality(parser)
	for accept(parser, .Arrow) {
		rhs := parser_parse_assignment(parser)
		if rhs.kind != .FunctionCall &&
		   rhs.kind != .Identifier &&
		   rhs.kind != .FieldAccess &&
		   rhs.kind != .IndexAccess {
			append(
				&g_message_list,
				message_create(
					.Error,
					fmt.aprintf("Expected function call, got {} at {}", rhs.kind, "TODO"),
					rhs.range,
				),
			)
			return lhs
		}
		if rhs.kind != .FunctionCall {
			rhs = node_create_function_call(rhs.range, rhs, nil)
		}
		inject_at(&rhs.children, 1, lhs)
		lhs = rhs
	}
	return lhs
}

@(private = "file")
parser_parse_equality :: proc(parser: ^Parser) -> ^Node {
	return parser_parse_binary_expression(parser, {.Equals, .NotEquals}, parser_parse_comparison)
}

@(private = "file")
parser_parse_comparison :: proc(parser: ^Parser) -> ^Node {
	return parser_parse_binary_expression(
		parser,
		{.LessThan, .LessThanOrEqual, .GreaterThan, .GreaterThanOrEqual},
		parser_parse_addition,
	)
}

@(private = "file")
parser_parse_addition :: proc(parser: ^Parser) -> ^Node {
	return parser_parse_binary_expression(parser, {.Add, .Subtract}, parser_parse_multiplication)
}

@(private = "file")
parser_parse_multiplication :: proc(parser: ^Parser) -> ^Node {
	return parser_parse_binary_expression(parser, {.Multiply, .Divide, .Modulo}, parser_parse_exponent)
}

@(private = "file")
parser_parse_exponent :: proc(parser: ^Parser) -> ^Node {
	return parser_parse_binary_expression(parser, {.Exponent}, parser_parse_bitwise)
}

@(private = "file")
parser_parse_bitwise :: proc(parser: ^Parser) -> ^Node {
	return parser_parse_binary_expression(parser, {.BitwiseAnd, .BitwiseOr, .BitwiseXOR}, parser_parse_bitwise1)
}

@(private = "file")
parser_parse_bitwise1 :: proc(parser: ^Parser) -> ^Node {
	return parser_parse_binary_expression(parser, {.BitwiseLeftShift, .BitwiseRightShift}, parser_parse_prefix_2)
}

@(private = "file")
parser_parse_prefix_2 :: proc(parser: ^Parser) -> ^Node {
	range := parser.tok.range
	if accept(parser, .Not) {
		rhs := parser_parse_suffix(parser)
		range.range.end = rhs.range.range.end
		return node_create_unary(.Not, range, rhs)
	} else if accept(parser, .BitwiseNot) {
		rhs := parser_parse_suffix(parser)
		range.range.end = rhs.range.range.end
		return node_create_unary(.BitwiseNot, range, rhs)
	} else if accept(parser, .Increment) {
		rhs := parser_parse_suffix(parser)
		range.range.end = rhs.range.range.end
		return node_create_unary(.Increment, range, rhs)
	} else if accept(parser, .Decrement) {
		rhs := parser_parse_suffix(parser)
		range.range.end = rhs.range.range.end
		return node_create_unary(.Decrement, range, rhs)
	} else if accept(parser, .BitwiseXOR) {
		rhs := parser_parse_suffix(parser)
		range.range.end = rhs.range.range.end
		return node_create_unary(.BitwiseXOR, range, rhs)
	}
	return parser_parse_suffix(parser)
}

@(private = "file")
parser_parse_suffix :: proc(parser: ^Parser) -> ^Node {
	range := parser.tok.range
	lhs := parser_parse_prefix(parser)
	range_op := parser.tok.range
	range.range.end = range_op.range.end
	if accept(parser, .OpenBracket) {
		rhs := parser_parse_expression(parser)
		range.range.end = rhs.range.range.end
		expect(parser, .CloseBracket)
		return node_create_index_access(range, lhs, rhs)
	} else if accept(parser, .Increment) {
		return node_create_unary(.Increment, range, lhs, true)
	} else if accept(parser, .Decrement) {
		return node_create_unary(.Decrement, range, lhs, true)
	} else if accept(parser, .As) {
		type := parser_parse_type(parser)
		range.range.end = type.range.range.end
		return node_create_cast(range, lhs, type)
	} else if accept(parser, .BitwiseAs) {
		type := parser_parse_type(parser)
		range.range.end = type.range.range.end
		return node_create_bitwise_cast(range, lhs, type)
	}
	return lhs
}

@(private = "file")
parser_parse_prefix :: proc(parser: ^Parser) -> ^Node {
	range := parser.tok.range
	if accept(parser, .Add) {
		return node_create_unary(.Add, range, parser_parse_prefix(parser))
	} else if accept(parser, .Subtract) {
		return node_create_unary(.Subtract, range, parser_parse_prefix(parser))
	}
	return parser_parse_factor(parser)
}

@(private = "file")
parser_parse_expression_list :: proc(parser: ^Parser, end: TokenKind = .CloseBrace) -> [dynamic]^Node {
	ret := [dynamic]^Node{}
	for parser.tok.kind != end && parser.tok.kind != .EOF {
		append(&ret, parser_parse_expression(parser))
	}
	return ret
}

@(private = "file")
parser_is_factor_token_or_prefix :: proc(tt: TokenKind) -> bool {
	#partial switch tt {
	case .Identifier, .Integer, .Float, .String, .Character, .OpenParen:
		return true
	case:
		return false
	}
}

@(private = "file")
parser_parse_factor :: proc(parser: ^Parser) -> (ret: ^Node) {
	ret = nil
	if parser.tok.kind == .Dot && parser.next.kind == .Identifier {
		range := parser.tok.range
		parser_next(parser)
		name := parser.tok.value.([dynamic]u8)
		parser_next(parser)

		expect(parser, .OpenBrace)
		args := parser_parse_expression_list(parser)
		range.range.end = parser.tok.range.range.end
		expect(parser, .CloseBrace)
		ret = node_create_struct_initializer(range, name, args)
	} else if parser.tok.kind == .Integer {
		ret = node_create_value(.Integer, parser.tok.range, parser.tok.value)
		parser_next(parser)
	} else if parser.tok.kind == .Float {
		ret = node_create_value(.Float, parser.tok.range, parser.tok.value)
		parser_next(parser)
	} else if parser.tok.kind == .Character {
		ret = node_create_value(.Character, parser.tok.range, parser.tok.value)
		parser_next(parser)
	} else if parser.tok.kind == .String {
		ret = node_create_value(.String, parser.tok.range, parser.tok.value)
		parser_next(parser)
	} else if parser.tok.kind == .Identifier {
		ret = node_create_value(.Identifier, parser.tok.range, parser.tok.value)
		parser_next(parser)
		prev := parser.can_be_function
		parser.can_be_function = false

		for accept(parser, .Dot) {
			if parser.tok.kind != .Identifier {
				append(
					&g_message_list,
					message_create(
						.Error,
						fmt.aprintf("Expected identifier after '.' but got {}", parser.tok.kind),
						parser.tok.range,
					),
				)
				continue
			}

			right := node_create_value(.Identifier, parser.tok.range, parser.tok.value)
			parser_next(parser)

			ret = node_create_field_access(
				SourceLocation{{ret.range.range.start, right.range.range.end}, ret.range.file},
				ret,
				right,
			)
		}

		parser.can_be_function = prev
		if parser.can_be_function &&
		   parser.tok.kind != .CloseParen &&
		   parser.tok.kind != .Semicolon &&
		   parser.tok.kind != .Arrow &&
		   parser.tok.kind != .Assign &&
		   parser.tok.kind != .EOF {
			prev := parser.can_be_function
			parser.can_be_function = false
			args: [dynamic]^Node
			for parser.tok.kind != .CloseParen &&
			    parser.tok.kind != .Semicolon &&
			    parser.tok.kind != .Arrow &&
			    parser.tok.kind != .EOF &&
			    parser.tok.kind != .Assign &&
			    parser_is_factor_token_or_prefix(parser.tok.kind) {
				append(&args, parser_parse_expression(parser))
			}
			ret = node_create_function_call(ret.range, ret, args)
			parser.can_be_function = prev
		}
	} else if accept(parser, .OpenParen) {
		prev := parser.can_be_function
		parser.can_be_function = true
		ret = parser_parse_expression(parser)
		parser.can_be_function = prev
		expect(parser, .CloseParen)
	} else {
		append(
			&g_message_list,
			message_create(
				.Error,
				fmt.aprintf("Unexpected factor token {} at {}", parser.tok.kind, "TODO"),
				parser.tok.range,
			),
		)
		parser_next(parser)
	}
	return
}
