vscode-cowriting-plugin/src/threadController.ts

/**
 * ThreadController — the thin editor-facing layer (spec §6.4). Wires
 * CoauthorStore + Anchorer + threadModel to vscode.comments: create-on-
 * selection, reply, resolve, render, live-range tracking, reload, and external
 * re-anchoring (SLICE-4) — driven by the SHARED sidecar watcher in extension.ts
 * via handleExternalSidecarChange (self-writes are suppressed in CoauthorStore).
 * Threads are NEVER silently moved — an unresolvable anchor renders as an
 * orphaned thread (INV-1).
 */
import * as vscode from "vscode";
import { CoauthorStore } from "./store";
import { emptyArtifact, type Artifact, type Provenance } from "./model";
import { buildFingerprint, resolve, shift, type OffsetRange } from "./anchorer";
import { addThread, appendMessage, setStatus } from "./threadModel";

/** Test-facing snapshot of what is currently rendered for a document. */
export interface RenderedThread {
  id: string;
  status: "open" | "resolved";
  orphaned: boolean;
  range: { startLine: number; endLine: number };
}

interface DocState {
  docPath: string;
  uri: vscode.Uri;
  artifact: Artifact;
  /** thread id -> vscode thread. */
  vsThreads: Map<string, vscode.CommentThread>;
  /** thread id -> live offset range (within-session optimization, INV-3). */
  live: Map<string, OffsetRange>;
  /** thread id -> orphaned flag for the current render. */
  orphaned: Map<string, boolean>;
}

export class ThreadController implements vscode.Disposable {
  private readonly controller: vscode.CommentController;
  private readonly disposables: vscode.Disposable[] = [];
  private readonly docs = new Map<string, DocState>(); // keyed by docPath

  constructor(private readonly store: CoauthorStore, private readonly rootDir: string) {
    this.controller = vscode.comments.createCommentController("cowriting.threads", "Coauthoring Threads");
    this.controller.commentingRangeProvider = {
      provideCommentingRanges: (document) => {
        if (!this.isInRoot(document.uri)) return [];
        return [new vscode.Range(0, 0, Math.max(0, document.lineCount - 1), 0)];
      },
    };
    this.disposables.push(this.controller);

    this.disposables.push(
      vscode.commands.registerCommand("cowriting.createThread", () => this.createThreadOnSelection()),
      vscode.commands.registerCommand("cowriting.reply", (r: vscode.CommentReply) => this.reply(r)),
      vscode.commands.registerCommand("cowriting.resolveThread", (t: vscode.CommentThread) =>
        this.setThreadStatus(t, "resolved"),
      ),
      vscode.commands.registerCommand("cowriting.reopenThread", (t: vscode.CommentThread) =>
        this.setThreadStatus(t, "open"),
      ),
      vscode.workspace.onDidChangeTextDocument((e) => this.onDidChange(e)),
      vscode.workspace.onDidSaveTextDocument((d) => this.onDidSave(d)),
    );
  }

  private isInRoot(uri: vscode.Uri): boolean {
    return uri.scheme === "file" && uri.fsPath.startsWith(this.rootDir);
  }

  private docPathOf(uri: vscode.Uri): string {
    return vscode.workspace.asRelativePath(uri, false);
  }

  private currentAuthor(): Provenance {
    const id = vscode.workspace.getConfiguration("git").get<string>("user.name") || process.env.USER || "human";
    return { kind: "human", id };
  }

  private persist(state: DocState): void {
    this.store.update(state.docPath, (a) => {
      a.threads = state.artifact.threads;
      for (const t of state.artifact.threads) {
        a.anchors[t.anchorId] = state.artifact.anchors[t.anchorId];
      }
    });
  }

  private ensureState(document: vscode.TextDocument): DocState {
    const docPath = this.docPathOf(document.uri);
    let state = this.docs.get(docPath);
    if (!state) {
      state = {
        docPath,
        uri: document.uri,
        artifact: this.store.load(docPath) ?? emptyArtifact(docPath),
        vsThreads: new Map(),
        live: new Map(),
        orphaned: new Map(),
      };
      this.docs.set(docPath, state);
    }
    return state;
  }

  // ---- PUC-1: create on selection -------------------------------------------------

  async createThreadOnSelection(firstBody = "New thread"): Promise<string | undefined> {
    const editor = vscode.window.activeTextEditor;
    if (!editor || editor.selection.isEmpty || !this.isInRoot(editor.document.uri)) return undefined;
    const document = editor.document;
    const state = this.ensureState(document);
    const offsets: OffsetRange = {
      start: document.offsetAt(editor.selection.start),
      end: document.offsetAt(editor.selection.end),
    };
    const fp = buildFingerprint(document.getText(), offsets);
    const { threadId } = addThread(state.artifact, fp, { author: this.currentAuthor(), body: firstBody });
    this.persist(state);
    this.renderThread(document, state, threadId, offsets, false);
    return threadId;
  }

  // ---- PUC-2: reply / resolve -----------------------------------------------------

  reply(r: vscode.CommentReply): void {
    const threadId = this.threadIdOf(r.thread);
    const state = this.stateOfThread(r.thread);
    if (!threadId || !state) return;
    appendMessage(state.artifact, threadId, { author: this.currentAuthor(), body: r.text });
    this.persist(state);
    this.refreshComments(r.thread, state, threadId);
  }

  private setThreadStatus(vsThread: vscode.CommentThread, status: "open" | "resolved"): void {
    const threadId = this.threadIdOf(vsThread);
    const state = this.stateOfThread(vsThread);
    if (!threadId || !state) return;
    setStatus(state.artifact, threadId, status);
    this.persist(state);
    vsThread.state = status === "resolved" ? vscode.CommentThreadState.Resolved : vscode.CommentThreadState.Unresolved;
    vsThread.contextValue = status;
  }

