import { getContext } from "../lib/context.ts";
import { TrackSystem } from "../track/system.ts";
import { Train } from "../train.old.ts";

export class StateMachine<T> {
  private _states: Map<T, State<T>> = new Map();
  private currentState?: State<T>;

  update(dt: number, ctx?: CanvasRenderingContext2D) {
    this.currentState?.update(dt, ctx);
  }

  optimizePerformance(percent: number) {
    const ctx = getContext() as { trains: Train[]; track: TrackSystem };
    if (percent < 0.5) {
      ctx.track.optimize(percent);
    }
  }

  get current() {
    return this.currentState;
  }

  get states() {
    return this._states;
  }

  addState(state: State<T>) {
    this.states.set(state.name, state);
  }

  transitionTo(state: T) {
    if (!this.current) {
      this.currentState = this._states.get(state)!;
      this.currentState.start();
      return;
    }
    if (this.current?.canTransitionTo(state) && this._states.has(state)) {
      this.current.stop();
      this.currentState = this._states.get(state)!;
      this.current.start();
    }
  }
}

export abstract class State<T> {
  protected stateMachine: StateMachine<T>;
  protected abstract validTransitions: Set<T>;

  abstract readonly name: T;

  constructor(
    stateMachine: StateMachine<T>,
  ) {
    this.stateMachine = stateMachine;
  }

  abstract update(dt: number, ctx?: CanvasRenderingContext2D): void;
  abstract start(): void;
  abstract stop(): void;

  canTransitionTo(state: T) {
    return this.validTransitions.has(state);
  }
}

export abstract class ExtensibleState<T> extends State<T> {
  extensions: Map<string, (...args: unknown[]) => void> = new Map();
  registerExtension(name: string, cb: (...args: unknown[]) => void) {
    this.extensions.set(name, cb);
  }

  constructor(stateMachine: StateMachine<T>) {
    super(stateMachine);
    const oldUpdate = this.update;
    this.update = function (dt: number, ctx?: CanvasRenderingContext2D) {
      oldUpdate.apply(this, [dt, ctx]);
      this.runExtensions(dt, ctx);
    };
  }

  runExtensions(...args: unknown[]) {
    for (const [name, cb] of this.extensions) {
      cb(...args);
    }
  }
}