package shader

import "core:fmt"

Lexer :: struct {
	source:      string,
	file:        string,
	pos:         int,
	line:        int,
	col:         int,
	tokens:      [dynamic]Token,
	diagnostics: [dynamic]Diagnostic,
}

lexer_init :: proc(source: string, file := "", allocator := context.allocator) -> Lexer {
	return Lexer{
		source      = source,
		file        = file,
		pos         = 0,
		line        = 1,
		col         = 1,
		tokens      = make([dynamic]Token, allocator),
		diagnostics = make([dynamic]Diagnostic, allocator),
	}
}

tokenize :: proc(l: ^Lexer) -> ([]Token, []Diagnostic) {
	for !is_at_end(l) {
		skip_whitespace_and_comments(l)
		if is_at_end(l) do break
		scan_token(l)
	}

	append(&l.tokens, Token{
		kind = .EOF,
		text = "",
		span = make_span(l, l.pos, l.pos),
	})

	return l.tokens[:], l.diagnostics[:]
}

// Internal

@(private = "file")
is_at_end :: proc(l: ^Lexer) -> bool {
	return l.pos >= len(l.source)
}

@(private = "file")
peek :: proc(l: ^Lexer) -> u8 {
	if is_at_end(l) do return 0
	return l.source[l.pos]
}

@(private = "file")
peek_next :: proc(l: ^Lexer) -> u8 {
	if l.pos + 1 >= len(l.source) do return 0
	return l.source[l.pos + 1]
}

@(private = "file")
advance :: proc(l: ^Lexer) -> u8 {
	ch := l.source[l.pos]
	l.pos += 1
	if ch == '\n' {
		l.line += 1
		l.col = 1
	} else {
		l.col += 1
	}
	return ch
}

@(private = "file")
make_span :: proc(l: ^Lexer, start_pos: int, end_pos: int) -> Source_Span {
	// Calculate start line/col by scanning from beginning
	// This is O(n) but only done per token — acceptable for shader-length sources
	start_line := 1
	start_col := 1
	for i in 0 ..< start_pos {
		if i < len(l.source) && l.source[i] == '\n' {
			start_line += 1
			start_col = 1
		} else {
			start_col += 1
		}
	}
	return Source_Span{
		file       = l.file,
		line_start = start_line,
		col_start  = start_col,
		line_end   = l.line,
		col_end    = l.col,
	}
}

@(private = "file")
skip_whitespace_and_comments :: proc(l: ^Lexer) {
	for !is_at_end(l) {
		ch := peek(l)
		switch ch {
		case ' ', '\t', '\r', '\n':
			advance(l)
		case '-':
			if peek_next(l) == '-' {
				skip_comment(l)
			} else {
				return
			}
		case:
			return
		}
	}
}

@(private = "file")
skip_comment :: proc(l: ^Lexer) {
	// Skip the two dashes
	advance(l)
	advance(l)

	// Check for block comment --[[ ... ]]
	if !is_at_end(l) && peek(l) == '[' && peek_next(l) == '[' {
		advance(l) // [
		advance(l) // [
		for !is_at_end(l) {
			if peek(l) == ']' && peek_next(l) == ']' {
				advance(l) // ]
				advance(l) // ]
				return
			}
			advance(l)
		}
		// Unterminated block comment — diagnostic added at EOF
		append(&l.diagnostics, Diagnostic{
			level   = .Error,
			message = "unterminated block comment",
			span    = make_span(l, l.pos, l.pos),
		})
		return
	}

	// Line comment — skip to end of line
	for !is_at_end(l) && peek(l) != '\n' {
		advance(l)
	}
}

@(private = "file")
scan_token :: proc(l: ^Lexer) {
	start := l.pos
	ch := peek(l)

	switch {
	case is_alpha(ch) || ch == '_':
		scan_identifier(l)
	case is_digit(ch):
		scan_number(l)
	case ch == '"':
		scan_string(l)
	case:
		scan_symbol(l)
	}
}

@(private = "file")
scan_identifier :: proc(l: ^Lexer) {
	start := l.pos
	for !is_at_end(l) && (is_alpha(peek(l)) || is_digit(peek(l)) || peek(l) == '_') {
		advance(l)
	}

	text := l.source[start:l.pos]
	kind: Token_Kind
	if kw, ok := keyword_kind(text); ok {
		kind = kw
	} else {
		kind = .Identifier
	}

	append(&l.tokens, Token{
		kind = kind,
		text = text,
		span = make_span(l, start, l.pos),
	})
}

