`nanomachine` is a type-safe state machine builder based on Nanostores and inspired by XState.

nanomachine.ts

import * as nanostores from "nanostores";
import type { MaybePromise } from "./types/MaybePromise.ts";
import type { MergeDeep } from "type-fest";

class StateMachineContextBuilder {
	context<CONTEXT_TYPE>() {
		const atom = nanostores.atom();
		return new StateMachineEventsBuilder<CONTEXT_TYPE>(atom as never);
	}
}

class StateMachineEventsBuilder<CONTEXT_TYPE> {
	private atom: nanostores.PreinitializedWritableAtom<CONTEXT_TYPE>;

	constructor(atom: nanostores.PreinitializedWritableAtom<CONTEXT_TYPE>) {
		this.atom = atom;
	}

	events<
		EVENTS extends {
			[key: string]: [] | [unknown];
		},
	>() {
		return new StateMachineStatesBuilder<
			CONTEXT_TYPE,
			{
				[key in keyof EVENTS]: EVENTS[key] extends []
					? undefined
					: EVENTS[key][0];
			}
		>(this.atom);
	}
}

class StateMachineStatesBuilder<
	CONTEXT_TYPE,
	EVENTS extends {
		[key: string]: unknown;
	},
> {
	private atom: nanostores.PreinitializedWritableAtom<CONTEXT_TYPE>;

	constructor(atom: nanostores.PreinitializedWritableAtom<CONTEXT_TYPE>) {
		this.atom = atom;
	}

	states<STATE_TYPES extends string>() {
		return new StateMachineInitialBuilder<CONTEXT_TYPE, STATE_TYPES, EVENTS>(
			this.atom,
		);
	}
}

class StateMachineStateBuilder<
	CONTEXT_TYPE,
	STATE_TYPES extends string,
	EVENTS extends {
		[key: string]: unknown;
	},
