archify-vscode-ext/src/preview.ts

import path from 'node:path';
import * as vscode from 'vscode';
import { readDiagramSource } from './detect';
import { renderDiagram } from './renderer';
import { evidenceMode, log, readText, runtimeFor, viewerTheme } from './runtime';
import { handleWebviewMessage, WebviewMessage } from './webviewMessages';
import { makeNonce, messagePage, prepareDiagramHtml } from './webviewHtml';

export const PREVIEW_VIEW_TYPE = 'archify.preview';

interface PreviewState {
  uri: string;
  hash?: string;
}

class DiagramPreview implements vscode.Disposable {
  private seq = 0;
  private hash: string | undefined;
  private lastHtml: string | undefined;
  private rendered = false;
  private timer: NodeJS.Timeout | undefined;
  private readonly disposables: vscode.Disposable[] = [];
  private readonly onDisposeEmitter = new vscode.EventEmitter<void>();
  readonly onDidDispose = this.onDisposeEmitter.event;

  constructor(
    private readonly context: vscode.ExtensionContext,
    readonly panel: vscode.WebviewPanel,
    readonly uri: vscode.Uri,
    state?: PreviewState,
  ) {
    this.hash = state?.hash;
    panel.webview.options = { enableScripts: true, localResourceRoots: [] };
    panel.iconPath = vscode.Uri.joinPath(context.extensionUri, 'media', 'preview.svg');
    panel.title = `Preview ${path.basename(uri.path)}`;

    this.disposables.push(
      panel.onDidDispose(() => this.dispose()),
      panel.webview.onDidReceiveMessage((message: WebviewMessage) => {
        if (message.type === 'hash') {
          this.hash = message.hash;
          return;
        }
        void handleWebviewMessage(message, this.uri, this.panel);
      }),
    );

    if (uri.scheme === 'file') {
      const watcher = vscode.workspace.createFileSystemWatcher(
        new vscode.RelativePattern(vscode.Uri.file(path.dirname(uri.fsPath)), path.basename(uri.fsPath)),
      );
      const onDisk = () => {
        // Open documents are tracked through text change events instead.
        if (!vscode.workspace.textDocuments.some((doc) => doc.uri.toString() === uri.toString())) this.scheduleUpdate();
      };
      this.disposables.push(watcher, watcher.onDidChange(onDisk), watcher.onDidCreate(onDisk));
    }

    this.showMessage('Rendering…', `Rendering ${path.basename(uri.path)} with Archify.`);
    void this.update();
  }

  private state(): PreviewState {
    return { uri: this.uri.toString(), hash: this.hash };
  }

  private showMessage(title: string, body: string, detail?: string): void {
    this.lastHtml = undefined;
    this.panel.webview.html = messagePage(
      { cspSource: this.panel.webview.cspSource, nonce: makeNonce(), state: this.state() },
      title,
      body,
      detail,
    );
  }

  scheduleUpdate(): void {
    const delay = vscode.workspace.getConfiguration('archify', this.uri).get<number>('preview.debounceMs', 400);
    clearTimeout(this.timer);
    this.timer = setTimeout(() => void this.update(), delay);
  }

