NeitzelFXDesigner/designer-intellij/src/main/java/de/neitzel/fxdesigner/intellij/IntellijDocuments.java

package de.neitzel.fxdesigner.intellij;

import com.intellij.openapi.application.*;
import com.intellij.openapi.command.WriteCommandAction;
import com.intellij.openapi.command.undo.UndoManager;
import com.intellij.openapi.editor.*;
import com.intellij.openapi.editor.event.*;
import com.intellij.openapi.fileEditor.*;
import com.intellij.openapi.fileEditor.impl.text.TextEditorProvider;
import com.intellij.openapi.project.Project;
import com.intellij.openapi.util.Disposer;
import com.intellij.openapi.vfs.*;
import de.neitzel.fxdesigner.contracts.*;
import de.neitzel.fxdesigner.host.*;
import java.net.URI;
import java.nio.file.*;
import java.util.*;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicLong;

/**
 * Adapts actual IntelliJ Documents and native commands/history; never creates a competing undo
 * store.
 */
public final class IntellijDocuments implements HostServices.Documents, AutoCloseable {
  private final java.util.concurrent.CopyOnWriteArrayList<java.util.function.Consumer<URI>>
      changes = new java.util.concurrent.CopyOnWriteArrayList<>();

  public AutoCloseable onChange(java.util.function.Consumer<URI> listener) {
    changes.add(listener);
    return () -> changes.remove(listener);
  }

  private final Project project;
  private final com.intellij.openapi.Disposable lifetime;
  private final Map<Path, State> states = new HashMap<>();
  private final Set<Handle> handles = new HashSet<>();
  private boolean closed;

  private static final class State {
    VirtualFile file;
    Document document;
    String disk;
    AtomicLong revision = new AtomicLong();
  }

  public IntellijDocuments(Project project) {
    this.project = Objects.requireNonNull(project);
    lifetime = this::disposeResources;
    Disposer.register(project, lifetime);
    EditorFactory.getInstance()
        .getEventMulticaster()
        .addDocumentListener(
            new DocumentListener() {
              public void documentChanged(DocumentEvent event) {
                var file = FileDocumentManager.getInstance().getFile(event.getDocument());
                if (file != null && file.isInLocalFileSystem())
                  changes.forEach(listener -> listener.accept(canonicalUri(file)));
              }
            },
            lifetime);
  }

  private <T> CompletionStage<T> edt(Callable<T> action) {
    var future = new CompletableFuture<T>();
    Runnable task =
        () ->
            WriteIntentReadAction.run(
                () -> {
                  try {
                    if (closed || project.isDisposed() || Disposer.isDisposed(lifetime))
                      throw new IllegalStateException("Designer host disposed");
                    future.complete(action.call());
                  } catch (Throwable e) {
                    future.completeExceptionally(e);
                  }
                });
    if (ApplicationManager.getApplication().isDispatchThread()) task.run();
    else ApplicationManager.getApplication().invokeLater(task, ModalityState.nonModal());
    return future;
  }

  public CompletionStage<DocumentSession> open(URI uri) {
    return edt(
        () -> {
          Path supplied = Path.of(uri).toAbsolutePath().normalize();
          Path path = supplied.toRealPath();
          State state = states.get(path);
          if (state == null) {
            state = new State();
            state.file = LocalFileSystem.getInstance().refreshAndFindFileByNioFile(supplied);
            if (state.file == null) throw new IllegalArgumentException("Missing file: " + path);
            state.document = FileDocumentManager.getInstance().getDocument(state.file);
            if (state.document == null) throw new IllegalArgumentException("Not a text document");
            state.disk = Files.readString(path);
            var observed = state;
            state.document.addDocumentListener(
                new DocumentListener() {
                  @Override
                  public void documentChanged(DocumentEvent e) {
                    observed.revision.incrementAndGet();
                  }
                },
                lifetime);
            states.put(path, state);
          }
          var handle = new Handle(state);
          handles.add(handle);
          return handle;
        });
  }

  public CompletionStage<List<DocumentSnapshot>> snapshots() {
    return edt(
        () -> {
          unsaved().toCompletableFuture().join();
          for (var file : FileEditorManager.getInstance(project).getOpenFiles())
            if (file.isInLocalFileSystem() && "fxml".equalsIgnoreCase(file.getExtension()))
              open(Path.of(file.getPath()).toUri()).toCompletableFuture().join().close();
          return states.values().stream().map(this::snapshot).toList();
        });
  }

