package parser

import "../ast"
import "../source"
import "../token"
import "core:strconv"

Parser :: struct {
	tokens:      ^token.Stream,
	diagnostics: ^source.Diagnostics,
	module:      ast.Module,
	cursor:      int,
	delimiter_depth: int,
}

current :: proc(parser: ^Parser) -> token.Token {
	return parser.tokens.items[min(parser.cursor, len(parser.tokens.items)-1)]
}

previous :: proc(parser: ^Parser) -> token.Token {
	return parser.tokens.items[max(parser.cursor-1, 0)]
}

advance :: proc(parser: ^Parser) -> token.Token {
	result := current(parser)
	if result.kind != .Eof {
		parser.cursor += 1
	}
	return result
}

allow :: proc(parser: ^Parser, kind: token.Kind) -> (token.Token, bool) {
	if current(parser).kind == kind {
		return advance(parser), true
	}
	return current(parser), false
}

skip_newlines :: proc(parser: ^Parser) {
	for current(parser).kind == .Newline {
		advance(parser)
	}
}

add_expr :: proc(parser: ^Parser, expr: ast.Expr) -> int {
	id := len(parser.module.exprs)
	append(&parser.module.exprs, expr)
	return id
}

invalid_expr :: proc(parser: ^Parser, span: source.Span, message: string) -> int {
	id := source.add(parser.diagnostics, span, message)
	return add_expr(parser, ast.Expr{
		kind=.Invalid,
		span=span,
		left=ast.INVALID_ID,
		right=ast.INVALID_ID,
		diagnostic=id,
	})
}

is_type_token :: proc(kind: token.Kind) -> bool {
	#partial switch kind {
	case .Keyword_Int, .Keyword_I8, .Keyword_I16, .Keyword_I32, .Keyword_I64, .Keyword_Void:
		return true
	}
	return false
}

parse_type :: proc(parser: ^Parser) -> ast.Type_Syntax {
	tok := current(parser)
	#partial switch tok.kind {
	case .Keyword_Int:
		advance(parser)
		return .Int
	case .Keyword_I8:
		advance(parser)
		return .I8
	case .Keyword_I16:
		advance(parser)
		return .I16
	case .Keyword_I32:
		advance(parser)
		return .I32
	case .Keyword_I64:
		advance(parser)
		return .I64
	case .Keyword_Void:
		advance(parser)
		return .Void
	}
	source.add(parser.diagnostics, tok.span, "expected a type")
	return .Invalid
}

parse_call :: proc(parser: ^Parser, name: token.Token) -> int {
	left_paren := advance(parser)
	parser.delimiter_depth += 1
	defer parser.delimiter_depth -= 1
	args: [dynamic]int
	args.allocator = parser.module.allocator
	skip_newlines(parser)
	for current(parser).kind != .Right_Paren && current(parser).kind != .Eof {
		append(&args, parse_expression(parser))
		skip_newlines(parser)
		if _, ok := allow(parser, .Comma); ok {
			skip_newlines(parser)
			continue
		}
		break
	}
	right_paren, ok := allow(parser, .Right_Paren)
	if !ok {
		source.add(parser.diagnostics, current(parser).span, "expected ')' after call arguments")
		right_paren = left_paren
	}
	return add_expr(parser, ast.Expr{
		kind=.Call,
		span=source.Span{start=name.span.start, end=right_paren.span.end},
		text=name.text,
		args=args[:],
		left=ast.INVALID_ID,
		right=ast.INVALID_ID,
		diagnostic=-1,
	})
}

parse_primary :: proc(parser: ^Parser) -> int {
	tok := current(parser)
	#partial switch tok.kind {
	case .Integer:
		advance(parser)
		value, ok := strconv.parse_i64(tok.text)
		if !ok {
			return invalid_expr(parser, tok.span, "integer literal does not fit in i64")
		}
		return add_expr(parser, ast.Expr{
			kind=.Integer,
			span=tok.span,
			integer=value,
			left=ast.INVALID_ID,
			right=ast.INVALID_ID,
			diagnostic=-1,
		})
	case .Identifier:
		advance(parser)
		if current(parser).kind == .Left_Paren {
			return parse_call(parser, tok)
		}
		return add_expr(parser, ast.Expr{
			kind=.Name,
			span=tok.span,
			text=tok.text,
			left=ast.INVALID_ID,
			right=ast.INVALID_ID,
			diagnostic=-1,
		})
	case .Underscore:
		advance(parser)
		return invalid_expr(parser, tok.span, "'_' is a write-only sink and cannot be read")
	case .Left_Paren:
		advance(parser)
		parser.delimiter_depth += 1
		defer parser.delimiter_depth -= 1
		skip_newlines(parser)
		expr := parse_expression(parser)
		skip_newlines(parser)
		if _, ok := allow(parser, .Right_Paren); !ok {
			source.add(parser.diagnostics, current(parser).span, "expected ')'")
		}
		return expr
	case .Invalid:
		advance(parser)
		return add_expr(parser, ast.Expr{
			kind=.Invalid,
			span=tok.span,
			left=ast.INVALID_ID,
			right=ast.INVALID_ID,
			diagnostic=tok.diagnostic,
		})
	}
	if tok.kind != .Newline && tok.kind != .Right_Brace && tok.kind != .Eof {
		advance(parser)
	}
	return invalid_expr(parser, tok.span, "expected an expression")
}

