/*---------------------------------------------------------------------------------------------
 *  Copyright (c) Microsoft Corporation. All rights reserved.
 *  Licensed under the MIT License. See License.txt in the project root for license information.
 *--------------------------------------------------------------------------------------------*/
import * as vscode from 'vscode';
import { generateUuid } from './uuid';

/**
 * View type that uniquely identifies the Mermaid chat output renderer.
 */
const viewType = 'vscode.chatMermaidDiagram';

/**
 * Mime type used to identify Mermaid diagram data in chat output.
 */
const mime = 'text/vnd.mermaid';

export function activate(context: vscode.ExtensionContext) {

	// Register tools
	context.subscriptions.push(
		vscode.lm.registerTool<{ markup: string }>('renderMermaidDiagram', {
			invoke: async (options, _token) => {
				const sourceCode = options.input.markup;
				return writeMermaidToolOutput(sourceCode);
			},
		})
	);

	// Register the chat output renderer for Mermaid diagrams.
	// This will be invoked with the data generated by the tools.
	// It can also be invoked when rendering old Mermaid diagrams in the chat history.
	context.subscriptions.push(
		vscode.chat.registerChatOutputRenderer(viewType, {
			async renderChatOutput({ value }, webview, _ctx, _token) {
				const mermaidSource = new TextDecoder().decode(value);

				// Set the options for the webview
				const mediaRoot = vscode.Uri.joinPath(context.extensionUri, 'chat-webview-out');
				webview.options = {
					enableScripts: true,
					localResourceRoots: [mediaRoot],
				};

				// Set the HTML content for the webview
				const nonce = generateUuid();
				const mermaidScript = vscode.Uri.joinPath(mediaRoot, 'index.js');

				webview.html = `
					<!DOCTYPE html>
					<html lang="en">

					<head>
						<meta charset="UTF-8">
						<meta name="viewport" content="width=device-width, initial-scale=1.0">
						<title>Mermaid Diagram</title>
						<meta http-equiv="Content-Security-Policy" content="default-src 'none'; script-src ${webview.cspSource} 'nonce-${nonce}'; style-src 'self' 'unsafe-inline';" />
					</head>

					<body>
						<pre class="mermaid">
							${escapeHtmlText(mermaidSource)}
						</pre>

						<script type="module" nonce="${nonce}" src="${webview.asWebviewUri(mermaidScript)}"></script>
					</body>
					</html>`;
			},
		}));
}


function writeMermaidToolOutput(sourceCode: string): vscode.LanguageModelToolResult {
	// Expose the source code as a tool result for the LM
	const result = new vscode.LanguageModelToolResult([
		new vscode.LanguageModelTextPart(sourceCode)
	]);

	// And store custom data in the tool result details to indicate that a custom renderer should be used for it.
	// In this case we just store the source code as binary data.

	// Add cast to use proposed API
	(result as vscode.ExtendedLanguageModelToolResult2).toolResultDetails2 = {
		mime,
		value: new TextEncoder().encode(sourceCode),
	};

	return result;
}

function escapeHtmlText(str: string): string {
	return str
		.replace(/&/g, '&amp;')
		.replace(/</g, '&lt;')
		.replace(/>/g, '&gt;')
		.replace(/"/g, '&quot;')
		.replace(/'/g, '&#39;');
}


