ecs-observable/src/bt/runner.ts

import type { World, Entity } from "../index";
import { query } from "../query";
import {
  Task,
  Scheduled,
  Running,
  Succeeded,
  Failed,
  Cancelled,
  TERMINAL_TAGS,
  ChildOf,
} from "./task";

// ── Types ─────────────────────────────────────────────
/** Callback invoked for each leaf task that becomes Scheduled. */
export type LeafHandler = (world: World, entity: Entity) => void;

/** Callback invoked when a task reaches a terminal status. */
export type TerminalHandler = (
  world: World,
  entity: Entity,
  status: "succeeded" | "failed" | "cancelled",
) => void;

// ── Helpers ───────────────────────────────────────────
function terminalStatus(
  world: World,
  entity: Entity,
): "succeeded" | "failed" | "cancelled" | null {
  if (world.has(entity, Succeeded)) return "succeeded";
  if (world.has(entity, Failed)) return "failed";
  if (world.has(entity, Cancelled)) return "cancelled";
  return null;
}

function isTerminal(world: World, entity: Entity): boolean {
  return terminalStatus(world, entity) !== null;
}

function clearStatus(world: World, entity: Entity): void {
  for (const tag of TERMINAL_TAGS) {
    if (world.has(entity, tag)) world.remove(entity, tag);
  }
  if (world.has(entity, Running)) world.remove(entity, Running);
}

/** Recursively clear status from an entity and all its descendants. */
function clearSubtree(world: World, entity: Entity): void {
  clearStatus(world, entity);
  for (const child of childrenOf(world, entity)) {
    clearSubtree(world, child);
  }
}

function childrenOf(world: World, parent: Entity): Entity[] {
  return world.getRelatedTo(parent, ChildOf);
}

function parentOf(world: World, child: Entity): Entity | null {
  return world.getRelated(child, ChildOf) ?? null;
}

// ── TaskRunner ────────────────────────────────────────
/**
 * Push-based behaviour-tree runner.
 *
 * Only tasks tagged with `Scheduled` are processed each tick.
 * When a task finishes, it notifies its parent, which may schedule
 * the next child (sequential), aggregate results (parallel), or
 * propagate upward.
 *
 * Leaves are dispatched to a user-provided `onLeaf` callback.
 * Terminal results are dispatched to a user-provided `onTerminal` callback.
 *
 * @example
 * ```ts
 * const runner = new TaskRunner(world);
 * runner.onLeaf = (world, entity) => {
 *   // do the leaf's work, then call:
 *   runner.succeed(entity);
 * };
 *
 * // Each frame:
 * runner.tick();
 * ```
 */
export class TaskRunner {
  private _world: World;

  /** Called when a leaf task becomes Scheduled. */
  onLeaf: LeafHandler = () => {};

  /** Called when any task reaches a terminal status. */
  onTerminal: TerminalHandler = () => {};

  constructor(world: World) {
    this._world = world;
  }

  // ── Public API ────────────────────────────────────

  /**
   * Process all Scheduled tasks.
   *
   * Call once per frame. Only entities with `Scheduled` are touched.
   */
  tick(): void {
    const scheduled = [...this._world.query(query(Task, Scheduled))];
    for (const entity of scheduled) {
      this._world.remove(entity, Scheduled);
      this._execute(entity);
    }
  }

  /** Mark a leaf task as succeeded and propagate upward. */
  succeed(entity: Entity): void {
    this._finish(entity, Succeeded);
  }

  /** Mark a leaf task as failed and propagate upward. */
  fail(entity: Entity): void {
    this._finish(entity, Failed);
  }

  /** Cancel a task and all its descendants. */
  cancel(entity: Entity): void {
    this._cancelTree(entity);
  }

  /** Schedule a task for execution next tick. */
  schedule(entity: Entity): void {
    if (this._world.has(entity, Task)) {
      this._world.add(entity, Scheduled);
    }
  }

  /** Reset a task to idle (removes all status tags). */
  reset(entity: Entity): void {
    clearStatus(this._world, entity);
    this._world.remove(entity, Scheduled);
  }

  // ── Internal execution ────────────────────────────

  private _execute(entity: Entity): void {
    const t = this._world.get(entity, Task);

    switch (t.kind) {
      case "leaf":
        this._executeLeaf(entity);
        break;
      case "sequential":
        this._executeSequential(entity);
        break;
      case "parallel":
        this._executeParallel(entity);
        break;
      case "random":
        this._executeRandom(entity);
        break;
      case "repeat":
        this._executeRepeat(entity);
        break;
      case "selector":
        this._executeSelector(entity);
        break;
    }
  }

