xo-stack/packages/stack/src/storage/activities.ts

import { BaseStore } from './base-store.js';
import { StoreComposite } from './store-composite.js';

import { EventEmitter, UUIDV7 } from '../utils/index.js';

// Make the given property in a type optional.
// TODO: This is an extremely useful utility. If we use it elsewhere, move to own "type-utils.ts" file.
type Optional<T, K extends keyof T> = Omit<T, K> & Partial<Pick<T, K>>;

// Base Activity interface
export interface Activity {
  /** Key (Type) of this activity. */
  key: string;

  /** Value for this activity. */
  value: unknown;

  /** Timestamp for this activity. */
  timestamp: bigint;

  /** Entity ID may refer to a transaction id or other specific identifier. */
  entityId?: string;
}

// Define the snapshot structure
export type SnapshotState = {
  [key: string]: {
    [entityId: string]: Activity;
  };
};

export type SnapshotData<T extends Activity = Activity> = {
  [keyOrEntityId: string]: T;
};

export type SnapshotIndexedByKey<T extends Activity = Activity> = {
  [key: string]: { [entityId: string]: T };
};

export type SnapshotIndexedByEntityId<T extends Activity = Activity> = {
  [entityId: string]: { [key: string]: T };
};

// Policy type definition
export type PolicyFunction = (
  existingActivity: Activity | undefined,
  newActivity: Activity,
) => Activity;

export type ActivitiesEvents = {
  snapshotUpdated: undefined;
  updated: { activity: Activity };
  [key: `updated:${string}`]: { activity: Activity };
};

/**
 * Implements an event sourcing pattern for managing application state
 * with a simplified interface focused on the current snapshot.
 *
 * @example
 * ```typescript
 * // Create Activities manager with a store implementation
 * const activities = await Activities.fromLocalStore(localStore);
 *
 * // Register a custom policy for wallet name conflicts
 * activities.registerPolicy('WalletName', (existing, newActivity) => {
 *   // Keep the most recent wallet name
 *   return existing?.timestamp > newActivity.timestamp ? existing : newActivity;
 * });
 *
 * // Add a new wallet name activity
 * const walletName: WalletNameInput = {
 *   name: 'WalletName',
 *   value: 'Test Wallet'
 * };
 * const savedActivity = await activities.add(walletName);
 * ```
 */
export class Activities extends EventEmitter<ActivitiesEvents> {
  // Dependencies.
  public readonly store: StoreComposite;

  // State.
  // TODO: If you can refactor this around your new Storage mechanism, things become a lot simpler:
  //       1. You don't need to define such a clunky snapshot format because you have indexes.
  //       2. Caching of snapshots likely becomes possible (because you can use a storage adapter).
  public snapshot: SnapshotState = {};

  // Indexes.
  // TODO: We want to be able to remove these in favor of the TODO note above.
  // public snapshotIndexedByEntity: SnapshotIndexedByEntityId = {};
  // public snapshotIndexedByKey: SnapshotIndexedByKey = {};

  protected policies = new Map<string, PolicyFunction>();
  protected storeKeyMap = new Map<string, string>();

  /**
   * Create an Activities instance from the given composite store.
   */
  static async from(stores?: Array<BaseStore>) {
    const compositeStore = await StoreComposite.from(stores);
    const activities = new Activities(compositeStore);
    await activities.compileSnapshot();
    return activities;
  }

  protected constructor(store: StoreComposite) {
    super();

    // Assign our dependencies.
    this.store = store;

    // Register callbacks.
    this.store.on(
      'set',
      async (payload: { key: string; value: Activity | null }) => {
        if (payload.value !== null) {
          await this.processActivity(payload.key, payload.value);
        }
      },
    );
  }

  // TODO: Additional argument to "add": Something like a quarum.
  //       But, instead of a quorum, it's a function that takes the resolve of all "set" promises.
  //       That way, we can implement custom storage policies?

  /**
   * Add a new activity to the store
   * @param activity The activity to add (id and timestamp are optional)
   * @returns The stored activity with all fields populated
   */
  public async add<T extends Activity>(
    activity: Optional<T, 'entityId' | 'timestamp'>,
  ): Promise<Activity> {
    // Generate timestamp if not provided
    const timestamp = activity.timestamp ?? BigInt(Date.now());

    // Generate a UUID based on the timestamp
    const uuid = UUIDV7.fromTimestamp(timestamp).toString();

    // If this doesn't reference a specific entity id, put it at root level (empty string: "")
    const id = activity.entityId ?? '';

    const completeActivity: Activity = {
      key: activity.key,
      value: activity.value,
      entityId: id,
      timestamp,
    };

    await this.store.set(uuid, completeActivity);

    return completeActivity;
  }

  public async remove<T extends Activity>(
    activity: Pick<T, 'key' | 'entityId'>,
  ) {
    // Generate timestamp.
    const timestamp = BigInt(Date.now());

    // Generate a UUID based on the timestamp
    const uuid = UUIDV7.fromTimestamp(timestamp).toString();

    // If this doesn't reference a specific entity id, put it at root level (empty string: "")
    const id = activity.entityId ?? '';

    const removedActivity: Activity = {
      key: activity.key,
      value: null,
      entityId: id,
      timestamp,
    };

    await this.store.set(uuid, removedActivity);

    return removedActivity;
  }

  public getSnapshot() {
    return this.snapshot;
  }

  /**
   * Get a specific activity by type and id with type inference
   */
  public get<T extends Activity>(
    type: T['key'],
    entityId: string = '',
  ): T | undefined {
    const activity = this.snapshot[type]?.[entityId];
    return activity as T | undefined;
  }