  // ---- PUC-3/4: render, reload, re-anchor -----------------------------------------

  /** Load + (re)render every thread for a document at its resolved anchor (or orphaned). */
  renderAll(document: vscode.TextDocument): void {
    const docPath = this.docPathOf(document.uri);
    const state = this.ensureState(document);
    // fresh artifact from disk (reload / external change)
    state.artifact = this.store.load(docPath) ?? emptyArtifact(docPath);
    for (const vsThread of state.vsThreads.values()) vsThread.dispose();
    state.vsThreads.clear();
    state.live.clear();
    state.orphaned.clear();
    const text = document.getText();
    for (const thread of state.artifact.threads) {
      const fp = state.artifact.anchors[thread.anchorId]?.fingerprint;
      const resolved = fp ? resolve(text, fp) : "orphaned";
      if (resolved === "orphaned") {
        const line = fp ? Math.min(fp.lineHint, Math.max(0, document.lineCount - 1)) : 0;
        const off = document.offsetAt(new vscode.Position(line, 0));
        this.renderThread(document, state, thread.id, { start: off, end: off }, true);
      } else {
        this.renderThread(document, state, thread.id, resolved, false);
      }
    }
  }

  /** Shared-watcher entry point (extension.ts): a sidecar changed externally. */
  handleExternalSidecarChange(uri: vscode.Uri): void {
    const p = uri.fsPath;
    for (const state of this.docs.values()) {
      if (this.store.sidecarPath(state.docPath) === p) {
        const doc = vscode.workspace.textDocuments.find((d) => this.docPathOf(d.uri) === state.docPath);
        if (doc) this.renderAll(doc);
      }
    }
  }

  private onDidChange(e: vscode.TextDocumentChangeEvent): void {
    const state = this.docs.get(this.docPathOf(e.document.uri));
    if (!state || state.live.size === 0) return;
    for (const change of e.contentChanges) {
      const edit = { start: change.rangeOffset, end: change.rangeOffset + change.rangeLength, newLength: change.text.length };
      for (const [id, range] of state.live) {
        const next = shift(range, edit);
        state.live.set(id, next);
        const vsThread = state.vsThreads.get(id);
        if (vsThread && !state.orphaned.get(id)) {
          vsThread.range = new vscode.Range(e.document.positionAt(next.start), e.document.positionAt(next.end));
        }
      }
    }
  }

  private onDidSave(document: vscode.TextDocument): void {
    const state = this.docs.get(this.docPathOf(document.uri));
    if (!state || state.live.size === 0) return;
    const text = document.getText();
    let changed = false;
    for (const thread of state.artifact.threads) {
      if (state.orphaned.get(thread.id)) continue;
      const range = state.live.get(thread.id);
      if (!range) continue;
      state.artifact.anchors[thread.anchorId] = { fingerprint: buildFingerprint(text, range) };
      changed = true;
    }
    if (changed) this.persist(state);
  }

  // ---- rendering helpers ----------------------------------------------------------

  private renderThread(
    document: vscode.TextDocument,
    state: DocState,
    threadId: string,
    offsets: OffsetRange,
    orphaned: boolean,
  ): void {
    const thread = state.artifact.threads.find((t) => t.id === threadId)!;
    const range = new vscode.Range(document.positionAt(offsets.start), document.positionAt(offsets.end));
    const vsThread = this.controller.createCommentThread(
      document.uri,
      range,
      thread.messages.map((m) => this.toComment(m.body, m.author.id)),
    );
    vsThread.label = orphaned ? "⚠ Orphaned thread (anchor not found)" : undefined;
    vsThread.contextValue = orphaned ? "orphaned" : thread.status;
    vsThread.state = thread.status === "resolved" ? vscode.CommentThreadState.Resolved : vscode.CommentThreadState.Unresolved;
    vsThread.collapsibleState = vscode.CommentThreadCollapsibleState.Collapsed;
    state.vsThreads.set(threadId, vsThread);
    state.live.set(threadId, offsets);
    state.orphaned.set(threadId, orphaned);
  }

  private refreshComments(vsThread: vscode.CommentThread, state: DocState, threadId: string): void {
    const thread = state.artifact.threads.find((t) => t.id === threadId)!;
    vsThread.comments = thread.messages.map((m) => this.toComment(m.body, m.author.id));
  }

  private toComment(body: string, authorName: string): vscode.Comment {
    return {
      body: new vscode.MarkdownString(body),
      mode: vscode.CommentMode.Preview,
      author: { name: authorName },
    };
  }

  private threadIdOf(vsThread: vscode.CommentThread): string | undefined {
    for (const state of this.docs.values()) {
      for (const [id, t] of state.vsThreads) if (t === vsThread) return id;
    }
    return undefined;
  }

  private stateOfThread(vsThread: vscode.CommentThread): DocState | undefined {
    for (const state of this.docs.values()) {
      for (const t of state.vsThreads.values()) if (t === vsThread) return state;
    }
    return undefined;
  }

  // ---- test-facing surface --------------------------------------------------------

  getRendered(docPath: string): RenderedThread[] {
    const state = this.docs.get(docPath);
    if (!state) return [];
    const out: RenderedThread[] = [];
    for (const [id, vsThread] of state.vsThreads) {
      const t = state.artifact.threads.find((x) => x.id === id)!;
      const r = vsThread.range;
      out.push({
        id,
        status: t.status,
        orphaned: !!state.orphaned.get(id),
        range: { startLine: r ? r.start.line : 0, endLine: r ? r.end.line : 0 },
      });
    }
    return out;
  }

  dispose(): void {
    for (const d of this.disposables) d.dispose();
  }
}