@(private = "file")
scan_number :: proc(l: ^Lexer) {
	start := l.pos
	is_float := false

	// Integer part
	for !is_at_end(l) && is_digit(peek(l)) {
		advance(l)
	}

	// Fractional part
	if !is_at_end(l) && peek(l) == '.' && is_digit(peek_next(l)) {
		is_float = true
		advance(l) // consume '.'
		for !is_at_end(l) && is_digit(peek(l)) {
			advance(l)
		}
	}

	// Exponent part
	if !is_at_end(l) && (peek(l) == 'e' || peek(l) == 'E') {
		is_float = true
		advance(l)
		if !is_at_end(l) && (peek(l) == '+' || peek(l) == '-') {
			advance(l)
		}
		for !is_at_end(l) && is_digit(peek(l)) {
			advance(l)
		}
	}

	append(&l.tokens, Token{
		kind = is_float ? .Float : .Integer,
		text = l.source[start:l.pos],
		span = make_span(l, start, l.pos),
	})
}

@(private = "file")
scan_string :: proc(l: ^Lexer) {
	start := l.pos
	advance(l) // skip opening quote

	for !is_at_end(l) && peek(l) != '"' && peek(l) != '\n' {
		if peek(l) == '\\' {
			advance(l) // skip escape char
		}
		advance(l)
	}

	if is_at_end(l) || peek(l) == '\n' {
		append(&l.diagnostics, Diagnostic{
			level   = .Error,
			message = "unterminated string",
			span    = make_span(l, start, l.pos),
		})
		return
	}

	advance(l) // skip closing quote

	append(&l.tokens, Token{
		kind = .String,
		text = l.source[start:l.pos],
		span = make_span(l, start, l.pos),
	})
}

@(private = "file")
scan_symbol :: proc(l: ^Lexer) {
	start := l.pos
	ch := advance(l)

	kind: Token_Kind
	switch ch {
	case '(':  kind = .Lparen
	case ')':  kind = .Rparen
	case '[':  kind = .Lbracket
	case ']':  kind = .Rbracket
	case '{':  kind = .Lbrace
	case '}':  kind = .Rbrace
	case '.':  kind = .Dot
	case ',':  kind = .Comma
	case ':':  kind = .Colon
	case '+':  kind = .Plus
	case '*':  kind = .Star
	case '/':  kind = .Slash
	case '%':  kind = .Percent
	case '@':  kind = .At
	case '#':  kind = .Hash
	case '-':
		if !is_at_end(l) && peek(l) == '>' {
			advance(l)
			kind = .Arrow
		} else {
			kind = .Minus
		}
	case '=':
		if !is_at_end(l) && peek(l) == '=' {
			advance(l)
			kind = .Eq_Eq
		} else {
			kind = .Eq
		}
	case '!':
		if !is_at_end(l) && peek(l) == '=' {
			advance(l)
			kind = .Not_Eq
		} else {
			append(&l.diagnostics, Diagnostic{
				level   = .Error,
				message = "unexpected character '!'",
				span    = make_span(l, start, l.pos),
			})
			return
		}
	case '<':
		if !is_at_end(l) && peek(l) == '=' {
			advance(l)
			kind = .Lt_Eq
		} else {
			kind = .Lt
		}
	case '>':
		if !is_at_end(l) && peek(l) == '=' {
			advance(l)
			kind = .Gt_Eq
		} else {
			kind = .Gt
		}
	case:
		append(&l.diagnostics, Diagnostic{
			level   = .Error,
			message = fmt.aprintf("unexpected character '%c'", ch),
			span    = make_span(l, start, l.pos),
		})
		return
	}

	append(&l.tokens, Token{
		kind = kind,
		text = l.source[start:l.pos],
		span = make_span(l, start, l.pos),
	})
}

@(private = "file")
is_alpha :: proc(ch: u8) -> bool {
	return (ch >= 'a' && ch <= 'z') || (ch >= 'A' && ch <= 'Z')
}

@(private = "file")
is_digit :: proc(ch: u8) -> bool {
	return ch >= '0' && ch <= '9'
}