parse_expression :: proc(parser: ^Parser) -> int {
	left := parse_primary(parser)
	if parser.delimiter_depth > 0 {
		skip_newlines(parser)
	}
	for current(parser).kind == .Plus {
		advance(parser)
		skip_newlines(parser)
		right := parse_primary(parser)
		left_expr := parser.module.exprs[left]
		right_expr := parser.module.exprs[right]
		left = add_expr(parser, ast.Expr{
			kind=.Add,
			span=source.Span{start=left_expr.span.start, end=right_expr.span.end},
			left=left,
			right=right,
			diagnostic=-1,
		})
		if parser.delimiter_depth > 0 {
			skip_newlines(parser)
		}
	}
	return left
}

finish_statement :: proc(parser: ^Parser) -> int {
	if current(parser).kind == .Newline {
		skip_newlines(parser)
		return -1
	}
	if current(parser).kind == .Eof {
		return -1
	}
	diagnostic := source.add(
		parser.diagnostics,
		current(parser).span,
		"completed statements must be followed by a newline",
	)
	for current(parser).kind != .Newline &&
	    current(parser).kind != .Right_Brace &&
	    current(parser).kind != .Eof {
		advance(parser)
	}
	skip_newlines(parser)
	return diagnostic
}

parse_return :: proc(parser: ^Parser) -> int {
	start := advance(parser)
	skip_newlines(parser)
	if current(parser).kind == .Underscore {
		end := advance(parser)
		id := len(parser.module.statements)
		append(&parser.module.statements, ast.Stmt{
			kind=.Return,
			span=source.Span{start=start.span.start, end=end.span.end},
			name="_",
			expr=ast.INVALID_ID,
			diagnostic=-1,
		})
		return id
	}
	expr := parse_expression(parser)
	id := len(parser.module.statements)
	append(&parser.module.statements, ast.Stmt{
		kind=.Return,
		span=source.Span{start=start.span.start, end=parser.module.exprs[expr].span.end},
		expr=expr,
		diagnostic=-1,
	})
	return id
}

parse_statement :: proc(parser: ^Parser) -> int {
	if current(parser).kind == .Keyword_Return {
		return parse_return(parser)
	}

	if current(parser).kind == .Identifier || current(parser).kind == .Underscore {
		start_cursor := parser.cursor
		name := advance(parser)
		type_syntax := ast.Type_Syntax.Invalid
		had_type := false
		if is_type_token(current(parser).kind) {
			type_syntax = parse_type(parser)
			had_type = true
		}
		operator := current(parser)
		if operator.kind == .Colon_Colon || operator.kind == .Equal {
			advance(parser)
			skip_newlines(parser)
			expr := parse_expression(parser)
			kind := ast.Stmt_Kind.Assignment
			immutable := false
			if operator.kind == .Colon_Colon || had_type {
				kind = .Declaration
				immutable = operator.kind == .Colon_Colon
			}
			id := len(parser.module.statements)
			append(&parser.module.statements, ast.Stmt{
				kind=kind,
				span=source.Span{start=name.span.start, end=parser.module.exprs[expr].span.end},
				name=name.text,
				type=type_syntax,
				immutable=immutable,
				expr=expr,
				diagnostic=-1,
			})
			return id
		}
		parser.cursor = start_cursor
	}

	expr := parse_expression(parser)
	id := len(parser.module.statements)
	append(&parser.module.statements, ast.Stmt{
		kind=.Expression,
		span=parser.module.exprs[expr].span,
		expr=expr,
		diagnostic=-1,
	})
	return id
}

