colanode/apps/desktop/src/main/services/node-service.ts

import {
  generateId,
  IdType,
  Node,
  NodeAttributes,
  NodeMutationContext,
  registry,
  ServerNodeCreateTransaction,
  ServerNodeDeleteTransaction,
  ServerNodeTransaction,
  ServerNodeUpdateTransaction,
} from '@colanode/core';
import { decodeState, YDoc } from '@colanode/crdt';
import { sql } from 'kysely';

import { SelectWorkspace } from '@/main/data/app/schema';
import { databaseService } from '@/main/data/database-service';
import {
  CreateDownload,
  CreateUpload,
  SelectDownload,
  SelectNode,
  SelectNodeTransaction,
  SelectUpload,
} from '@/main/data/workspace/schema';
import { createLogger } from '@/main/logger';
import { interactionService } from '@/main/services/interaction-service';
import {
  fetchNodeAncestors,
  mapDownload,
  mapNode,
  mapTransaction,
  mapUpload,
} from '@/main/utils';
import { eventBus } from '@/shared/lib/event-bus';

export type CreateNodeInput = {
  id: string;
  attributes: NodeAttributes;
  upload?: CreateUpload;
  download?: CreateDownload;
};

export type UpdateNodeResult =
  | 'success'
  | 'not_found'
  | 'unauthorized'
  | 'failed'
  | 'invalid_attributes';

class NodeService {
  private readonly logger = createLogger('node-service');

  public async fetchNode(nodeId: string, userId: string): Promise<Node | null> {
    const workspaceDatabase =
      await databaseService.getWorkspaceDatabase(userId);

    const nodeRow = await workspaceDatabase
      .selectFrom('nodes')
      .where('id', '=', nodeId)
      .selectAll()
      .executeTakeFirst();

    if (!nodeRow) {
      return null;
    }

    return mapNode(nodeRow);
  }

