module main

import term
import strings
import v.scanner
import v.ast
import v.token
import v.pref

const highlight_builtin_types = ['bool', 'string', 'i8', 'i16', 'int', 'i64', 'i128', 'isize',
	'byte', 'u8', 'u16', 'u32', 'u64', 'usize', 'u128', 'rune', 'f32', 'f64', 'byteptr', 'voidptr',
	'any']

fn color_highlight(code string, tb &ast.Table) string {
	highlight_code := fn (tok token.Token, typ HighlightTokenTyp) string {
		mut lit := ''
		match typ {
			.unone, .operator, .punctuation {
				lit = tok.kind.str()
			}
			.string {
				use_double_quote := tok.lit.contains("'") && !tok.lit.contains('"')
				unescaped_val := tok.lit.replace('\\\\', '\x01').replace_each(["\\'", "'", '\\"',
					'"'])
				if use_double_quote {
					s := unescaped_val.replace_each(['\x01', '\\\\', '"', '\\"'])
					lit = term.yellow('"${s}"')
				} else {
					s := unescaped_val.replace_each(['\x01', '\\\\', "'", "\\'"])
					lit = term.yellow("'${s}'")
				}
			}
			.char {
				lit = term.yellow('`${tok.lit}`')
			}
			.comment {
				lit = if tok.lit != '' && tok.lit[0] == 1 {
					term.gray('//${tok.lit[1..]}')
				} else {
					term.gray('//${tok.lit}')
				}
			}
			.keyword {
				lit = term.bright_blue(tok.lit)
			}
			.builtin, .symbol {
				lit = term.green(tok.lit)
			}
			.function {
				lit = term.cyan(tok.lit)
			}
			.number, .module_ {
				lit = term.bright_blue(tok.lit)
			}
			.boolean {
				lit = term.bright_magenta(tok.lit)
			}
			.none_ {
				lit = term.red(tok.lit)
			}
			.prefix {
				lit = term.magenta(tok.lit)
			}
			else {
				lit = tok.lit
			}
		}
		return lit
	}
	mut s := scanner.new_scanner(code, .parse_comments, &pref.Preferences{ is_fmt: true })
	mut prev_prev := token.Token{}
	mut prev := token.Token{}
	mut tok := s.scan()
	mut next_tok := s.scan()
	mut buf := strings.new_builder(200)
	mut i := 0
	for i < code.len {
		if i == tok.pos {
			mut tok_typ := HighlightTokenTyp.unone
			match tok.kind {
				.name {
					if (tok.lit in highlight_builtin_types || tb.known_type(tok.lit))
						&& (next_tok.kind != .lpar || prev.kind !in [.key_fn, .rpar]) {
						tok_typ = .builtin
					} else if
						(next_tok.kind in [.lcbr, .rpar, .eof, .comma, .pipe, .name, .rcbr, .assign, .key_pub, .key_mut, .pipe, .comma, .comment, .lt, .lsbr]
						&& next_tok.lit !in highlight_builtin_types)
						&& (prev.kind in [.name, .amp, .lcbr, .rsbr, .key_type, .assign, .dot, .not, .question, .rpar, .key_struct, .key_enum, .pipe, .key_interface, .comment, .ellipsis, .comma]
						&& prev.lit !in highlight_builtin_types)
						&& ((tok.lit != '' && tok.lit[0].is_capital())
						|| prev_prev.lit in ['C', 'JS']) {
						tok_typ = .symbol
					} else if tok.lit[0].is_capital() && prev.kind == .lpar
						&& next_tok.kind == .comma {
						tok_typ = .symbol
					} else if next_tok.kind == .lpar
						|| (!(tok.lit != '' && tok.lit[0].is_capital())
						&& next_tok.kind in [.lt, .lsbr] && next_tok.pos == tok.pos + tok.lit.len) {
						tok_typ = .function
					} else if next_tok.kind == .dot {
						if tok.lit in ['C', 'JS'] {
							tok_typ = .prefix
						} else {
							if tok.lit != '' && tok.lit[0].is_capital() {
								tok_typ = .symbol
							} else {
								tok_typ = .module_
							}
						}
					} else if tok.lit in ['r', 'c'] && next_tok.kind == .string {
						tok_typ = .prefix
					} else {
						tok_typ = .name
					}
				}
				.comment {
					tok_typ = .comment
				}
				.chartoken {
					tok_typ = .char
				}
				.string {
					tok_typ = .string
				}
				.number {
					tok_typ = .number
				}
				.key_true, .key_false {
					tok_typ = .boolean
				}
				.lpar, .lcbr, .rpar, .rcbr, .lsbr, .rsbr, .semicolon, .colon, .comma, .dot,
				.dotdot, .ellipsis {
					tok_typ = .punctuation
				}
				.key_none {
					tok_typ = .none_
				}
				else {
					if token.is_key(tok.lit) || token.is_decl(tok.kind) {
						tok_typ = .keyword
					} else if tok.kind == .decl_assign || tok.kind.is_assign() || tok.is_unary()
						|| tok.kind.is_relational() || tok.kind.is_infix() || tok.kind.is_postfix() {
						tok_typ = .operator
					}
				}
			}
			buf.write_string(highlight_code(tok, tok_typ))
			if prev_prev.kind == .eof || prev.kind == .eof || next_tok.kind == .eof {
				break
			}
			prev_prev = prev
			prev = tok
			i = tok.pos + tok.len
			tok = next_tok
			next_tok = s.scan()
		} else {
			buf.write_u8(code[i])
			i++
		}
	}
	return buf.str()
}