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

import path from 'node:path';
import * as vscode from 'vscode';
import { renderDiagram } from '../archify/cli';
import { evidenceMode, runtimeFor } from '../archify/runtime';
import { isMainFile, isRegistryFile } from '../model/format';
import { ModelPanel, panelFromIndex } from '../model/panel';
import { prepareDiagram } from '../model/prepare';
import { ModelService } from '../model/service';
import { isArchifyFile } from '../sources/detect';
import { readText } from '../util/files';
import { log } from '../util/log';
import { viewerTheme } from '../util/settings';
import { handleWebviewMessage, WebviewMessage } from '../webview/messages';
import { makeNonce, messagePage, modelPage, prepareDiagramHtml } from '../webview/page';
import { viewer3dOption, viewerRoots } from '../webview/viewer3d';

import { DIAGRAM_EDITOR_VIEW_TYPE, PREVIEW_VIEW_TYPE, REGISTRY_VIEW_TYPE } from './viewTypes';

export { DIAGRAM_EDITOR_VIEW_TYPE, PREVIEW_VIEW_TYPE };

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;
  /** Root of the model this file belongs to, once known. */
  root: vscode.Uri | undefined;

  constructor(
    private readonly context: vscode.ExtensionContext,
    private readonly models: ModelService,
    readonly panel: vscode.WebviewPanel,
    readonly uri: vscode.Uri,
    state?: PreviewState,
    /** True when this is the custom editor of the file, not a side preview. */
    readonly isEditor = false,
  ) {
    this.hash = state?.hash;
    panel.webview.options = { enableScripts: true, localResourceRoots: viewerRoots(this.context.extensionUri) };
    if (!isEditor) {
      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;
        }
        if (message.type === 'openDiagram') {
          void this.openLinked(message.path, DIAGRAM_EDITOR_VIEW_TYPE);
          return;
        }
        if (message.type === 'openRegistry') {
          void this.openLinked(message.file, REGISTRY_VIEW_TYPE);
          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,
    );
  }

  /** Opens a diagram of the model, given its path relative to the model root. */
  private async openLinked(relativePath: string, viewType: string): Promise<void> {
    if (!this.root) return;
    const target = vscode.Uri.joinPath(this.root, ...relativePath.split('/'));
    try {
      await vscode.workspace.fs.stat(target);
    } catch {
      vscode.window.showWarningMessage(`Archify: ${relativePath} does not exist in this model.`);
      return;
    }
    await vscode.commands.executeCommand('vscode.openWith', target, viewType, vscode.ViewColumn.Active);
  }

  private showModelPage(model: ModelPanel): void {
    const fingerprint = `model\u0000${JSON.stringify(model)}`;
    if (fingerprint === this.lastHtml) return;
    this.lastHtml = fingerprint;
    this.rendered = false;
    this.panel.webview.html = modelPage({ cspSource: this.panel.webview.cspSource, nonce: makeNonce(), state: this.state() }, model);
  }

  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;
    }
    // main.archify: the model page wraps the diagram, or stands alone.
    const main = isMainFile(this.uri.path) && this.uri.scheme === 'file';
    let model: ModelPanel | undefined;
    if (main) {
      this.root = vscode.Uri.file(path.dirname(this.uri.fsPath));
      try {
        model = panelFromIndex(await this.models.index(this.root));
      } catch (error) {
        log().warn(`Model index failed for ${this.root.fsPath}: ${error instanceof Error ? error.message : String(error)}`);
      }
    }
    let prepared;
    try {
      prepared = await prepareDiagram(this.models, this.uri, text);
    } catch (error) {
      log().warn(`Could not resolve ${this.uri.fsPath}: ${error instanceof Error ? error.message : String(error)}`);
    }
    if (seq !== this.seq) return;
    if (!main) this.root = prepared?.root;
    if (!prepared) {
      if (model) {
        this.showModelPage(model);
        return;
      }
      this.showMessage(
        'Nothing to show yet',
        `Set "diagram_type" to architecture, workflow, sequence, dataflow or lifecycle, and make sure the file is valid ${isArchifyFile(this.uri.path) ? 'YAML' : 'JSON'}.`,
      );
      return;
    }
    const source = prepared.source;

    let result;
    try {
      const runtime = await runtimeFor(this.context, this.uri);
      result = await renderDiagram(runtime, source, 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,
      viewer3d: viewer3dOption(this.panel.webview, this.context.extensionUri),
      links: Object.fromEntries(prepared.links),
      descriptions: Object.fromEntries(prepared.descriptions),
      model,
    });
    // Nonces differ per render, so compare the renderer output instead.
    const fingerprint = `${viewerTheme()}\u0000${result.evidenceNotice ?? ''}\u0000${JSON.stringify([model ?? null, [...prepared.links], [...prepared.descriptions]])}\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 DiagramViewManager
  implements vscode.Disposable, vscode.WebviewPanelSerializer<PreviewState>, vscode.CustomTextEditorProvider
{
  /** Side previews, one per source file. */
  private readonly previews = new Map<string, DiagramPreview>();
  /** Custom-editor tabs of .archify files (several per file are allowed). */
  private readonly editors = new Set<DiagramPreview>();
  private active: DiagramPreview | undefined;
  private readonly disposables: vscode.Disposable[] = [];

  constructor(
    private readonly context: vscode.ExtensionContext,
    private readonly models: ModelService,
  ) {
    this.disposables.push(
      // A registry or main.archify change re-renders every diagram of the model;
      // any other file change only refreshes the model page's list.
      models.onDidChangeModel(({ root, file }) => {
        const whole = isRegistryFile(file.path) || isMainFile(file.path);
        for (const preview of [...this.previews.values(), ...this.editors]) {
          if (preview.root?.toString() !== root.toString()) continue;
          if (whole || isMainFile(preview.uri.path)) preview.scheduleUpdate();
        }
      }),
      vscode.window.registerWebviewPanelSerializer(PREVIEW_VIEW_TYPE, this),
      vscode.window.registerCustomEditorProvider(DIAGRAM_EDITOR_VIEW_TYPE, this, {
        webviewOptions: { retainContextWhenHidden: true, enableFindWidget: true },
        supportsMultipleEditorsPerDocument: true,
      }),
      vscode.workspace.onDidChangeTextDocument((event) => {
        if (!event.contentChanges.length || !this.liveUpdate(event.document.uri)) return;
        for (const preview of this.forUri(event.document.uri)) preview.scheduleUpdate();
      }),
      vscode.workspace.onDidSaveTextDocument((doc) => {
        for (const preview of this.forUri(doc.uri)) void preview.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;
  }

  private forUri(uri: vscode.Uri): DiagramPreview[] {
    const key = uri.toString();
    return [...this.previews.values(), ...this.editors].filter((preview) => preview.uri.toString() === key);
  }

  /** Opening a .archify file lands here: the diagram is the editor. */
  async resolveCustomTextEditor(document: vscode.TextDocument, panel: vscode.WebviewPanel): Promise<void> {
    panel.webview.options = { enableScripts: true, localResourceRoots: viewerRoots(this.context.extensionUri) };
    this.track(new DiagramPreview(this.context, this.models, panel, document.uri, undefined, true));
  }

  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: viewerRoots(this.context.extensionUri) },
    );
    this.track(new DiagramPreview(this.context, this.models, panel, uri));
  }

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

  private track(preview: DiagramPreview): void {
    const key = preview.uri.toString();
    if (preview.isEditor) this.editors.add(preview);
    else {
      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));
      if (preview.panel.active) void vscode.commands.executeCommand('setContext', 'archify.inModel', Boolean(preview.root));
    };
    preview.panel.onDidChangeViewState(onViewState);
    onViewState();
    preview.onDidDispose(() => {
      this.editors.delete(preview);
      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(), ...this.editors]) void preview.update(force);
  }

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