  public async createNode(
    userId: string,
    input: CreateNodeInput | CreateNodeInput[]
  ) {
    this.logger.trace(`Creating ${Array.isArray(input) ? 'nodes' : 'node'}`);
    const workspace = await this.fetchWorkspace(userId);

    const inputs = Array.isArray(input) ? input : [input];
    const createdNodes: SelectNode[] = [];
    const createdNodeTransactions: SelectNodeTransaction[] = [];
    const createdUploads: SelectUpload[] = [];
    const createdDownloads: SelectDownload[] = [];

    const workspaceDatabase =
      await databaseService.getWorkspaceDatabase(userId);

    await workspaceDatabase.transaction().execute(async (transaction) => {
      for (const inputItem of inputs) {
        const model = registry.getNodeModel(inputItem.attributes.type);
        if (!model.schema.safeParse(inputItem.attributes).success) {
          throw new Error('Invalid attributes');
        }

        let ancestors: Node[] = [];
        if (inputItem.attributes.parentId) {
          const ancestorRows = await fetchNodeAncestors(
            transaction,
            inputItem.attributes.parentId
          );
          ancestors = ancestorRows.map(mapNode);
        }

        const context = new NodeMutationContext(
          workspace.account_id,
          workspace.workspace_id,
          userId,
          workspace.role,
          ancestors
        );

        if (!model.canCreate(context, inputItem.attributes)) {
          throw new Error('Insufficient permissions');
        }

        const ydoc = new YDoc();
        const update = ydoc.updateAttributes(
          model.schema,
          inputItem.attributes
        );

        const createdAt = new Date().toISOString();
        const transactionId = generateId(IdType.Transaction);

        const createdNode = await transaction
          .insertInto('nodes')
          .returningAll()
          .values({
            id: inputItem.id,
            attributes: JSON.stringify(inputItem.attributes),
            created_at: createdAt,
            created_by: context.userId,
            transaction_id: transactionId,
          })
          .executeTakeFirst();

        if (!createdNode) {
          throw new Error('Failed to create node');
        }

        createdNodes.push(createdNode);

        const createdTransaction = await transaction
          .insertInto('node_transactions')
          .returningAll()
          .values({
            id: transactionId,
            node_id: inputItem.id,
            operation: 'create',
            node_type: inputItem.attributes.type,
            data: update,
            created_at: createdAt,
            created_by: context.userId,
            retry_count: 0,
            status: 'pending',
          })
          .executeTakeFirst();

        if (!createdTransaction) {
          throw new Error('Failed to create transaction');
        }

        createdNodeTransactions.push(createdTransaction);

        if (inputItem.upload) {
          const createdUpload = await transaction
            .insertInto('uploads')
            .returningAll()
            .values(inputItem.upload)
            .executeTakeFirst();

          if (!createdUpload) {
            throw new Error('Failed to create upload');
          }

          createdUploads.push(createdUpload);
        }

        if (inputItem.download) {
          const createdDownload = await transaction
            .insertInto('downloads')
            .returningAll()
            .values(inputItem.download)
            .executeTakeFirst();

          if (!createdDownload) {
            throw new Error('Failed to create download');
          }

          createdDownloads.push(createdDownload);
        }
      }
    });

    for (const createdNode of createdNodes) {
      this.logger.trace(
        `Created node ${createdNode.id} with type ${createdNode.type}`
      );

      eventBus.publish({
        type: 'node_created',
        userId,
        node: mapNode(createdNode),
      });
    }

    for (const createdTransaction of createdNodeTransactions) {
      this.logger.trace(
        `Created transaction ${createdTransaction.id} for node ${createdTransaction.node_id} with operation ${createdTransaction.operation}`
      );

      eventBus.publish({
        type: 'node_transaction_created',
        userId,
        transaction: mapTransaction(createdTransaction),
      });

      await interactionService.setInteraction(
        userId,
        createdTransaction.node_id,
        createdTransaction.node_type,
        'lastReceivedTransactionId',
        createdTransaction.id
      );
    }

    for (const createdUpload of createdUploads) {
      this.logger.trace(
        `Created upload ${createdUpload.upload_id} for node ${createdUpload.node_id}`
      );

      eventBus.publish({
        type: 'upload_created',
        userId,
        upload: mapUpload(createdUpload),
      });
    }

    for (const createdDownload of createdDownloads) {
      this.logger.trace(
        `Created download ${createdDownload.upload_id} for node ${createdDownload.node_id}`
      );

      eventBus.publish({
        type: 'download_created',
        userId,
        download: mapDownload(createdDownload),
      });
    }
  }

  public async updateNode(
    nodeId: string,
    userId: string,
    updater: (attributes: NodeAttributes) => NodeAttributes
  ): Promise<UpdateNodeResult> {
    let count = 0;
    while (count++ < 20) {
      const result = await this.tryUpdateNode(nodeId, userId, updater);
      if (result) {
        return result;
      }
    }

    return 'failed';
  }

