mnote/wolai-frontend/src/lib/tree-stream/tree-delta.ts

import type { SidebarInitialData } from "@/components/sidebar/types";
import type { SidebarDatasetListQueryResult } from "@/lib/sidebar-data";
import {
  buildSidebarDatasetListQueryResult,
  mapSidebarDatasetListQueryResultToInitialData,
} from "@/lib/sidebar-data";
import type { DocumentRecord } from "@/lib/documents";
import type { WorkspaceSummary } from "@/lib/workspaces";
import type { MediaAsset } from "@/types/media";

export type TreeStreamDeltaOp =
  | "upsert_document"
  | "remove_document"
  | "replace_documents"
  | "replace_sidebar";

export type TreeStreamDeltaEvent = {
  op: TreeStreamDeltaOp;
  node?: DocumentRecord | null;
  document?: DocumentRecord | null;
  documentId?: string | null;
  documents?: DocumentRecord[] | null;
  sidebar?: SidebarDatasetListQueryResult | SidebarInitialData | null;
};

function cloneSidebarData(data: SidebarInitialData): SidebarInitialData {
  return {
    ...data,
    workspaces: [...data.workspaces],
    documents: [...data.documents],
    kernelFileTreeProjection: {
      ...data.kernelFileTreeProjection,
      items: [...data.kernelFileTreeProjection.items],
      edges: [...data.kernelFileTreeProjection.edges],
    },
    kernelSidebarProjection: {
      ...data.kernelSidebarProjection,
      items: [...data.kernelSidebarProjection.items],
      edges: [...data.kernelSidebarProjection.edges],
    },
    kernelSidebarTree: [...data.kernelSidebarTree],
    trashedDocuments: [...data.trashedDocuments],
    trashedMediaAssets: [...(data.trashedMediaAssets ?? [])],
    trashedMindmapAssets: [...(data.trashedMindmapAssets ?? [])],
    trashedTableAssets: [...(data.trashedTableAssets ?? [])],
    mindmapDocs: [...(data.mindmapDocs ?? [])],
    mindmapAssets: [...(data.mindmapAssets ?? [])],
    mindmapAssetChildren: { ...(data.mindmapAssetChildren ?? {}) },
    tableAssets: [...(data.tableAssets ?? [])],
    mediaAssets: [...(data.mediaAssets ?? [])],
  };
}

function buildSidebarFromDocuments(input: {
  base: SidebarInitialData;
  documents: DocumentRecord[];
}): SidebarInitialData {
  const queryResult = buildSidebarDatasetListQueryResult({
    activeWorkspaceId: input.base.activeWorkspaceId,
    workspaces: input.base.workspaces as WorkspaceSummary[],
    documents: input.documents,
    trashedDocuments: input.base.trashedDocuments,
    mindmaps: (input.base.mindmapAssets ?? []).map((asset) => ({
      mindmap_id: asset.id,
      workspace_id: asset.workspace_id,
      document_id: asset.document_id,
      created_at: asset.created_at,
      updated_at: asset.updated_at,
      deleted_at: asset.deleted_at ?? null,
      deleted_by: asset.deleted_by ?? null,
    })),
    mediaAssets: input.base.mediaAssets as MediaAsset[] | null,
    trashedMediaAssets: input.base.trashedMediaAssets as MediaAsset[] | null,
    tables: (input.base.tableAssets ?? []).map((asset) => ({
      id: asset.id,
      workspace_id: asset.workspace_id,
      document_id: asset.document_id,
      title: asset.file_name ?? null,
      created_at: asset.created_at,
      updated_at: asset.updated_at,
      deleted_at: asset.deleted_at ?? null,
      deleted_by: asset.deleted_by ?? null,
      purged_at: asset.purged_at ?? null,
      is_archived: Boolean(asset.deleted_at),
    })),
  });

  return mapSidebarDatasetListQueryResultToInitialData(queryResult);
}

function normalizeUpsertDocument(event: TreeStreamDeltaEvent): DocumentRecord | null {
  const candidate = event.node ?? event.document ?? null;
  return candidate && typeof candidate === "object" ? candidate : null;
}

function normalizeDocumentId(event: TreeStreamDeltaEvent): string | null {
  const candidate = typeof event.documentId === "string" ? event.documentId.trim() : "";
  return candidate || null;
}

export function applyTreeStreamDelta(
  base: SidebarInitialData,
  event: TreeStreamDeltaEvent,
): SidebarInitialData {
  if (event.op === "replace_sidebar" && event.sidebar) {
    if ("activeWorkspaceId" in event.sidebar) {
      return cloneSidebarData(event.sidebar as SidebarInitialData);
    }
    return mapSidebarDatasetListQueryResultToInitialData(event.sidebar as SidebarDatasetListQueryResult);
  }

  if (event.op === "replace_documents" && Array.isArray(event.documents)) {
    return buildSidebarFromDocuments({
      base,
      documents: [...event.documents],
    });
  }

  if (event.op === "upsert_document") {
    const nextDocument = normalizeUpsertDocument(event);
    if (!nextDocument) {
      return base;
    }
    const nextDocuments = [...base.documents];
    const existingIndex = nextDocuments.findIndex((item) => item.id === nextDocument.id);
    if (existingIndex >= 0) {
      nextDocuments[existingIndex] = nextDocument;
    } else {
      nextDocuments.push(nextDocument);
    }
    return buildSidebarFromDocuments({
      base,
      documents: nextDocuments,
    });
  }

  if (event.op === "remove_document") {
    const documentId = normalizeDocumentId(event);
    if (!documentId) {
      return base;
    }
    const removedIds = new Set<string>([documentId]);
    let changed = true;
    while (changed) {
      changed = false;
      for (const document of base.documents) {
        if (document.parent_id && removedIds.has(document.parent_id) && !removedIds.has(document.id)) {
          removedIds.add(document.id);
          changed = true;
        }
      }
    }
    return buildSidebarFromDocuments({
      base,
      documents: base.documents.filter((item) => !removedIds.has(item.id)),
    });
  }

  return base;
}