> {
	private atom: nanostores.PreinitializedWritableAtom<CONTEXT_TYPE>;
	private stateAtom: nanostores.PreinitializedWritableAtom<STATE_TYPES>;
	private event: keyof EVENTS;
	private payload: EVENTS[keyof EVENTS];
	private emit: <K extends keyof EVENTS>(
		...[event, payload]: EVENTS[K] extends undefined ? [K] : [K, EVENTS[K]]
	) => void;
	private stopPropagation = false;
	private enabledEvents: ("onReceive" | "onEntry" | "after")[] = [
		"onEntry",
		"onReceive",
		"after",
	];
	private resolve: () => void;
	private reject: (error: Error) => void;

	constructor(
		atom: nanostores.PreinitializedWritableAtom<CONTEXT_TYPE>,
		stateAtom: nanostores.PreinitializedWritableAtom<STATE_TYPES>,
		event: keyof EVENTS,
		payload: EVENTS[keyof EVENTS],
		emit: typeof this.emit,
		resolve: () => void,
		reject: (error: Error) => void,
		enabledEvents?: typeof this.enabledEvents,
	) {
		this.atom = atom;
		this.stateAtom = stateAtom;
		this.event = event;
		this.payload = payload;
		this.emit = emit;
		this.enabledEvents = enabledEvents ?? this.enabledEvents;
		this.resolve = resolve;
		this.reject = reject;
	}

	localEvent<
		EVENT_NAME extends string,
		EVENT_PAYLOAD,
	>(): StateMachineStateBuilder<
		CONTEXT_TYPE,
		STATE_TYPES,
		{
			[key in keyof EVENTS]: EVENTS[key];
		} & {
			[key in EVENT_NAME]: EVENT_PAYLOAD;
		}
	> {
		return this as never;
	}

	guard(
		evalGuard: (state: CONTEXT_TYPE) => boolean,
		fallbackState: STATE_TYPES,
	): typeof this {
		try {
			this.stopPropagation =
				this.stopPropagation || !evalGuard(this.atom.get());
			if (this.stopPropagation) this.stateAtom.set(fallbackState);
		} catch (error) {
			this.reject(new Error(`Guard evaluation failed: ${error}`));
			this.stopPropagation = true;
		}
		return this;
	}

	guardContext<K extends keyof CONTEXT_TYPE, T extends CONTEXT_TYPE[K]>(
		property: K,
		evalGuard: (state: CONTEXT_TYPE[K]) => state is T,
		fallbackState: STATE_TYPES,
	): StateMachineStateBuilder<
		// { [key in K]: T } & Omit<CONTEXT_TYPE, K>,
		MergeDeep<CONTEXT_TYPE, { [key in K]: T }>,
		STATE_TYPES,
		EVENTS
	> {
		try {
			const currentState = this.atom.get();

			this.stopPropagation =
				this.stopPropagation || !evalGuard(currentState[property]);
			if (this.stopPropagation) this.stateAtom.set(fallbackState);
		} catch (error) {
			this.reject(new Error(`Guard evaluation failed: ${error}`));
			this.stopPropagation = true;
		}

		return this as never;
	}

	onEntry(
		action: (
			state: CONTEXT_TYPE,
			set: (value: Partial<CONTEXT_TYPE>) => void,
			emit: typeof this.emit,
		) => MaybePromise<void>,
	): typeof this {
		if (!this.enabledEvents.includes("onEntry")) return this;
		if (this.stopPropagation) return this;

		// TODO: Lock to prevent changes while the state is already changed
		Promise.resolve(
			action(
				this.atom.get(),
				(value) => this.atom.set({ ...this.atom.get(), ...value }),
				this.emit.bind(this),
			),
		).catch((err) => this.reject(err));
		return this;
	}

	onReceive(
		events: {
			[key in keyof EVENTS]?: (
				...[context, set, payload]: EVENTS[key] extends undefined
					? [CONTEXT_TYPE, (value: Partial<CONTEXT_TYPE>) => void]
					: [CONTEXT_TYPE, (value: Partial<CONTEXT_TYPE>) => void, EVENTS[key]]
			) => STATE_TYPES | "$_END" | false | undefined;
		},
	): typeof this {
		if (!this.enabledEvents.includes("onReceive")) return this;
		if (this.stopPropagation) return this;

		const newState = (
			events[this.event] as (
				...args: unknown[]
			) => ReturnType<NonNullable<(typeof events)[keyof EVENTS]>> | undefined
		)?.(
			this.atom.get(),
			(value: CONTEXT_TYPE) => this.atom.set({ ...this.atom.get(), ...value }),
			this.payload,
		);

		if (newState === "$_END") {
			this.resolve();
			this.stopPropagation = true;
			return this;
		}

		if (newState) this.stateAtom.set(newState as never);
		return this;
	}

	after(timeout: number, stateType: STATE_TYPES): typeof this {
		if (!this.enabledEvents.includes("after")) return this;
		if (this.stopPropagation) return this;

		const expectedState = this.stateAtom.get();
		setTimeout(() => {
			if (this.stateAtom.get() === expectedState) {
				this.stateAtom.set(stateType);
			}
		}, timeout);
		return this;
	}
}

class StateMachineInitialBuilder<
	CONTEXT_TYPE,
	STATE_TYPES extends string,
	EVENTS extends {
		[key: string]: unknown;
	},
> {
	private atom: nanostores.PreinitializedWritableAtom<CONTEXT_TYPE>;

	constructor(atom: nanostores.PreinitializedWritableAtom<CONTEXT_TYPE>) {
		this.atom = atom;
	}

	initial(initialState: STATE_TYPES) {
		return new StateMachineFinalBuilder<CONTEXT_TYPE, STATE_TYPES, EVENTS>(
			this.atom,
			initialState,
		);
	}
}

class StateMachineFinalBuilder<
	CONTEXT_TYPE,
	STATE_TYPES extends string,
	EVENTS extends {
		[key: string]: unknown;
	},