  private async tryUpdateNode(
    nodeId: string,
    userId: string,
    updater: (attributes: NodeAttributes) => NodeAttributes
  ): Promise<UpdateNodeResult | null> {
    this.logger.trace(`Updating node ${nodeId}`);

    const workspace = await this.fetchWorkspace(userId);

    const workspaceDatabase =
      await databaseService.getWorkspaceDatabase(userId);

    const ancestorRows = await fetchNodeAncestors(workspaceDatabase, nodeId);
    const nodeRow = ancestorRows.find((ancestor) => ancestor.id === nodeId);
    if (!nodeRow) {
      return 'not_found';
    }

    const ancestors = ancestorRows.map(mapNode);
    const node = mapNode(nodeRow);

    if (!node) {
      return 'not_found';
    }

    const context = new NodeMutationContext(
      workspace.account_id,
      workspace.workspace_id,
      userId,
      workspace.role,
      ancestors
    );

    const transactionId = generateId(IdType.Transaction);
    const updatedAt = new Date().toISOString();
    const updatedAttributes = updater(node.attributes);

    const model = registry.getNodeModel(node.type);
    if (!model.schema.safeParse(updatedAttributes).success) {
      return 'invalid_attributes';
    }

    if (!model.canUpdate(context, node, updatedAttributes)) {
      return 'unauthorized';
    }

    const ydoc = new YDoc();
    const previousTransactions = await workspaceDatabase
      .selectFrom('node_transactions')
      .where('node_id', '=', nodeId)
      .selectAll()
      .execute();

    for (const previousTransaction of previousTransactions) {
      if (previousTransaction.data === null) {
        return 'not_found';
      }

      ydoc.applyUpdate(previousTransaction.data);
    }

    const update = ydoc.updateAttributes(model.schema, updatedAttributes);

    const { updatedNode, createdTransaction } = await workspaceDatabase
      .transaction()
      .execute(async (trx) => {
        const updatedNode = await trx
          .updateTable('nodes')
          .returningAll()
          .set({
            attributes: JSON.stringify(ydoc.getAttributes()),
            updated_at: updatedAt,
            updated_by: context.userId,
            transaction_id: transactionId,
          })
          .where('id', '=', nodeId)
          .where('transaction_id', '=', node.transactionId)
          .executeTakeFirst();

        if (updatedNode) {
          const createdTransaction = await trx
            .insertInto('node_transactions')
            .returningAll()
            .values({
              id: transactionId,
              node_id: nodeId,
              node_type: node.type,
              operation: 'update',
              data: update,
              created_at: updatedAt,
              created_by: context.userId,
              retry_count: 0,
              status: 'pending',
            })
            .executeTakeFirst();

          return { updatedNode, createdTransaction };
        }

        return {
          updatedNode: undefined,
          createdTransaction: undefined,
        };
      });

    if (updatedNode) {
      this.logger.trace(
        `Updated node ${updatedNode.id} with type ${updatedNode.type}`
      );

      eventBus.publish({
        type: 'node_updated',
        userId,
        node: mapNode(updatedNode),
      });
    } else {
      this.logger.trace(`Failed to update node ${nodeId}`);
    }

    if (createdTransaction) {
      this.logger.trace(
        `Created transaction ${createdTransaction.id} for node ${nodeId}`
      );

      eventBus.publish({
        type: 'node_transaction_created',
        userId,
        transaction: mapTransaction(createdTransaction),
      });

      await interactionService.setInteraction(
        userId,
        createdTransaction.node_id,
        createdTransaction.node_type,
        'lastReceivedTransactionId',
        createdTransaction.id
      );
    } else {
      this.logger.trace(`Failed to create transaction for node ${nodeId}`);
    }

    if (updatedNode) {
      return 'success';
    }

    return null;
  }

  public async deleteNode(nodeId: string, userId: string) {
    const workspace = await this.fetchWorkspace(userId);
    const workspaceDatabase =
      await databaseService.getWorkspaceDatabase(userId);

    const ancestorRows = await fetchNodeAncestors(workspaceDatabase, nodeId);
    const ancestors = ancestorRows.map(mapNode);

    const node = ancestors.find((ancestor) => ancestor.id === nodeId);
    if (!node) {
      throw new Error('Node not found');
    }

    const model = registry.getNodeModel(node.type);
    const context = new NodeMutationContext(
      workspace.account_id,
      workspace.workspace_id,
      userId,
      workspace.role,
      ancestors
    );

    if (!model.canDelete(context, node)) {
      throw new Error('Insufficient permissions');
    }

    const { deletedNode, createdTransaction } = await workspaceDatabase
      .transaction()
      .execute(async (trx) => {
        const deletedNode = await trx
          .deleteFrom('nodes')
          .returningAll()
          .where('id', '=', nodeId)
          .executeTakeFirst();

        if (!deletedNode) {
          return { deletedNode: undefined, createdTransaction: undefined };
        }

        await trx
          .deleteFrom('node_transactions')
          .where('node_id', '=', nodeId)
          .execute();

        await trx
          .deleteFrom('collaborations')
          .where('node_id', '=', nodeId)
          .execute();

        await trx
          .deleteFrom('interaction_events')
          .where('node_id', '=', nodeId)
          .execute();

        await trx
          .deleteFrom('interactions')
          .where('node_id', '=', nodeId)
          .execute();

        const createdTransaction = await trx
          .insertInto('node_transactions')
          .returningAll()
          .values({
            id: generateId(IdType.Transaction),
            node_id: nodeId,
            node_type: node.type,
            operation: 'delete',
            data: null,
            created_at: new Date().toISOString(),
            created_by: context.userId,
            retry_count: 0,
            status: 'pending',
          })
          .executeTakeFirst();

        return { deletedNode, createdTransaction };
      });

    if (deletedNode) {
      this.logger.trace(
        `Deleted node ${deletedNode.id} with type ${deletedNode.type}`
      );

      eventBus.publish({
        type: 'node_deleted',
        userId,
        node: mapNode(deletedNode),
      });
    } else {
      this.logger.trace(`Failed to delete node ${nodeId}`);
    }

    if (createdTransaction) {
      this.logger.trace(
        `Created transaction ${createdTransaction.id} for node ${nodeId}`
      );

      eventBus.publish({
        type: 'node_transaction_created',
        userId,
        transaction: mapTransaction(createdTransaction),
      });
    } else {
      this.logger.trace(`Failed to create transaction for node ${nodeId}`);
    }
  }

