import { EventObject } from 'xstate';

interface StateSchema<TEvents extends EventObject, TAllStates extends string> {
  states?: { [K in TAllStates]: StateNode<TEvents, TAllStates> };
}
interface StateNode<
  TEvents extends EventObject,
  TAllStates extends string,
  TContext = any
> extends StateSchema<TEvents, TAllStates> {
  context: TContext;
  transitions: Array<Transition<TEvents, TAllStates>>;
}
interface Transition<TEvents extends EventObject, TAllStates extends string> {
  to: TAllStates;
  event: TEvents;
}

// I'm sorry Anders... this isn't your fault
type ContextMapFromStateSchema<
  TStateSchema extends StateSchema<any, any>
> = TStateSchema['states'] extends Record<infer Keys, any>
  ? { [K in Keys]: TStateSchema['states'][K]['context'] }
  : never;

type ContextUnionFromStateSchema<
  TStateSchema extends StateSchema<any, any>
> = TStateSchema['states'] extends Record<any, infer TStateNode>
  ? TStateNode extends { context: infer TCtx }
    ? TCtx
    : never
  : never;

type TransitionConfig<T extends StateSchema<any, any>> = {
  [K in keyof T['states']]: {
    // oh god oh no
    on: {
      [E in T['states'][K]['transitions'][number]['event']['type']]: ActionFunction<
        ContextMapFromStateSchema<T>[K],
        Extract<T['states'][K]['transitions'][number]['event'], { type: E }>,
        ContextMapFromStateSchema<T>[Extract<
          T['states'][K]['transitions'][number],
          { event: { type: E } }
        >['to']]
      >
    };
  }
};

type ActionFunction<TContext, TEvent extends EventObject, TReturnContext> = (
  ctx: TContext,
  event: TEvent
) => TReturnContext;

export enum States {
  EVEN = 'EVEN',
  ODD = 'ODD',
}

export enum EventTypes {
  CLICK = 'CLICK',
  DOUBLE_CLICK = 'DOUBLE_CLICK',
  CLEAR = 'CLEAR',
}

interface ClickEvent {
  type: EventTypes.CLICK;
}

interface DoubleClickEvent {
  type: EventTypes.DOUBLE_CLICK;
}

interface ExampleSchema {
  states: {
    [States.EVEN]: {
      context: {
        value: 'EVEN';
      };
      transitions: [
        { to: States.ODD; event: ClickEvent },
        { to: States.EVEN; event: DoubleClickEvent }
      ];
    };
    [States.ODD]: {
      context: {
        value: 'ODD';
      };
      transitions: [
        { to: States.EVEN; event: ClickEvent },
        { to: States.ODD; event: DoubleClickEvent }
      ];
    };
  };
}

export const config: TransitionConfig<ExampleSchema> = {
  [States.ODD]: {
    on: {
      CLICK: (_ctx /* ctx is { value: 'ODD' } */, _event) => {
        return { value: 'EVE' };
      },
      DOUBLE_CLICK: (_ctx, _event) => {
        return { value: 'ODD' };
      },
    },
  },
  [States.EVEN]: {
    on: {
      CLICK: (_ctx /* ctx is { value: 'EVEN' }, */, _event) => {
        return { value: 'ODD' };
      },
    },
  },
};