const keywordsByLanguage: Record<string, string[]> = {
	js: [
		'const', 'let', 'var', 'function', 'return', 'if', 'else', 'for', 'while', 'class',
		'extends', 'new', 'import', 'from', 'export', 'default', 'async', 'await', 'try',
		'catch', 'finally', 'throw', 'typeof', 'instanceof', 'in', 'of', 'switch', 'case',
		'break', 'continue', 'do', 'yield', 'static', 'get', 'set', 'this', 'super', 'null',
		'undefined', 'true', 'false'
	],
	python: [
		'def', 'return', 'if', 'elif', 'else', 'for', 'while', 'class', 'import', 'from', 'as',
		'try', 'except', 'finally', 'raise', 'with', 'lambda', 'yield', 'pass', 'break',
		'continue', 'and', 'or', 'not', 'in', 'is', 'None', 'True', 'False', 'self', 'async',
		'await', 'global', 'nonlocal'
	],
	bash: [
		'if', 'then', 'else', 'elif', 'fi', 'for', 'while', 'do', 'done', 'function', 'return',
		'local', 'export', 'case', 'esac', 'in', 'echo'
	],
	go: [
		'func', 'return', 'if', 'else', 'for', 'range', 'switch', 'case', 'break', 'continue',
		'package', 'import', 'var', 'const', 'type', 'struct', 'interface', 'map', 'chan', 'go',
		'defer', 'nil', 'true', 'false'
	],
	rust: [
		'fn', 'let', 'mut', 'return', 'if', 'else', 'for', 'while', 'loop', 'match', 'struct',
		'enum', 'impl', 'trait', 'pub', 'use', 'mod', 'const', 'static', 'self', 'Self', 'true',
		'false'
	]
};

keywordsByLanguage.ts = [
	...keywordsByLanguage.js,
	'interface', 'type', 'enum', 'implements', 'public', 'private', 'protected', 'readonly',
	'namespace', 'declare', 'as'
];
keywordsByLanguage.jsx = keywordsByLanguage.js;
keywordsByLanguage.tsx = keywordsByLanguage.ts;
keywordsByLanguage.sh = keywordsByLanguage.bash;
keywordsByLanguage.py = keywordsByLanguage.python;

const defaultKeywords = keywordsByLanguage.js;

export function escapeHtml(text: string): string {
	return text
		.replace(/&/g, '&amp;')
		.replace(/</g, '&lt;')
		.replace(/>/g, '&gt;')
		.replace(/"/g, '&quot;');
}

const tokenPattern =
	/(\/\/.*$)|(#.*$)|(\/\*[\s\S]*?\*\/)|("(?:[^"\\]|\\.)*")|('(?:[^'\\]|\\.)*')|(`(?:[^`\\]|\\.)*`)|(\b\d+(?:\.\d+)?\b)|([A-Za-z_$][A-Za-z0-9_$]*)/gm;

export function highlightCode(code: string, lang: string): string {
	const keywords = new Set(keywordsByLanguage[lang.toLowerCase()] ?? defaultKeywords);
	let result = '';
	let lastIndex = 0;
	let match: RegExpExecArray | null;

	tokenPattern.lastIndex = 0;
	while ((match = tokenPattern.exec(code)) !== null) {
		result += escapeHtml(code.slice(lastIndex, match.index));

		const [full, lineComment, hashComment, blockComment, dquote, squote, backtick, number, word] = match;

		if (lineComment || hashComment || blockComment) {
			result += `<span class="tok-comment">${escapeHtml(full)}</span>`;
		} else if (dquote || squote || backtick) {
			result += `<span class="tok-string">${escapeHtml(full)}</span>`;
		} else if (number) {
			result += `<span class="tok-number">${escapeHtml(full)}</span>`;
		} else if (word) {
			result += keywords.has(word)
				? `<span class="tok-keyword">${word}</span>`
				: escapeHtml(word);
		}

		lastIndex = match.index + full.length;
	}

	result += escapeHtml(code.slice(lastIndex));
	return result;
}