  public async applyServerTransaction(
    userId: string,
    transaction: ServerNodeTransaction
  ) {
    if (transaction.operation === 'create') {
      await this.applyServerCreateTransaction(userId, transaction);
    } else if (transaction.operation === 'update') {
      await this.applyServerUpdateTransaction(userId, transaction);
    } else if (transaction.operation === 'delete') {
      await this.applyServerDeleteTransaction(userId, transaction);
    }
  }

  public async replaceTransactions(
    userId: string,
    nodeId: string,
    transactions: ServerNodeTransaction[]
  ): Promise<boolean> {
    const workspaceDatabase =
      await databaseService.getWorkspaceDatabase(userId);

    const firstTransaction = transactions[0];
    if (!firstTransaction) {
      return false;
    }

    const lastTransaction = transactions[transactions.length - 1];
    if (!lastTransaction) {
      return false;
    }

    const ydoc = new YDoc();
    for (const transaction of transactions) {
      if (transaction.operation === 'delete') {
        await this.applyServerDeleteTransaction(userId, transaction);
        return true;
      }

      ydoc.applyUpdate(transaction.data);
    }

    const attributes = ydoc.getAttributes<NodeAttributes>();
    const attributesJson = JSON.stringify(attributes);

    await workspaceDatabase.transaction().execute(async (trx) => {
      await trx
        .insertInto('nodes')
        .values({
          id: nodeId,
          attributes: attributesJson,
          created_at: firstTransaction.createdAt,
          created_by: firstTransaction.createdBy,
          updated_at:
            firstTransaction.id !== lastTransaction.id
              ? lastTransaction.createdAt
              : null,
          updated_by:
            firstTransaction.id !== lastTransaction.id
              ? lastTransaction.createdBy
              : null,
          transaction_id: lastTransaction.id,
        })
        .onConflict((oc) =>
          oc.columns(['id']).doUpdateSet({
            attributes: attributesJson,
            updated_at: lastTransaction.createdAt,
            updated_by: lastTransaction.createdBy,
            transaction_id: lastTransaction.id,
          })
        )
        .execute();

      await trx
        .insertInto('node_transactions')
        .values(
          transactions.map((t) => ({
            id: t.id,
            node_id: t.nodeId,
            node_type: t.nodeType,
            operation: t.operation,
            data:
              t.operation !== 'delete' && t.data ? decodeState(t.data) : null,
            created_at: t.createdAt,
            created_by: t.createdBy,
            retry_count: 0,
            status: 'synced',
            version: BigInt(t.version),
            server_created_at: t.serverCreatedAt,
          }))
        )
        .onConflict((oc) =>
          oc.columns(['id']).doUpdateSet({
            status: 'synced',
            version: sql`excluded.version`,
            server_created_at: sql`excluded.server_created_at`,
          })
        )
        .execute();
    });

    return true;
  }