  public getByEntity(entityId: string = '') {
    return Object.keys(this.snapshot).reduce(
      (result, key) => {
        const activity = this.snapshot[key][entityId];
        if (activity) {
          result[key] = activity;
        }
        return result;
      },
      {} as Record<string, Activity>,
    );
  }

  /**
   * Get all activities of a specific type
   */
  public getByKey<T extends Activity>(key: T['key']): SnapshotData<T> {
    const activities = this.snapshot[key];
    if (!activities) {
      return {};
    }

    // NOTE: We want to return a new object, not a reference to our current object.
    //       But we only need to "clone" the first level as each activity itself is immutable.
    return { ...activities } as SnapshotData<T>;
  }

  /*
  public getByEntityId<T extends Activity>(
    entityId: string = '',
  ): SnapshotData<T> {
    const activities = this.snapshotIndexedByEntity[entityId];

    if (!activities) {
      return {};
    }

    // NOTE: We want to return a new object, not a reference to our current object.
    //       But we only need to "clone" the first level as each activity itself is immutable.
    return { ...activities } as SnapshotData<T>;
  }
  */

  /**
   * Register a policy for handling conflicts between activities of the same type
   * @param type Activity type to register policy for
   * @param policyFn Function that decides which activity to keep
   */
  public registerPolicy(type: string, policyFn: PolicyFunction): void {
    this.policies.set(type, policyFn);
  }

  /**
   * Process a single activity and update the snapshot
   */
  private async processActivity(
    storeKey: string,
    activity: Activity,
  ): Promise<void> {
    if (!activity) {
      return;
    }
    // Ensure timestamp exists by extracting from UUID if needed
    if (!activity.timestamp) {
      activity.timestamp = UUIDV7.from(storeKey).getTimestamp();
    }

    const snapshotKey = this.getSnapshotKey(activity);
    const existingActivity = this.getActivityFromSnapshot(
      activity.key,
      activity.entityId,
    );

    // Apply policy to determine which activity to keep
    const keptActivity = this.applyPolicy(existingActivity, activity);

    // Update snapshot with the kept activity
    this.setActivityInSnapshot(keptActivity);

    if (keptActivity === activity) {
      // If we kept the new activity, update the store key mapping
      const oldStoreKey = this.storeKeyMap.get(snapshotKey);
      if (oldStoreKey && oldStoreKey !== storeKey) {
        // Delete the old activity from the store
        // await this.store.set(oldStoreKey, null);
      }
      this.storeKeyMap.set(snapshotKey, storeKey);

      // Emit the updated activity.
      this.emit('updated', { activity });

      // Emit the updated activity.
      this.emit(`updated:${activity.key}`, { activity });

      // Emit an event indicating our snapshot has been updated.
      this.emit('snapshotUpdated', undefined);
    } else if (storeKey !== this.storeKeyMap.get(snapshotKey)) {
      // If we kept the existing activity, delete the new activity from the store
      // await this.store.set(storeKey, null);
    }
  }

  /**
   * Compile a full snapshot from scratch
   * This happens automatically on initialization
   */
  private async compileSnapshot(): Promise<SnapshotState> {
    // Clear existing snapshot
    this.snapshot = {};
    this.storeKeyMap.clear();

    // Get all activity data from our store
    const activities = await this.store.all<Activity>();
    const entries = Object.entries(activities);

    // Filter for null payloads
    const filteredEntries = entries.filter(([_, payload]) => payload);

    // Sort activities by timestamp using UUID v7's embedded timestamp
    const activitiesSorted = filteredEntries.sort(([keyA], [keyB]) => {
      const timestampA = UUIDV7.from(keyA).getTimestamp();
      const timestampB = UUIDV7.from(keyB).getTimestamp();
      return Number(timestampA - timestampB);
    });

    // Process each activity in order
    for (const [storeKey, activity] of activitiesSorted) {
      await this.processActivity(storeKey, activity);
    }

    this.emit('snapshotUpdated', undefined);

    return this.snapshot;
  }

  /**
   * Get an activity from the snapshot by type and id
   */
  private getActivityFromSnapshot(
    type: string,
    entityId: string = '',
  ): Activity | undefined {
    return this.snapshot[type]?.[entityId];
  }

  /**
   * Set an activity in the snapshot, creating necessary object structure
   */
  private setActivityInSnapshot(activity: Activity): void {
    // Assign the key for convenience.
    const key = activity.key;

    // Assign the entityId to the provided entityId - or to the root ("") entity if none provided.
    const entityId = activity.entityId || '';

    if (!this.snapshot[key]) {
      this.snapshot[key] = {};
    }

    this.snapshot[key][entityId] = activity;

    /*
    // Index the activity by entityId:key.
    if (!this.snapshotIndexedByEntity[entityId]) {
      this.snapshot[entityId] = {};
    }

    this.snapshotIndexedByEntity[entityId][key] = activity;

    // Index the activity by key:entityId.
    if (!this.snapshotIndexedByKey[key]) {
      this.snapshot[key] = {};
    }

    this.snapshotIndexedByKey[key][entityId] = activity;
    */
  }

  /**
   * Get the snapshot key for an activity
   */
  private getSnapshotKey(activity: Activity): string {
    return `${activity.key}:${activity.entityId}`;
  }

  /**
   * Apply policy to determine which activity should be kept
   */
  private applyPolicy(
    existingActivity: Activity | undefined,
    newActivity: Activity,
  ): Activity {
    const policy = this.policies.get(newActivity.key);

    if (!policy) {
      // Default policy: compare timestamps
      if (!existingActivity) return newActivity;
      return existingActivity.timestamp >= newActivity.timestamp
        ? existingActivity
        : newActivity;
    }

    return policy(existingActivity, newActivity);
  }
}