import {EventEmitter} from 'node:events';
import {AsyncLocalStorage} from 'node:async_hooks';
import type {Bot} from '@wp2508/mineflayer';

type Cancellation = {generation: number; events: EventEmitter};
const scopes = new AsyncLocalStorage<{bot: Bot; cancellation: Cancellation}>();
const states = new WeakMap<Bot, Cancellation & {scopes: Set<Cancellation>}>();
function state(bot: Bot) {
    let value = states.get(bot);
    if (!value) states.set(bot, value = {generation: 0, events: new EventEmitter(), scopes: new Set()});
    return value;
}

function current(bot: Bot): Cancellation {
    const scope = scopes.getStore();
    return scope?.bot === bot ? scope.cancellation : state(bot);
}

export function cancellation_generation(bot: Bot): number {
    return current(bot).generation;
}

export function on_function_cancel(bot: Bot, listener: () => void): () => void {
    const {events} = current(bot);
    events.on('cancel', listener);
    return () => { events.removeListener('cancel', listener); };
}

/** Existing helpers inherit the call's cancellation scope across awaits. */
export function create_function_cancellation(bot: Bot) {
    const value = state(bot);
    const cancellation: Cancellation = {generation: 0, events: new EventEmitter()};
    value.scopes.add(cancellation);
    return {
        run<T>(operation: () => T): T { return scopes.run({bot, cancellation}, operation); },
        cancel() {
            cancellation.generation++;
            cancellation.events.emit('cancel');
        },
        dispose() { value.scopes.delete(cancellation); },
    };
}

export function cancel_functions(bot: Bot): void {
    const value = state(bot);
    value.generation++;
    value.events.emit('cancel');
    for (const cancellation of value.scopes) {
        cancellation.generation++;
        cancellation.events.emit('cancel');
    }
}