> {
	private atom: nanostores.PreinitializedWritableAtom<CONTEXT_TYPE>;
	private initialState: STATE_TYPES;

	constructor(
		atom: nanostores.PreinitializedWritableAtom<CONTEXT_TYPE>,
		initialState: STATE_TYPES,
	) {
		this.atom = atom;
		this.initialState = initialState;
	}

	$: {
		[key in STATE_TYPES]: (
			stateBuilder: StateMachineStateBuilder<CONTEXT_TYPE, STATE_TYPES, EVENTS>,
		) => StateMachineStateBuilder<any, STATE_TYPES, any>;
	} = {} as never;

	implement(
		states: {
			[key in STATE_TYPES]: (
				stateBuilder: StateMachineStateBuilder<
					CONTEXT_TYPE,
					STATE_TYPES,
					EVENTS
				>,
			) => StateMachineStateBuilder<CONTEXT_TYPE, STATE_TYPES, EVENTS>;
		},
	) {
		return new StateMachine<CONTEXT_TYPE, STATE_TYPES, EVENTS>(
			this.atom,
			states,
			this.initialState,
		);
	}
}

class StateMachine<
	CONTEXT_TYPE,
	STATE_TYPES extends string,
	EVENTS extends {
		[key: string]: unknown;
	},
> {
	private $stateAtom: nanostores.PreinitializedWritableAtom<STATE_TYPES>;
	private started = false;
	private atom: nanostores.PreinitializedWritableAtom<CONTEXT_TYPE>;
	private states: {
		[key in STATE_TYPES]: (
			stateBuilder: StateMachineStateBuilder<CONTEXT_TYPE, STATE_TYPES, EVENTS>,
		) => void;
	};
	private resolve: () => void;
	private reject: (error: Error) => void;
	private promise: Promise<void>;

	constructor(
		atom: nanostores.PreinitializedWritableAtom<CONTEXT_TYPE>,
		states: {
			[key in STATE_TYPES]: (
				stateBuilder: StateMachineStateBuilder<
					CONTEXT_TYPE,
					STATE_TYPES,
					EVENTS
				>,
			) => void;
		},
		initialState: STATE_TYPES,
	) {
		this.$stateAtom = nanostores.atom(initialState);
		this.atom = atom;
		this.states = states;

		const { promise, resolve, reject } = Promise.withResolvers<void>();
		this.promise = promise;
		this.resolve = resolve;
		this.reject = reject;
	}

	start(context: CONTEXT_TYPE) {
		if (this.started) {
			return this;
		}
		this.started = true;
		this.atom.set(context);
		this.$stateAtom.subscribe((state) => {
			this.states[state]?.(
				new StateMachineStateBuilder(
					this.atom,
					this.$stateAtom,
					null!,
					null!,
					this.emit.bind(this),
					this.resolve,
					this.reject,
				),
			);
		});

		this.promise.catch((error) => {
			this.started = false;
		});

		return this;
	}

	get() {
		return this.atom.get();
	}

	set(value: CONTEXT_TYPE) {
		this.atom.set(value);
	}

	emit<K extends keyof EVENTS>(
		...[event, payload]: EVENTS[K] extends undefined ? [K] : [K, EVENTS[K]]
	) {
		if (!this.started) {
			return;
		}
		const state = this.$stateAtom.get();
		this.states[state]?.(
			new StateMachineStateBuilder(
				this.$atom,
				this.$stateAtom,
				event,
				payload as never,
				this.emit.bind(this),
				this.resolve,
				this.reject,
				["onReceive"],
			),
		);
	}

	subscribe(listener: (value: CONTEXT_TYPE) => void): () => void {
		return this.atom.subscribe(listener);
	}

	get $atom() {
		return this.atom;
	}

	get $promise() {
		return this.promise;
	}
}

export function createStateMachine() {
	return new StateMachineContextBuilder();
}