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)
			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
	} 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 {
		ret = parser_parse_block(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 {
		ret = parser_parse_expression(parser)
		expect(parser, .Semicolon)
	}
	if ret != nil {
		ret.range.start = range_beg.start
	}
	return
}

@(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)
	}
	expect(parser, .OpenBrace)
	fields := parser_parse_definitions(parser, .CloseBrace)
	expect(parser, .CloseBrace)
	return node_create_struct_enum_or_union(range, .Struct, name, fields)
}

@(private = "file")
parser_parse_enum_definition :: proc(parser: ^Parser) -> ^Node {
	range := parser.tok.range
	expect(parser, .Enum)
	panic("TODO, enum not implemented yet")
}

@(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_enum_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 {
		type = parser_parse_type(parser)
	}
	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) {
	// FIXME: Add more types
	range := parser.tok.range
	if parser.tok.kind == .Identifier {
		ret = node_create_value(.Identifier, range, parser.tok.value.([dynamic]u8))
		parser_next(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
	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,
) -> ^Node {
	lhs := next(parser)
	for kind in kinds {
		for accept(parser, kind) {
			rhs := next(parser)
			lhs = node_create_binary(kind, lhs.range, lhs, rhs)
			lhs^.range.end = rhs.range.end
		}
	}
	return lhs
}

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

@(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_prefix_2)
}

@(private = "file")
parser_parse_bitwise :: proc(parser: ^Parser) -> ^Node {
	return parser_parse_binary_expression(
		parser,
		{.BitwiseAnd, .BitwiseOr, .BitwiseXOR, .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.end = rhs.range.end
		return node_create_unary(.Not, range, rhs)
	} else if accept(parser, .BitwiseNot) {
		rhs := parser_parse_suffix(parser)
		range.end = rhs.range.end
		return node_create_unary(.BitwiseNot, range, rhs)
	} else if accept(parser, .Increment) {
		rhs := parser_parse_suffix(parser)
		range.end = rhs.range.end
		return node_create_unary(.Increment, range, rhs)
	} else if accept(parser, .Decrement) {
		rhs := parser_parse_suffix(parser)
		range.end = rhs.range.end
		return node_create_unary(.Decrement, range, rhs)
	} else if accept(parser, .BitwiseXOR) {
		rhs := parser_parse_suffix(parser)
		range.end = rhs.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.end = range_op.end
	if accept(parser, .OpenBracket) {
		rhs := parser_parse_expression(parser)
		range.end = rhs.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)
	} else if accept(parser, .Decrement) {
		return node_create_unary(.Decrement, range, lhs)
	}
	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_is_factor_token_or_prefix :: proc(tt: TokenKind) -> bool {
	#partial switch tt {
	case .Identifier, .Integer, .Float, .String, .Character, .OpenParen:
		return true
	case .Not, .Subtract, .BitwiseNot, .Increment, .Decrement:
		// Assuming these can be used as prefixes in unary operations
		return true
	case:
		return false
	}
}

@(private = "file")
parser_parse_factor :: proc(parser: ^Parser) -> (ret: ^Node) {
	ret = nil
	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
		if accept(parser, .Dot) {
			ret = node_create_field_access(
				{ret.range.start, parser.tok.range.start},
				ret,
				parser_parse_factor(parser),
			)
		}
		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 != .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_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,
			),
		)
	}
	return
}