  private _executeLeaf(entity: Entity): void {
    this._world.add(entity, Running);
    this.onLeaf(this._world, entity);
  }

  private _executeSequential(entity: Entity): void {
    const children = childrenOf(this._world, entity);

    // Find the first non-terminal child
    for (const child of children) {
      if (isTerminal(this._world, child)) {
        const status = terminalStatus(this._world, child)!;
        if (status === "failed" || status === "cancelled") {
          this._finish(entity, status === "cancelled" ? Cancelled : Failed);
          return;
        }
        // succeeded — continue to next child
        continue;
      }

      // Found a child that hasn't run yet — schedule it
      this._world.add(child, Scheduled);
      return;
    }

    // All children succeeded
    this._finish(entity, Succeeded);
  }

  private _executeParallel(entity: Entity): void {
    const children = childrenOf(this._world, entity);
    let allDone = true;

    for (const child of children) {
      if (isTerminal(this._world, child)) {
        const status = terminalStatus(this._world, child)!;
        if (status === "failed" || status === "cancelled") {
          this._finish(entity, status === "cancelled" ? Cancelled : Failed);
          return;
        }
        // succeeded — this child is done
        continue;
      }

      allDone = false;

      // Schedule if not already running or scheduled
      if (
        !this._world.has(child, Running) &&
        !this._world.has(child, Scheduled)
      ) {
        this._world.add(child, Scheduled);
      }
    }

    if (allDone) {
      this._finish(entity, Succeeded);
    }
  }

  private _executeRandom(entity: Entity): void {
    const children = childrenOf(this._world, entity);

    // Check if any child already reached a terminal
    for (const child of children) {
      if (isTerminal(this._world, child)) {
        const status = terminalStatus(this._world, child)!;
        this._finish(
          entity,
          status === "succeeded"
            ? Succeeded
            : status === "cancelled"
              ? Cancelled
              : Failed,
        );
        return;
      }
    }

    // Pick a random child that isn't already running
    const eligible = children.filter(
      (c) => !this._world.has(c, Running) && !this._world.has(c, Scheduled),
    );

    if (eligible.length > 0) {
      const pick = eligible[Math.floor(Math.random() * eligible.length)];
      this._world.add(pick, Scheduled);
    }
    // If no eligible children (all running), wait for one to finish
  }

  private _executeRepeat(entity: Entity): void {
    const children = childrenOf(this._world, entity);

    // Repeat expects exactly one child
    if (children.length === 0) return;

    const child = children[0];

    // If child reached a terminal, reset the entire subtree and schedule again
    if (isTerminal(this._world, child)) {
      clearSubtree(this._world, child);
      this._world.add(child, Scheduled);
      return;
    }

    // Schedule child if not already running or scheduled
    if (
      !this._world.has(child, Running) &&
      !this._world.has(child, Scheduled)
    ) {
      this._world.add(child, Scheduled);
    }
    // Repeat itself never terminates — it just keeps the child going
  }

  private _executeSelector(entity: Entity): void {
    const children = childrenOf(this._world, entity);

    // Find the first non-terminal child
    for (const child of children) {
      if (isTerminal(this._world, child)) {
        const status = terminalStatus(this._world, child)!;
        if (status === "succeeded") {
          // First success — selector succeeds
          this._finish(entity, Succeeded);
          return;
        }
        // failed or cancelled — continue to next child
        continue;
      }

      // Found a child that hasn't run yet — schedule it
      this._world.add(child, Scheduled);
      return;
    }

    // All children failed
    this._finish(entity, Failed);
  }

  // ── Completion propagation ────────────────────────

  private _finish(
    entity: Entity,
    tag: typeof Succeeded | typeof Failed | typeof Cancelled,
  ): void {
    if (!this._world.has(entity, Task)) return;

    clearStatus(this._world, entity);
    this._world.add(entity, tag);

    const status = terminalStatus(this._world, entity)!;
    this.onTerminal(this._world, entity, status);

    // Notify parent
    const parent = parentOf(this._world, entity);
    if (parent) {
      this._world.add(parent, Scheduled);
    }
  }

  private _cancelTree(entity: Entity): void {
    if (!this._world.has(entity, Task)) return;

    // Cancel children first
    for (const child of childrenOf(this._world, entity)) {
      this._cancelTree(child);
    }

    // Cancel this node
    clearStatus(this._world, entity);
    this._world.add(entity, Cancelled);
    this.onTerminal(this._world, entity, "cancelled");

    // Notify parent
    const parent = parentOf(this._world, entity);
    if (parent) {
      this._world.add(parent, Scheduled);
    }
  }
}