  private async applyServerCreateTransaction(
    userId: string,
    transaction: ServerNodeCreateTransaction
  ) {
    this.logger.trace(
      `Applying server create transaction ${transaction.id} for node ${transaction.nodeId}`
    );

    const workspaceDatabase =
      await databaseService.getWorkspaceDatabase(userId);

    const version = BigInt(transaction.version);
    const existingTransaction = await workspaceDatabase
      .selectFrom('node_transactions')
      .select(['id', 'status', 'version', 'server_created_at'])
      .where('id', '=', transaction.id)
      .executeTakeFirst();

    if (existingTransaction) {
      if (
        existingTransaction.status === 'synced' &&
        existingTransaction.version === version &&
        existingTransaction.server_created_at === transaction.serverCreatedAt
      ) {
        this.logger.trace(
          `Server create transaction ${transaction.id} for node ${transaction.nodeId} is already synced`
        );
        return;
      }

      await workspaceDatabase
        .updateTable('node_transactions')
        .set({
          status: 'synced',
          version,
          server_created_at: transaction.serverCreatedAt,
        })
        .where('id', '=', transaction.id)
        .execute();

      this.logger.trace(
        `Server create transaction ${transaction.id} for node ${transaction.nodeId} has been synced`
      );
      return;
    }

    const ydoc = new YDoc();
    ydoc.applyUpdate(transaction.data);

    const attributes = ydoc.getAttributes();

    const { createdNode } = await workspaceDatabase
      .transaction()
      .execute(async (trx) => {
        const createdNode = await trx
          .insertInto('nodes')
          .returningAll()
          .values({
            id: transaction.nodeId,
            attributes: JSON.stringify(attributes),
            created_at: transaction.createdAt,
            created_by: transaction.createdBy,
            transaction_id: transaction.id,
          })
          .executeTakeFirst();

        await trx
          .insertInto('node_transactions')
          .values({
            id: transaction.id,
            node_id: transaction.nodeId,
            node_type: transaction.nodeType,
            operation: 'create',
            data: decodeState(transaction.data),
            created_at: transaction.createdAt,
            created_by: transaction.createdBy,
            retry_count: 0,
            status: createdNode ? 'synced' : 'incomplete',
            version,
            server_created_at: transaction.serverCreatedAt,
          })
          .execute();

        return { createdNode };
      });

    if (createdNode) {
      this.logger.trace(
        `Created node ${createdNode.id} with type ${createdNode.type} with transaction ${transaction.id}`
      );

      eventBus.publish({
        type: 'node_created',
        userId,
        node: mapNode(createdNode),
      });

      await interactionService.setInteraction(
        userId,
        createdNode.id,
        createdNode.type,
        'lastReceivedTransactionId',
        transaction.id
      );
    } else {
      this.logger.trace(
        `Server create transaction ${transaction.id} for node ${transaction.nodeId} is incomplete`
      );

      eventBus.publish({
        type: 'node_transaction_incomplete',
        userId,
        transactionId: transaction.id,
      });
    }
  }

