package c_frontend_tokenizer

import "core:fmt"
import "core:os"
import "core:strings"
import "core:unicode/utf8"


Error_Handler :: #type proc(pos: Pos, fmt: string, args: ..any)


Tokenizer :: struct {
	// Immutable data
	path: string,
	src:  []byte,


	// Tokenizing state
	ch:          rune,
	offset:      int,
	read_offset: int,
	line_offset: int,
	line_count:  int,

	// Extra information for tokens
	at_bol:    bool,
	has_space: bool,

	// Mutable data
	err:  Error_Handler,
	warn: Error_Handler,
	error_count:   int,
	warning_count: int,
}

init_defaults :: proc(t: ^Tokenizer, err: Error_Handler = default_error_handler, warn: Error_Handler = default_warn_handler) {
	t.err = err
	t.warn = warn
}


@(private)
offset_to_pos :: proc(t: ^Tokenizer, offset: int) -> (pos: Pos) {
	pos.file = t.path
	pos.offset = offset
	pos.line = t.line_count
	pos.column = offset - t.line_offset + 1
	return
}

default_error_handler :: proc(pos: Pos, msg: string, args: ..any) {
	fmt.eprintf("%s(%d:%d) ", pos.file, pos.line, pos.column)
	fmt.eprintf(msg, ..args)
	fmt.eprintf("\n")
}

default_warn_handler :: proc(pos: Pos, msg: string, args: ..any) {
	fmt.eprintf("%s(%d:%d) warning: ", pos.file, pos.line, pos.column)
	fmt.eprintf(msg, ..args)
	fmt.eprintf("\n")
}

error_offset :: proc(t: ^Tokenizer, offset: int, msg: string, args: ..any) {
	pos := offset_to_pos(t, offset)
	if t.err != nil {
		t.err(pos, msg, ..args)
	}
	t.error_count += 1
}

warn_offset :: proc(t: ^Tokenizer, offset: int, msg: string, args: ..any) {
	pos := offset_to_pos(t, offset)
	if t.warn != nil {
		t.warn(pos, msg, ..args)
	}
	t.warning_count += 1
}

error :: proc(t: ^Tokenizer, tok: ^Token, msg: string, args: ..any) {
	pos := tok.pos
	if t.err != nil {
		t.err(pos, msg, ..args)
	}
	t.error_count += 1
}

warn :: proc(t: ^Tokenizer, tok: ^Token, msg: string, args: ..any) {
	pos := tok.pos
	if t.warn != nil {
		t.warn(pos, msg, ..args)
	}
	t.warning_count += 1
}


advance_rune :: proc(t: ^Tokenizer) {
	if t.read_offset < len(t.src) {
		t.offset = t.read_offset
		if t.ch == '\n' {
			t.at_bol = true
			t.line_offset = t.offset
			t.line_count += 1
		}
		r, w := rune(t.src[t.read_offset]), 1
		switch {
		case r == 0:
			error_offset(t, t.offset, "illegal character NUL")
		case r >= utf8.RUNE_SELF:
			r, w = utf8.decode_rune(t.src[t.read_offset:])
			if r == utf8.RUNE_ERROR && w == 1 {
				error_offset(t, t.offset, "illegal UTF-8 encoding")
			} else if r == utf8.RUNE_BOM && t.offset > 0 {
				error_offset(t, t.offset, "illegal byte order mark")
			}
		}
		t.read_offset += w
		t.ch = r
	} else {
		t.offset = len(t.src)
		if t.ch == '\n' {
			t.at_bol = true
			t.line_offset = t.offset
			t.line_count += 1
		}
		t.ch = -1
	}
}

advance_rune_n :: proc(t: ^Tokenizer, n: int) {
	for in 0..<n {
		advance_rune(t)
	}
}

is_digit :: proc(r: rune) -> bool {
	return '0' <= r && r <= '9'
}

skip_whitespace :: proc(t: ^Tokenizer) {
	for {
		switch t.ch {
		case ' ', '\t', '\r', '\v', '\f', '\n':
			t.has_space = true
			advance_rune(t)
		case:
			return
		}
	}
}

