1/**
2 * Rehype plugin to enhance code blocks with wrapper, header, and action buttons.
3 *
4 * Wraps <pre><code> elements with a container that includes:
5 * - Language label
6 * - Copy button
7 * - Preview button (for HTML code blocks)
8 *
9 * This operates directly on the HAST tree for better performance,
10 * avoiding the need to stringify and re-parse HTML.
11 */
12
13import type { Plugin } from 'unified';
14import type { Root, Element, ElementContent } from 'hast';
15import { visit } from 'unist-util-visit';
16
17declare global {
18 interface Window {
19 idxCodeBlock?: number;
20 }
21}
22
23const COPY_ICON_SVG = `<svg xmlns="http://www.w3.org/2000/svg" width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round" class="lucide lucide-copy-icon lucide-copy"><rect width="14" height="14" x="8" y="8" rx="2" ry="2"/><path d="M4 16c-1.1 0-2-.9-2-2V4c0-1.1.9-2 2-2h10c1.1 0 2 .9 2 2"/></svg>`;
24
25const PREVIEW_ICON_SVG = `<svg xmlns="http://www.w3.org/2000/svg" width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round" class="lucide lucide-eye lucide-eye-icon"><path d="M2.062 12.345a1 1 0 0 1 0-.69C3.5 7.73 7.36 5 12 5s8.5 2.73 9.938 6.655a1 1 0 0 1 0 .69C20.5 16.27 16.64 19 12 19s-8.5-2.73-9.938-6.655"/><circle cx="12" cy="12" r="3"/></svg>`;
26
27/**
28 * Creates an SVG element node from raw SVG string.
29 * Since we can't parse HTML in HAST directly, we use the raw property.
30 */
31function createRawHtmlElement(html: string): Element {
32 return {
33 type: 'element',
34 tagName: 'span',
35 properties: {},
36 children: [{ type: 'raw', value: html } as unknown as ElementContent]
37 };
38}
39
40function createCopyButton(codeId: string): Element {
41 return {
42 type: 'element',
43 tagName: 'button',
44 properties: {
45 className: ['copy-code-btn'],
46 'data-code-id': codeId,
47 title: 'Copy code',
48 type: 'button'
49 },
50 children: [createRawHtmlElement(COPY_ICON_SVG)]
51 };
52}
53
54function createPreviewButton(codeId: string): Element {
55 return {
56 type: 'element',
57 tagName: 'button',
58 properties: {
59 className: ['preview-code-btn'],
60 'data-code-id': codeId,
61 title: 'Preview code',
62 type: 'button'
63 },
64 children: [createRawHtmlElement(PREVIEW_ICON_SVG)]
65 };
66}
67
68function createHeader(language: string, codeId: string): Element {
69 const actions: Element[] = [createCopyButton(codeId)];
70
71 if (language.toLowerCase() === 'html') {
72 actions.push(createPreviewButton(codeId));
73 }
74
75 return {
76 type: 'element',
77 tagName: 'div',
78 properties: { className: ['code-block-header'] },
79 children: [
80 {
81 type: 'element',
82 tagName: 'span',
83 properties: { className: ['code-language'] },
84 children: [{ type: 'text', value: language }]
85 },
86 {
87 type: 'element',
88 tagName: 'div',
89 properties: { className: ['code-block-actions'] },
90 children: actions
91 }
92 ]
93 };
94}
95
96function createWrapper(header: Element, preElement: Element): Element {
97 return {
98 type: 'element',
99 tagName: 'div',
100 properties: { className: ['code-block-wrapper'] },
101 children: [header, preElement]
102 };
103}
104
105function extractLanguage(codeElement: Element): string {
106 const className = codeElement.properties?.className;
107 if (!Array.isArray(className)) return 'text';
108
109 for (const cls of className) {
110 if (typeof cls === 'string' && cls.startsWith('language-')) {
111 return cls.replace('language-', '');
112 }
113 }
114
115 return 'text';
116}
117
118/**
119 * Generates a unique code block ID using a global counter.
120 */
121function generateCodeId(): string {
122 if (typeof window !== 'undefined') {
123 return `code-${(window.idxCodeBlock = (window.idxCodeBlock ?? 0) + 1)}`;
124 }
125 // Fallback for SSR - use timestamp + random
126 return `code-${Date.now()}-${Math.random().toString(36).slice(2, 7)}`;
127}
128
129/**
130 * Rehype plugin to enhance code blocks with wrapper, header, and action buttons.
131 * This plugin wraps <pre><code> elements with a container that includes:
132 * - Language label
133 * - Copy button
134 * - Preview button (for HTML code blocks)
135 */
136export const rehypeEnhanceCodeBlocks: Plugin<[], Root> = () => {
137 return (tree: Root) => {
138 visit(tree, 'element', (node: Element, index, parent) => {
139 if (node.tagName !== 'pre' || !parent || index === undefined) return;
140
141 const codeElement = node.children.find(
142 (child): child is Element => child.type === 'element' && child.tagName === 'code'
143 );
144
145 if (!codeElement) return;
146
147 const language = extractLanguage(codeElement);
148 const codeId = generateCodeId();
149
150 codeElement.properties = {
151 ...codeElement.properties,
152 'data-code-id': codeId
153 };
154
155 const header = createHeader(language, codeId);
156 const wrapper = createWrapper(header, node);
157
158 // Replace pre with wrapper in parent
159 (parent.children as ElementContent[])[index] = wrapper;
160 });
161 };
162};