ecs-observable/src/observable/observe.ts

import { Subject } from "rxjs";
import type { Query } from "../query";
import type { Entity } from "../entity";
import type { WorldEvent, QueryUpdate, RelationshipUpdate } from "./events";
import type { RelationshipDef } from "../relationship";
import type { ComponentDef } from "../component";

// ── Internal state ───────────────────────────────────
interface QueryObserverState {
  query: Query;
  matched: Set<Entity>;
  subject: Subject<QueryUpdate>;
}

interface RelationshipObserverState {
  rel: RelationshipDef;
  edges: Set<string>;
  subject: Subject<RelationshipUpdate>;
}

// ── Observable layer ─────────────────────────────────
export class ObservableLayer {
  readonly events$ = new Subject<WorldEvent>();

  private _observers: QueryObserverState[] = [];
  private _relObservers: RelationshipObserverState[] = [];

  // ── Observer index: component key → observers that care ──
  private _compIndex = new Map<symbol, Set<QueryObserverState>>();

  // ── Query observers ─────────────────────────────

  observe(query: Query): Subject<QueryUpdate> {
    const existing = this._observers.find(
      (o) => o.query === query || queriesEqual(o.query, query),
    );
    if (existing) return existing.subject;

    const state: QueryObserverState = {
      query,
      matched: new Set(),
      subject: new Subject<QueryUpdate>(),
    };
    this._observers.push(state);
    this._indexObserver(state);
    return state.subject;
  }

  seed(query: Query, entities: Entity[]): void {
    const obs = this._observers.find(
      (o) => o.query === query || queriesEqual(o.query, query),
    );
    if (!obs) return;
    for (const e of entities) obs.matched.add(e);
  }

  // ── Relationship observers ───────────────────────

  observeRelated(rel: RelationshipDef): Subject<RelationshipUpdate> {
    const existing = this._relObservers.find((o) => o.rel._key === rel._key);
    if (existing) return existing.subject;

    const state: RelationshipObserverState = {
      rel,
      edges: new Set(),
      subject: new Subject<RelationshipUpdate>(),
    };
    this._relObservers.push(state);
    return state.subject;
  }

  seedRelated(
    rel: RelationshipDef,
    edges: { source: Entity; target: Entity }[],
  ): void {
    const obs = this._relObservers.find((o) => o.rel._key === rel._key);
    if (!obs) return;
    for (const { source, target } of edges) {
      obs.edges.add(edgeKey(source, target));
    }
  }

  // ── Event dispatch ───────────────────────────────

  onEvent(
    event: WorldEvent,
    queryMatches: (query: Query, e: Entity) => boolean,
  ): void {
    this.events$.next(event);

    this._dispatchToObservers(event, queryMatches);

    for (const o of this._relObservers) {
      this._updateRelObserver(o, event);
    }
  }

  // ── Private: observer indexing ───────────────────

  /** Add an observer to the component index. */
  private _indexObserver(state: QueryObserverState): void {
    for (const def of state.query.with) {
      this._addToIndex(def._key, state);
    }
    for (const def of state.query.not) {
      this._addToIndex(def._key, state);
    }
    // Also index under a well-known symbol for spawn/destroy events
    // (those always fan out to all observers).
    this._addToIndex(ANY_KEY, state);
  }

  /** Remove an observer from the component index. */
  private _unindexObserver(state: QueryObserverState): void {
    for (const def of state.query.with) {
      this._remFromIndex(def._key, state);
    }
    for (const def of state.query.not) {
      this._remFromIndex(def._key, state);
    }
    this._remFromIndex(ANY_KEY, state);
  }

  private _addToIndex(key: symbol, state: QueryObserverState): void {
    let set = this._compIndex.get(key);
    if (!set) {
      set = new Set();
      this._compIndex.set(key, set);
    }
    set.add(state);
  }

  private _remFromIndex(key: symbol, state: QueryObserverState): void {
    const set = this._compIndex.get(key);
    if (!set) return;
    set.delete(state);
    if (set.size === 0) this._compIndex.delete(key);
  }