scan_comment :: proc(t: ^Tokenizer) -> string {
	offset := t.offset-1
	next := -1
	general: {
		if t.ch == '/'{ // line comments
			advance_rune(t)
			for t.ch != '\n' && t.ch >= 0 {
				advance_rune(t)
			}

			next = t.offset
			if t.ch == '\n' {
				next += 1
			}
			break general
		}

		/* style comment */
		advance_rune(t)
		for t.ch >= 0 {
			ch := t.ch
			advance_rune(t)
			if ch == '*' && t.ch == '/' {
				advance_rune(t)
				next = t.offset
				break general
			}
		}

		error_offset(t, offset, "comment not terminated")
	}

	lit := t.src[offset : t.offset]

	// NOTE(bill): Strip CR for line comments
	for len(lit) > 2 && lit[1] == '/' && lit[len(lit)-1] == '\r' {
		lit = lit[:len(lit)-1]
	}


	return string(lit)
}

scan_identifier :: proc(t: ^Tokenizer) -> string {
	offset := t.offset

	for is_ident1(t.ch) {
		advance_rune(t)
	}

	return string(t.src[offset : t.offset])
}

scan_string :: proc(t: ^Tokenizer) -> string {
	offset := t.offset-1

	for {
		ch := t.ch
		if ch == '\n' || ch < 0 {
			error_offset(t, offset, "string literal was not terminated")
			break
		}
		advance_rune(t)
		if ch == '"' {
			break
		}
		if ch == '\\' {
			scan_escape(t)
		}
	}

	return string(t.src[offset : t.offset])
}

digit_val :: proc(r: rune) -> int {
	switch r {
	case '0'..='9':
		return int(r-'0')
	case 'A'..='F':
		return int(r-'A' + 10)
	case 'a'..='f':
		return int(r-'a' + 10)
	}
	return 16
}

scan_escape :: proc(t: ^Tokenizer) -> bool {
	offset := t.offset

	esc := t.ch
	n: int
	base, max: u32
	switch esc {
	case 'a', 'b', 'e', 'f', 'n', 't', 'v', 'r', '\\', '\'', '"':
		advance_rune(t)
		return true

	case '0'..='7':
		for digit_val(t.ch) < 8 {
			advance_rune(t)
		}
		return true
	case 'x':
		advance_rune(t)
		for digit_val(t.ch) < 16 {
			advance_rune(t)
		}
		return true
	case 'u':
		advance_rune(t)
		n, base, max = 4, 16, utf8.MAX_RUNE
	case 'U':
		advance_rune(t)
		n, base, max = 8, 16, utf8.MAX_RUNE
	case:
		if t.ch < 0 {
			error_offset(t, offset, "escape sequence was not terminated")
		} else {
			break
		}
		return false
	}

	x: u32
	main_loop: for n > 0 {
		d := u32(digit_val(t.ch))
		if d >= base {
			if t.ch == '"' || t.ch == '\'' {
				break main_loop
			}
			if t.ch < 0 {
				error_offset(t, t.offset, "escape sequence was not terminated")
			} else {
				error_offset(t, t.offset, "illegal character '%r' : %d in escape sequence", t.ch, t.ch)
			}
			return false
		}

		x = x*base + d
		advance_rune(t)
		n -= 1
	}

	if x > max || 0xd800 <= x && x <= 0xe000 {
		error_offset(t, offset, "escape sequence is an invalid Unicode code point")
		return false
	}
	return true
}

scan_rune :: proc(t: ^Tokenizer) -> string {
	offset := t.offset-1
	valid := true
	n := 0
	for {
		ch := t.ch
		if ch == '\n' || ch < 0 {
			if valid {
				error_offset(t, offset, "rune literal not terminated")
				valid = false
			}
			break
		}
		advance_rune(t)
		if ch == '\'' {
			break
		}
		n += 1
		if ch == '\\' {
			if !scan_escape(t)  {
				valid = false
			}
		}
	}

	if valid && n != 1 {
		error_offset(t, offset, "illegal rune literal")
	}

	return string(t.src[offset : t.offset])
}