parse_params :: proc(parser: ^Parser) -> []ast.Param {
	params: [dynamic]ast.Param
	params.allocator = parser.module.allocator
	skip_newlines(parser)
	for current(parser).kind != .Right_Paren && current(parser).kind != .Eof {
		names: [dynamic]token.Token
		names.allocator = parser.module.allocator
		for {
			if current(parser).kind != .Identifier {
				source.add(parser.diagnostics, current(parser).span, "expected parameter name")
				break
			}
			append(&names, advance(parser))
			if is_type_token(current(parser).kind) {
				break
			}
			if _, ok := allow(parser, .Comma); !ok {
				source.add(parser.diagnostics, current(parser).span, "expected ',' or parameter type")
				break
			}
			skip_newlines(parser)
		}
		type_syntax := parse_type(parser)
		for name in names {
			append(&params, ast.Param{name=name.text, span=name.span, type=type_syntax})
		}
		delete(names)
		skip_newlines(parser)
		if _, ok := allow(parser, .Comma); ok {
			skip_newlines(parser)
			continue
		}
		break
	}
	return params[:]
}

parse_function :: proc(parser: ^Parser, name: token.Token, c_abi: bool) {
	func_token := advance(parser)
	if _, ok := allow(parser, .Left_Paren); !ok {
		source.add(parser.diagnostics, current(parser).span, "expected '(' after 'func'")
	}
	params := parse_params(parser)
	if _, ok := allow(parser, .Right_Paren); !ok {
		source.add(parser.diagnostics, current(parser).span, "expected ')' after parameters")
	}
	skip_newlines(parser)
	result := parse_type(parser)
	skip_newlines(parser)
	if _, ok := allow(parser, .Left_Brace); !ok {
		source.add(parser.diagnostics, current(parser).span, "expected '{' before function body")
	}

	body: [dynamic]int
	body.allocator = parser.module.allocator
	skip_newlines(parser)
	for current(parser).kind != .Right_Brace && current(parser).kind != .Eof {
		append(&body, parse_statement(parser))
		if diagnostic := finish_statement(parser); diagnostic >= 0 {
			statement_id := len(parser.module.statements)
			append(&parser.module.statements, ast.Stmt{
				kind=.Invalid,
				span=current(parser).span,
				expr=ast.INVALID_ID,
				diagnostic=diagnostic,
			})
			append(&body, statement_id)
		}
	}
	end := current(parser)
	if _, ok := allow(parser, .Right_Brace); !ok {
		source.add(parser.diagnostics, current(parser).span, "expected '}' after function body")
		end = func_token
	}
	append(&parser.module.functions, ast.Function{
		span=source.Span{start=name.span.start, end=end.span.end},
		name=name.text,
		c_abi=c_abi,
		params=params,
		result=result,
		body=body[:],
		diagnostic=-1,
	})
}

parse_top_level :: proc(parser: ^Parser) {
	if current(parser).kind != .Identifier {
		source.add(parser.diagnostics, current(parser).span, "expected a top-level declaration")
		for current(parser).kind != .Newline && current(parser).kind != .Eof {
			advance(parser)
		}
		_ = finish_statement(parser)
		return
	}
	name := advance(parser)
	type_syntax := ast.Type_Syntax.Invalid
	if is_type_token(current(parser).kind) {
		type_syntax = parse_type(parser)
	}
	operator := current(parser)
	if operator.kind != .Colon_Colon && operator.kind != .Equal {
		source.add(parser.diagnostics, operator.span, "expected '::' or '=' after top-level name")
		_ = finish_statement(parser)
		return
	}
	advance(parser)
	skip_newlines(parser)

	c_abi := false
	if operator.kind == .Colon_Colon && current(parser).kind == .Keyword_C {
		c_abi = true
		advance(parser)
		skip_newlines(parser)
	}
	if operator.kind == .Colon_Colon && current(parser).kind == .Keyword_Func {
		parse_function(parser, name, c_abi)
		return
	}

	expr := parse_expression(parser)
	append(&parser.module.globals, ast.Global{
		span=source.Span{start=name.span.start, end=parser.module.exprs[expr].span.end},
		name=name.text,
		type=type_syntax,
		immutable=operator.kind == .Colon_Colon,
		expr=expr,
		diagnostic=-1,
	})
	_ = finish_statement(parser)
}

parse :: proc(
	stream: ^token.Stream,
	diagnostics: ^source.Diagnostics,
	allocator := context.allocator,
) -> ast.Module {
	parser := Parser{
		tokens=stream,
		diagnostics=diagnostics,
		module=ast.init_module(allocator),
	}
	skip_newlines(&parser)
	for current(&parser).kind != .Eof {
		parse_top_level(&parser)
		skip_newlines(&parser)
	}
	return parser.module
}