  public CompletionStage<List<DocumentSnapshot>> unsaved() {
    return edt(
        () -> {
          var result = new ArrayList<DocumentSnapshot>();
          var manager = FileDocumentManager.getInstance();
          for (var document : manager.getUnsavedDocuments()) {
            var file = manager.getFile(document);
            if (file == null || !file.isInLocalFileSystem()) continue;
            var handle =
                (Handle) open(Path.of(file.getPath()).toUri()).toCompletableFuture().join();
            try {
              result.add(snapshot(handle.s));
            } finally {
              handle.release();
            }
          }
          return List.copyOf(result);
        });
  }

  private static URI canonicalUri(VirtualFile file) {
    try {
      return Path.of(file.getPath()).toRealPath().toUri();
    } catch (java.io.IOException e) {
      return Path.of(file.getPath()).toAbsolutePath().normalize().toUri();
    }
  }

  private DocumentSnapshot snapshot(State s) {
    return new DocumentSnapshot(
        canonicalUri(s.file),
        s.revision.get(),
        s.document.getText(),
        s.document.isWritable() && s.file.isWritable(),
        FileDocumentManager.getInstance().isDocumentUnsaved(s.document));
  }

  private final class Handle implements DocumentSession {
    private final State s;
    private boolean released;
    private Editor editor;

    Handle(State s) {
      this.s = s;
    }

    private <T> CompletionStage<T> run(Callable<T> action) {
      return edt(
          () -> {
            if (released) throw new IllegalStateException("Document handle closed");
            return action.call();
          });
    }

    public CompletionStage<DocumentSnapshot> snapshot() {
      return run(() -> IntellijDocuments.this.snapshot(s));
    }

    public CompletionStage<DocumentSnapshot> apply(EditBatch edits) {
      return run(
          () -> {
            edits.applyTo(IntellijDocuments.this.snapshot(s));
            var ranges = new ArrayList<>(edits.replacements());
            ranges.sort(Comparator.comparingInt(EditBatch.Replacement::start));
            Collections.reverse(ranges);
            WriteCommandAction.runWriteCommandAction(
                project,
                edits.name(),
                java.util.UUID.randomUUID().toString(),
                () -> {
                  for (var range : ranges)
                    s.document.replaceString(range.start(), range.end(), range.text());
                });
            return IntellijDocuments.this.snapshot(s);
          });
    }

    public CompletionStage<DocumentSnapshot> undo() {
      return history(false);
    }

    public CompletionStage<DocumentSnapshot> redo() {
      return history(true);
    }

    private CompletionStage<DocumentSnapshot> history(boolean redo) {
      return run(
          () -> {
            if (!s.document.isWritable() || !s.file.isWritable())
              throw new IllegalStateException("Read-only document");
            if (editor == null)
              editor = EditorFactory.getInstance().createEditor(s.document, project);
            var textEditor = TextEditorProvider.getInstance().getTextEditor(editor);
            var undo = UndoManager.getInstance(project);
            if (redo) {
              if (undo.isRedoAvailable(textEditor)) undo.redo(textEditor);
            } else if (undo.isUndoAvailable(textEditor)) undo.undo(textEditor);
            return IntellijDocuments.this.snapshot(s);
          });
    }

    public CompletionStage<DocumentSnapshot> save(long revision) {
      return run(
          () -> {
            if (s.revision.get() != revision)
              throw new ConcurrentModificationException("Stale save");
            if (!s.document.isWritable() || !s.file.isWritable())
              throw new IllegalStateException("Read-only document");
            Path file = Path.of(s.file.getPath());
            String onDisk = Files.readString(file);
            if (!FileDocumentManager.getInstance().isDocumentUnsaved(s.document)
                && onDisk.equals(s.document.getText())) {
              s.disk = onDisk;
              return IntellijDocuments.this.snapshot(s);
            }
            if (!onDisk.equals(s.disk))
              throw new ConcurrentModificationException("External change; draft retained");
            FileDocumentManager.getInstance().saveDocument(s.document);
            if (FileDocumentManager.getInstance().isDocumentUnsaved(s.document))
              throw new IllegalStateException("IDE did not save document");
            s.disk = Files.readString(file);
            return IntellijDocuments.this.snapshot(s);
          });
    }

    private void release() {
      if (released) return;
      released = true;
      if (editor != null) EditorFactory.getInstance().releaseEditor(editor);
      handles.remove(this);
    }

    public void close() {
      edt(
          () -> {
            release();
            return null;
          });
    }
  }

  private void disposeResources() {
    for (var h : List.copyOf(handles)) h.release();
    closed = true;
    states.clear();
  }

  public void close() {
    edt(
        () -> {
          Disposer.dispose(lifetime);
          return null;
        });
  }
}