scan_number :: proc(t: ^Tokenizer, seen_decimal_point: bool) -> (Token_Kind, string) {
	scan_mantissa :: proc(t: ^Tokenizer, base: int) {
		for digit_val(t.ch) < base {
			advance_rune(t)
		}
	}
	scan_exponent :: proc(t: ^Tokenizer) {
		if t.ch == 'e' || t.ch == 'E' || t.ch == 'p' || t.ch == 'P' {
			advance_rune(t)
			if t.ch == '-' || t.ch == '+' {
				advance_rune(t)
			}
			if digit_val(t.ch) < 10 {
				scan_mantissa(t, 10)
			} else {
				error_offset(t, t.offset, "illegal floating-point exponent")
			}
		}
	}
	scan_fraction :: proc(t: ^Tokenizer) -> (early_exit: bool) {
		if t.ch == '.' && peek(t) == '.' {
			return true
		}
		if t.ch == '.' {
			advance_rune(t)
			scan_mantissa(t, 10)
		}
		return false
	}

	check_end := true


	offset := t.offset
	seen_point := seen_decimal_point

	if seen_point {
		offset -= 1
		scan_mantissa(t, 10)
		scan_exponent(t)
	} else {
		if t.ch == '0' {
			int_base :: proc(t: ^Tokenizer, base: int, msg: string) {
				prev := t.offset
				advance_rune(t)
				scan_mantissa(t, base)
				if t.offset - prev <= 1 {
					error_offset(t, t.offset, msg)
				}
			}

			advance_rune(t)
			switch t.ch {
			case 'b', 'B':
				int_base(t, 2, "illegal binary integer")
			case 'x', 'X':
				int_base(t, 16, "illegal hexadecimal integer")
			case:
				seen_point = false
				scan_mantissa(t, 10)
				if t.ch == '.' {
					seen_point = true
					if scan_fraction(t) {
						check_end = false
					}
				}
				if check_end {
					scan_exponent(t)
					check_end = false
				}
			}
		}
	}

	if check_end {
		scan_mantissa(t, 10)

		if !scan_fraction(t) {
			scan_exponent(t)
		}
	}

	return .Number, string(t.src[offset : t.offset])
}

scan_punct :: proc(t: ^Tokenizer, ch: rune) -> (kind: Token_Kind) {
	kind = .Punct
	switch ch {
	case:
		kind = .Invalid

	case '<', '>':
		if t.ch == ch {
			advance_rune(t)
		}
		if t.ch == '=' {
			advance_rune(t)
		}
	case '!', '+', '-', '*', '/', '%', '^', '=':
		if t.ch == '=' {
			advance_rune(t)
		}
	case '#':
		if t.ch == '#' {
			advance_rune(t)
		}
	case '&':
		if t.ch == '=' || t.ch == '&' {
			advance_rune(t)
		}
	case '|':
		if t.ch == '=' || t.ch == '|' {
			advance_rune(t)
		}
	case '(', ')', '[', ']', '{', '}':
		// okay
	case '~', ',', ':', ';', '?':
		// okay
	case '`':
		// okay
	case '.':
		if t.ch == '.' && peek(t) == '.' {
			advance_rune(t)
			advance_rune(t) // consume last '.'
		}
	}
	return
}

peek :: proc(t: ^Tokenizer) -> byte {
	if t.read_offset < len(t.src) {
		return t.src[t.read_offset]
	}
	return 0
}
peek_str :: proc(t: ^Tokenizer, str: string) -> bool {
	if t.read_offset < len(t.src) {
		return strings.has_prefix(string(t.src[t.offset:]), str)
	}
	return false
}

scan_literal_prefix :: proc(t: ^Tokenizer, str: string, prefix: ^string) -> bool {
	if peek_str(t, str) {
		offset := t.offset
		for _ in str {
			advance_rune(t)
		}
		prefix^ = string(t.src[offset:][:len(str)-1])
		return true
	}
	return false
}


allow_next_to_be_newline :: proc(t: ^Tokenizer) -> bool {
	if t.ch == '\n' {
		advance_rune(t)
		return true
	} else if t.ch == '\r' && peek(t) == '\n' { // allow for MS-DOS style line endings
		advance_rune(t) // \r
		advance_rune(t) // \n
		return true
	}
	return false
}

