import hljs from "highlight.js/lib/core";
import javascript from "highlight.js/lib/languages/javascript";
import typescript from "highlight.js/lib/languages/typescript";
import python from "highlight.js/lib/languages/python";
import bash from "highlight.js/lib/languages/bash";
import json from "highlight.js/lib/languages/json";
import markdown from "highlight.js/lib/languages/markdown";
import css from "highlight.js/lib/languages/css";
import xml from "highlight.js/lib/languages/xml";
import { Marked, type Tokens } from "marked";
import "highlight.js/styles/github.min.css";

hljs.registerLanguage("javascript", javascript);
hljs.registerLanguage("typescript", typescript);
hljs.registerLanguage("python", python);
hljs.registerLanguage("bash", bash);
hljs.registerLanguage("json", json);
hljs.registerLanguage("markdown", markdown);
hljs.registerLanguage("css", css);
hljs.registerLanguage("xml", xml);
hljs.registerLanguage("html", xml);

function escapeMarkdownCode(raw: string): string {
	return String(raw)
		.replace(/&/g, "&")
		.replace(/</g, "&lt;")
		.replace(/>/g, "&gt;")
		.replace(/"/g, "&quot;");
}

function renderAssistantCodeBlock(token: Tokens.Code): string {
	const langSlug = (((token.lang || "").match(/^\S+/) || [""])[0]).trim().toLowerCase();
	const codeRaw = String(token.text ?? "").replace(/\n+$/, "");
	const skipHl = !!token.escaped;

	let innerHtml: string;
	if (skipHl) {
		innerHtml = codeRaw;
	} else {
		try {
			if (langSlug && hljs.getLanguage(langSlug)) {
				innerHtml = hljs.highlight(codeRaw, { language: langSlug }).value;
			} else {
				innerHtml = hljs.highlightAuto(codeRaw).value;
			}
		} catch {
			innerHtml = escapeMarkdownCode(codeRaw);
		}
	}

	const safeLangSlug = /^[a-z][a-z0-9_-]*$/i.test(langSlug) ? langSlug : "";
	const codeClass = safeLangSlug ? `hljs language-${safeLangSlug}` : "hljs";
	return `<pre class="assistant-pre"><code class="${codeClass}">${innerHtml}</code></pre>\n`;
}

const marked = new Marked({
	breaks: true,
	gfm: true,
	renderer: {
		code: renderAssistantCodeBlock,
	},
});

export function renderMarkdown(text: string): string {
	if (!text) return "";
	return marked.parse(text) as string;
}