  private async applyServerUpdateTransaction(
    userId: string,
    transaction: ServerNodeUpdateTransaction
  ) {
    const workspaceDatabase =
      await databaseService.getWorkspaceDatabase(userId);

    const version = BigInt(transaction.version);
    const existingTransaction = await workspaceDatabase
      .selectFrom('node_transactions')
      .select(['id', 'status', 'version', 'server_created_at'])
      .where('id', '=', transaction.id)
      .executeTakeFirst();

    if (existingTransaction) {
      if (
        existingTransaction.status === 'synced' &&
        existingTransaction.version === version &&
        existingTransaction.server_created_at === transaction.serverCreatedAt
      ) {
        this.logger.trace(
          `Server update transaction ${transaction.id} for node ${transaction.nodeId} is already synced`
        );
        return;
      }

      await workspaceDatabase
        .updateTable('node_transactions')
        .set({
          status: 'synced',
          version,
          server_created_at: transaction.serverCreatedAt,
        })
        .where('id', '=', transaction.id)
        .execute();

      this.logger.trace(
        `Server update transaction ${transaction.id} for node ${transaction.nodeId} has been synced`
      );
      return;
    }

    const previousTransactions = await workspaceDatabase
      .selectFrom('node_transactions')
      .selectAll()
      .where('node_id', '=', transaction.nodeId)
      .orderBy('id', 'asc')
      .execute();

    const ydoc = new YDoc();

    for (const previousTransaction of previousTransactions) {
      if (previousTransaction.data) {
        ydoc.applyUpdate(previousTransaction.data);
      }
    }

    ydoc.applyUpdate(transaction.data);
    const attributes = ydoc.getAttributes();

    const { updatedNode } = await workspaceDatabase
      .transaction()
      .execute(async (trx) => {
        const updatedNode = await trx
          .updateTable('nodes')
          .returningAll()
          .set({
            attributes: JSON.stringify(attributes),
            updated_at: transaction.createdAt,
            updated_by: transaction.createdBy,
            transaction_id: transaction.id,
          })
          .where('id', '=', transaction.nodeId)
          .executeTakeFirst();

        await trx
          .insertInto('node_transactions')
          .values({
            id: transaction.id,
            node_id: transaction.nodeId,
            node_type: transaction.nodeType,
            operation: 'update',
            data: decodeState(transaction.data),
            created_at: transaction.createdAt,
            created_by: transaction.createdBy,
            retry_count: 0,
            status: updatedNode ? 'synced' : 'incomplete',
            version,
            server_created_at: transaction.serverCreatedAt,
          })
          .execute();

        return { updatedNode };
      });

    if (updatedNode) {
      this.logger.trace(
        `Updated node ${updatedNode.id} with type ${updatedNode.type} with transaction ${transaction.id}`
      );

      eventBus.publish({
        type: 'node_updated',
        userId,
        node: mapNode(updatedNode),
      });

      await interactionService.setInteraction(
        userId,
        updatedNode.id,
        updatedNode.type,
        'lastReceivedTransactionId',
        transaction.id
      );
    } else {
      this.logger.trace(
        `Server update transaction ${transaction.id} for node ${transaction.nodeId} is incomplete`
      );

      eventBus.publish({
        type: 'node_transaction_incomplete',
        userId,
        transactionId: transaction.id,
      });
    }
  }

  private async applyServerDeleteTransaction(
    userId: string,
    transaction: ServerNodeDeleteTransaction
  ) {
    this.logger.trace(
      `Applying server delete transaction ${transaction.id} for node ${transaction.nodeId}`
    );

    const workspaceDatabase =
      await databaseService.getWorkspaceDatabase(userId);

    const result = await workspaceDatabase
      .transaction()
      .execute(async (trx) => {
        await trx
          .deleteFrom('node_transactions')
          .where('node_id', '=', transaction.nodeId)
          .execute();

        await trx
          .deleteFrom('interactions')
          .where('node_id', '=', transaction.nodeId)
          .execute();

        await trx
          .deleteFrom('interaction_events')
          .where('node_id', '=', transaction.nodeId)
          .execute();

        await trx
          .deleteFrom('collaborations')
          .where('node_id', '=', transaction.nodeId)
          .execute();

        const nodeRow = await trx
          .deleteFrom('nodes')
          .returningAll()
          .where('id', '=', transaction.nodeId)
          .executeTakeFirst();

        return nodeRow;
      });

    if (result) {
      this.logger.trace(
        `Deleted node ${result.id} with type ${result.type} with transaction ${transaction.id}`
      );

      eventBus.publish({
        type: 'node_deleted',
        userId,
        node: mapNode(result),
      });
    }
  }

  async fetchWorkspace(userId: string): Promise<SelectWorkspace> {
    const workspace = await databaseService.appDatabase
      .selectFrom('workspaces')
      .selectAll()
      .where('user_id', '=', userId)
      .executeTakeFirst();

    if (!workspace) {
      throw new Error('Workspace not found');
    }

    return workspace;
  }
}

export const nodeService = new NodeService();