scan :: proc(t: ^Tokenizer, f: ^File) -> ^Token {
	skip_whitespace(t)

	offset := t.offset

	kind: Token_Kind
	lit: string
	prefix: string

	switch ch := t.ch; {
	case scan_literal_prefix(t, `u8"`, &prefix):
		kind = .String
		lit = scan_string(t)
	case scan_literal_prefix(t, `u"`, &prefix):
		kind = .String
		lit = scan_string(t)
	case scan_literal_prefix(t, `L"`, &prefix):
		kind = .String
		lit = scan_string(t)
	case scan_literal_prefix(t, `U"`, &prefix):
		kind = .String
		lit = scan_string(t)
	case scan_literal_prefix(t, `u'`, &prefix):
		kind = .Char
		lit = scan_rune(t)
	case scan_literal_prefix(t, `L'`, &prefix):
		kind = .Char
		lit = scan_rune(t)
	case scan_literal_prefix(t, `U'`, &prefix):
		kind = .Char
		lit = scan_rune(t)

	case is_ident0(ch):
		lit = scan_identifier(t)
		kind = .Ident
	case '0' <= ch && ch <= '9':
		kind, lit = scan_number(t, false)
	case:
		advance_rune(t)
		switch ch {
		case -1:
			kind = .EOF
		case '\\':
			kind = .Punct
			if allow_next_to_be_newline(t) {
				t.at_bol = true
				t.has_space = false
				return scan(t, f)
			}

		case '.':
			if is_digit(t.ch) {
				kind, lit = scan_number(t, true)
			} else {
				kind = scan_punct(t, ch)
			}
		case '"':
			kind = .String
			lit = scan_string(t)
		case '\'':
			kind = .Char
			lit = scan_rune(t)
		case '/':
			if t.ch == '/' || t.ch == '*' {
				kind = .Comment
				lit = scan_comment(t)
				t.has_space = true
				break
			}
			fallthrough
		case:
			kind = scan_punct(t, ch)
			if kind == .Invalid && ch != utf8.RUNE_BOM {
				error_offset(t, t.offset, "illegal character '%r': %d", ch, ch)
			}
		}
	}

	if lit == "" {
		lit = string(t.src[offset : t.offset])
	}

	if kind == .Comment {
		return scan(t, f)
	}

	tok := new(Token)
	tok.kind = kind
	tok.lit = lit
	tok.pos = offset_to_pos(t, offset)
	tok.file = f
	tok.prefix = prefix
	tok.at_bol = t.at_bol
	tok.has_space = t.has_space

	t.at_bol, t.has_space = false, false

	return tok
}

tokenize :: proc(t: ^Tokenizer, f: ^File) -> ^Token {
	setup_tokenizer: {
		t.src = f.src
		t.ch = ' '
		t.offset = 0
		t.read_offset = 0
		t.line_offset = 0
		t.line_count = len(t.src) > 0 ? 1 : 0
		t.error_count = 0
		t.path = f.name


		advance_rune(t)
		if t.ch == utf8.RUNE_BOM {
			advance_rune(t)
		}
	}


	t.at_bol = true
	t.has_space = false

	head: Token
	curr := &head
	for {
		tok := scan(t, f)
		if tok == nil {
			break
		}
		curr.next = tok
		curr = curr.next
		if tok.kind == .EOF {
			break
		}
	}

	return head.next
}

add_new_file :: proc(t: ^Tokenizer, name: string, src: []byte, id: int) -> ^File {
	file := new(File)
	file.id = id
	file.src = src
	file.name = name
	file.display_name = name
	return file
}

tokenize_file :: proc(t: ^Tokenizer, path: string, id: int, loc := #caller_location) -> ^Token {
	src, ok := os.read_entire_file(path)
	if !ok {
		return nil
	}
	return tokenize(t, add_new_file(t, path, src, id))
}


inline_tokenize :: proc(t: ^Tokenizer, tok: ^Token, src: []byte) -> ^Token {
	file := new(File)
	file.src = src
	if tok.file != nil {
		file.id = tok.file.id
		file.name = tok.file.name
		file.display_name = tok.file.name
	}

	return tokenize(t, file)
}