  /** Dispatch to only the relevant observers. */
  private _dispatchToObservers(
    event: WorldEvent,
    queryMatches: (query: Query, e: Entity) => boolean,
  ): void {
    if (!("entity" in event)) return;

    // Determine which component keys are relevant
    let keys: symbol[] = [];

    switch (event.type) {
      case "spawned":
      case "destroyed":
        // Check all observers (via the ANY_KEY set)
        keys = [ANY_KEY];
        break;

      case "componentAdded":
      case "componentRemoved":
      case "componentChanged":
        keys = [event.component._key];
        break;
    }

    // Collect unique observers to update (deduplicate across keys)
    const seen = new Set<QueryObserverState>();
    for (const key of keys) {
      const set = this._compIndex.get(key);
      if (!set) continue;
      for (const o of set) {
        if (!seen.has(o)) {
          seen.add(o);
          this._updateObserver(o, event, queryMatches);
        }
      }
    }
  }

  // ── Private: observer update logic ───────────────

  private _updateObserver(
    obs: QueryObserverState,
    event: WorldEvent,
    queryMatches: (query: Query, e: Entity) => boolean,
  ): void {
    const e = event.entity!;
    const wasMatched = obs.matched.has(e);
    const nowMatches = queryMatches(obs.query, e);

    switch (event.type) {
      case "spawned":
        break;

      case "destroyed":
        if (wasMatched) {
          obs.matched.delete(e);
          obs.subject.next({ added: [], removed: [e], changed: [] });
        }
        break;

      case "componentAdded":
      case "componentRemoved": {
        if (wasMatched && !nowMatches) {
          obs.matched.delete(e);
          obs.subject.next({ added: [], removed: [e], changed: [] });
        } else if (!wasMatched && nowMatches) {
          obs.matched.add(e);
          obs.subject.next({ added: [e], removed: [], changed: [] });
        }
        break;
      }

      case "componentChanged": {
        if (wasMatched && nowMatches) {
          obs.subject.next({ added: [], removed: [], changed: [e] });
        }
        break;
      }
    }
  }

  private _updateRelObserver(
    obs: RelationshipObserverState,
    event: WorldEvent,
  ): void {
    switch (event.type) {
      case "relationshipAdded": {
        if (event.relationship._key !== obs.rel._key) break;
        const key = edgeKey(event.source, event.target);
        if (obs.edges.has(key)) break;
        obs.edges.add(key);
        obs.subject.next({
          added: [{ source: event.source, target: event.target }],
          removed: [],
        });
        break;
      }

      case "relationshipRemoved": {
        if (event.relationship._key !== obs.rel._key) break;
        const key = edgeKey(event.source, event.target);
        if (!obs.edges.has(key)) break;
        obs.edges.delete(key);
        obs.subject.next({
          added: [],
          removed: [{ source: event.source, target: event.target }],
        });
        break;
      }

      case "destroyed": {
        const removed: { source: Entity; target: Entity }[] = [];
        for (const key of obs.edges) {
          const [si, ti] = key.split(":").map(Number);
          const idx = event.entity & 0xfffff;
          if (si === idx || ti === idx) {
            removed.push({ source: si as Entity, target: ti as Entity });
          }
        }
        for (const r of removed) {
          obs.edges.delete(edgeKey(r.source, r.target));
        }
        if (removed.length > 0) {
          obs.subject.next({ added: [], removed });
        }
        break;
      }
    }
  }

  // ── Teardown ─────────────────────────────────────

  reset(): void {
    for (const o of this._observers) {
      o.subject.complete();
      o.matched.clear();
    }
    this._observers = [];
    this._compIndex.clear();

    for (const o of this._relObservers) {
      o.subject.complete();
      o.edges.clear();
    }
    this._relObservers = [];
  }

  complete(): void {
    this.events$.complete();
    for (const o of this._observers) o.subject.complete();
    this._observers = [];
    this._compIndex.clear();
    for (const o of this._relObservers) o.subject.complete();
    this._relObservers = [];
  }
}

// ── Helpers ─────────────────────────────────────

/** Sentinel key for observers that must be notified on spawn/destroy. */
const ANY_KEY = Symbol("any");

function queriesEqual(a: Query, b: Query): boolean {
  if (a.with.length !== b.with.length) return false;
  if (a.not.length !== b.not.length) return false;
  return (
    a.with.every((c, i) => c === b.with[i]) &&
    a.not.every((c, i) => c === b.not[i])
  );
}

function edgeKey(source: Entity, target: Entity): string {
  return `${source & 0xfffff}:${target & 0xfffff}`;
}