archify-vscode-ext/src/views/htmlViewer.ts

import path from 'node:path';
import * as vscode from 'vscode';
import { isArchifyHtml } from '../sources/detect';
import { viewerTheme } from '../util/settings';
import { handleWebviewMessage, WebviewMessage } from '../webview/messages';
import { makeNonce, messagePage, prepareDiagramHtml } from '../webview/page';
import { viewer3dOption, viewerRoots } from '../webview/viewer3d';

export const HTML_VIEW_TYPE = 'archify.htmlViewer';

class HtmlDocument implements vscode.CustomDocument {
  constructor(readonly uri: vscode.Uri) {}
  dispose(): void {}
}

/** Read-only custom editor that shows rendered Archify HTML files. */
export class HtmlViewerProvider implements vscode.CustomReadonlyEditorProvider<HtmlDocument>, vscode.Disposable {
  private readonly viewers = new Map<vscode.WebviewPanel, () => Promise<void>>();
  private active: { uri: vscode.Uri; panel: vscode.WebviewPanel } | undefined;
  private readonly disposables: vscode.Disposable[];

  constructor(private readonly extensionUri: vscode.Uri) {
    this.disposables = [
      vscode.window.registerCustomEditorProvider(HTML_VIEW_TYPE, this, {
        webviewOptions: { retainContextWhenHidden: true, enableFindWidget: true },
        supportsMultipleEditorsPerDocument: true,
      }),
      vscode.window.onDidChangeActiveColorTheme(() => this.reloadAll()),
      vscode.workspace.onDidChangeConfiguration((event) => {
        if (event.affectsConfiguration('archify.theme')) this.reloadAll();
      }),
    ];
  }

  /** File shown in the focused Archify HTML viewer, if any. */
  get activeUri(): vscode.Uri | undefined {
    return this.active?.uri;
  }

  dispose(): void {
    for (const disposable of this.disposables.splice(0)) disposable.dispose();
  }

  openCustomDocument(uri: vscode.Uri): HtmlDocument {
    return new HtmlDocument(uri);
  }

  async resolveCustomEditor(document: HtmlDocument, panel: vscode.WebviewPanel): Promise<void> {
    const uri = document.uri;
    let hash: string | undefined;
    panel.webview.options = { enableScripts: true, localResourceRoots: viewerRoots(this.extensionUri) };

    const reload = async () => {
      let text: string;
      try {
        text = new TextDecoder('utf-8').decode(await vscode.workspace.fs.readFile(uri));
      } catch (error) {
        panel.webview.html = messagePage(
          { cspSource: panel.webview.cspSource, nonce: makeNonce() },
          'Cannot open file',
          uri.fsPath,
          error instanceof Error ? error.message : String(error),
        );
        return;
      }
      panel.webview.html = prepareDiagramHtml(text, {
        cspSource: panel.webview.cspSource,
        nonce: makeNonce(),
        theme: viewerTheme(),
        hash,
        notice: isArchifyHtml(text)
          ? undefined
          : { kind: 'info', text: 'This file was not generated by Archify. It is shown in a sandbox without network access.' },
        viewer3d: isArchifyHtml(text) ? viewer3dOption(panel.webview, this.extensionUri) : undefined,
      });
    };

    this.viewers.set(panel, reload);

    const disposables: vscode.Disposable[] = [
      panel.webview.onDidReceiveMessage((message: WebviewMessage) => {
        if (message.type === 'hash') {
          hash = message.hash;
          return;
        }
        void handleWebviewMessage(message, uri, panel);
      }),
      panel.onDidChangeViewState(() => {
        if (panel.active) this.active = { uri, panel };
        else if (this.active?.panel === panel) this.active = undefined;
      }),
    ];
    if (uri.scheme === 'file') {
      const watcher = vscode.workspace.createFileSystemWatcher(
        new vscode.RelativePattern(vscode.Uri.file(path.dirname(uri.fsPath)), path.basename(uri.fsPath)),
      );
      let timer: NodeJS.Timeout | undefined;
      const onChange = () => {
        clearTimeout(timer);
        timer = setTimeout(() => void reload(), 300);
      };
      disposables.push(watcher, watcher.onDidChange(onChange), watcher.onDidCreate(onChange), {
        dispose: () => clearTimeout(timer),
      });
    }
    panel.onDidDispose(() => {
      this.viewers.delete(panel);
      if (this.active?.panel === panel) this.active = undefined;
      for (const disposable of disposables) disposable.dispose();
    });
    if (panel.active) this.active = { uri, panel };
    await reload();
  }

  reloadAll(): void {
    for (const reload of this.viewers.values()) void reload();
  }
}