import { remark } from "remark";
import remarkGfm from "remark-gfm";
import remarkRehype from "remark-rehype";
import { toJsxRuntime } from "hast-util-to-jsx-runtime";
import {
	Children,
	type ComponentProps,
	type ReactElement,
	type ReactNode,
	Suspense,
	use,
	useDeferredValue,
} from "react";
import { Fragment, jsx, jsxs } from "react/jsx-runtime";
import { DynamicCodeBlock } from "fumadocs-ui/components/dynamic-codeblock";
import defaultMdxComponents from "fumadocs-ui/mdx";
import { visit } from "unist-util-visit";
import type { ElementContent, Root, RootContent } from "hast";

export interface Processor {
	process: (content: string) => Promise<ReactNode>;
}

export function rehypeWrapWords() {
	return (tree: Root) => {
		visit(tree, ["text", "element"], (node, index, parent) => {
			if (node.type === "element" && node.tagName === "pre") return "skip";
			if (node.type !== "text" || !parent || index === undefined) return;

			const words = node.value.split(/(?=\s)/);

			// Create new span nodes for each word and whitespace
			const newNodes: ElementContent[] = words.flatMap((word) => {
				if (word.length === 0) return [];

				return {
					type: "element",
					tagName: "span",
					properties: {
						class: "animate-fd-fade-in",
					},
					children: [{ type: "text", value: word }],
				};
			});

			Object.assign(node, {
				type: "element",
				tagName: "span",
				properties: {},
				children: newNodes,
			} satisfies RootContent);
			return "skip";
		});
	};
}

function createProcessor(): Processor {
	const processor = remark()
		.use(remarkGfm)
		.use(remarkRehype)
		.use(rehypeWrapWords);

	return {
		async process(content) {
			const nodes = processor.parse({ value: content });
			const hast = await processor.run(nodes);

			return toJsxRuntime(hast, {
				development: false,
				jsx,
				jsxs,
				Fragment,
				components: {
					...defaultMdxComponents,
					pre: Pre,
					img: undefined, // use JSX
				},
			});
		},
	};
}

function Pre(props: ComponentProps<"pre">) {
	const code = Children.only(props.children) as ReactElement;
	const codeProps = code.props as ComponentProps<"code">;
	const content = codeProps.children;
	if (typeof content !== "string") return null;

	let lang =
		codeProps.className
			?.split(" ")
			.find((v) => v.startsWith("language-"))
			?.slice("language-".length) ?? "text";

	if (lang === "mdx") lang = "md";

	return <DynamicCodeBlock lang={lang} code={content.trimEnd()} />;
}

const processor = createProcessor();

export function Markdown({ text }: { text: string }) {
	const deferredText = useDeferredValue(text);

	return (
		<Suspense fallback={<p className="invisible">{text}</p>}>
			<Renderer text={deferredText} />
		</Suspense>
	);
}

const cache = new Map<string, Promise<ReactNode>>();

function Renderer({ text }: { text: string }) {
	const result = cache.get(text) ?? processor.process(text);
	cache.set(text, result);

	return use(result);
}