  async update(force = false): Promise<void> {
    clearTimeout(this.timer);
    const seq = ++this.seq;
    let text: string;
    try {
      text = await readText(this.uri);
    } catch (error) {
      this.showMessage('Cannot read diagram', this.uri.fsPath, error instanceof Error ? error.message : String(error));
      return;
    }
    const source = readDiagramSource(text, this.uri.path);
    if (!source) {
      this.showMessage(
        'Nothing to show yet',
        'Set "diagram_type" to architecture, workflow, sequence, dataflow or lifecycle, and make sure the file is valid JSON.',
      );
      return;
    }

    let result;
    try {
      const runtime = await runtimeFor(this.context, this.uri);
      result = await renderDiagram(runtime, { ...source, text }, evidenceMode(this.uri));
    } catch (error) {
      result = { ok: false, error: error instanceof Error ? error.message : String(error) };
    }
    if (seq !== this.seq) return; // a newer render superseded this one

    if (!result.ok || !result.html) {
      log().warn(`Render failed for ${this.uri.fsPath}: ${result.error}`);
      if (this.rendered && this.lastHtml) {
        void this.panel.webview.postMessage({
          type: 'notice',
          notice: { kind: 'error', text: `Render failed; showing the last successful render.\n${firstLines(result.error)}` },
        });
      } else {
        this.showMessage('Archify could not render this diagram', 'Fix the problems below (also listed in the Problems panel) and save.', result.error);
      }
      return;
    }

    const html = prepareDiagramHtml(result.html, {
      cspSource: this.panel.webview.cspSource,
      nonce: makeNonce(),
      theme: viewerTheme(),
      hash: this.hash,
      state: this.state(),
      notice: result.evidenceNotice ? { kind: 'warning', text: result.evidenceNotice } : undefined,
    });
    // Nonces differ per render, so compare the renderer output instead.
    const fingerprint = `${viewerTheme()}\u0000${result.evidenceNotice ?? ''}\u0000${result.html}`;
    if (!force && this.rendered && fingerprint === this.lastHtml) {
      void this.panel.webview.postMessage({ type: 'notice', notice: undefined });
      return;
    }
    this.lastHtml = fingerprint;
    this.rendered = true;
    this.panel.webview.html = html;
  }

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

function firstLines(text: string | undefined, count = 6): string {
  return (text ?? '').split('\n').slice(0, count).join('\n');
}

export class PreviewManager implements vscode.Disposable, vscode.WebviewPanelSerializer<PreviewState> {
  private readonly previews = new Map<string, DiagramPreview>();
  private active: DiagramPreview | undefined;
  private readonly disposables: vscode.Disposable[] = [];

  constructor(private readonly context: vscode.ExtensionContext) {
    this.disposables.push(
      vscode.window.registerWebviewPanelSerializer(PREVIEW_VIEW_TYPE, this),
      vscode.workspace.onDidChangeTextDocument((event) => {
        const preview = this.previews.get(event.document.uri.toString());
        if (preview && event.contentChanges.length && this.liveUpdate(preview.uri)) preview.scheduleUpdate();
      }),
      vscode.workspace.onDidSaveTextDocument((doc) => void this.previews.get(doc.uri.toString())?.update()),
      vscode.workspace.onDidChangeConfiguration((event) => {
        if (event.affectsConfiguration('archify')) this.refreshAll();
      }),
      vscode.window.onDidChangeActiveColorTheme(() => this.refreshAll()),
    );
  }

  private liveUpdate(uri: vscode.Uri): boolean {
    return vscode.workspace.getConfiguration('archify', uri).get<boolean>('preview.liveUpdate', true);
  }

  get activePreview(): DiagramPreview | undefined {
    return this.active;
  }

  show(uri: vscode.Uri, viewColumn: vscode.ViewColumn, preserveFocus = false): void {
    const existing = this.previews.get(uri.toString());
    if (existing) {
      existing.panel.reveal(viewColumn, preserveFocus);
      return;
    }
    const panel = vscode.window.createWebviewPanel(
      PREVIEW_VIEW_TYPE,
      `Preview ${path.basename(uri.path)}`,
      { viewColumn, preserveFocus },
      { enableScripts: true, retainContextWhenHidden: true, enableFindWidget: true, localResourceRoots: [] },
    );
    this.track(new DiagramPreview(this.context, panel, uri));
  }

  async deserializeWebviewPanel(panel: vscode.WebviewPanel, state: PreviewState | undefined): Promise<void> {
    if (!state?.uri) {
      panel.dispose();
      return;
    }
    this.track(new DiagramPreview(this.context, panel, vscode.Uri.parse(state.uri), state));
  }

  private track(preview: DiagramPreview): void {
    const key = preview.uri.toString();
    this.previews.get(key)?.panel.dispose();
    this.previews.set(key, preview);
    const onViewState = () => {
      if (preview.panel.active) this.active = preview;
      else if (this.active === preview) this.active = undefined;
      void vscode.commands.executeCommand('setContext', 'archifyPreviewFocus', Boolean(this.active));
    };
    preview.panel.onDidChangeViewState(onViewState);
    onViewState();
    preview.onDidDispose(() => {
      if (this.previews.get(key) === preview) this.previews.delete(key);
      if (this.active === preview) {
        this.active = undefined;
        void vscode.commands.executeCommand('setContext', 'archifyPreviewFocus', false);
      }
    });
  }

  refreshAll(force = false): void {
    for (const preview of this.previews.values()) void preview.update(force);
  }

  dispose(): void {
    for (const preview of this.previews.values()) preview.panel.dispose();
    this.previews.clear();
    for (const disposable of this.disposables.splice(0)) disposable.dispose();
  }
}