import { assign } from './actions';
import type {
  ActionObject,
  AssignAction,
  Assigner,
  PropertyAssigner,
  ExtractEvent,
  EventObject
} from './types';

export interface ContextModel<TContext, TEvent extends EventObject> {
  initialContext: TContext;
  actions: {
    [key: string]: ActionObject<TContext, TEvent>;
  };
  assign: <TEventType extends TEvent['type'] = TEvent['type']>(
    assigner:
      | Assigner<TContext, ExtractEvent<TEvent, TEventType>>
      | PropertyAssigner<TContext, ExtractEvent<TEvent, TEventType>>,
    eventType?: TEventType
  ) => AssignAction<TContext, ExtractEvent<TEvent, TEventType>>;
}

export function createModel<TContext, TEvent extends EventObject>(
  initialContext: TContext
): ContextModel<TContext, TEvent> {
  const model: ContextModel<TContext, TEvent> = {
    initialContext,
    actions: {},
    assign
  };

  return model;
}

export function assertEvent<
  TEvent extends EventObject,
  TEventType extends TEvent['type']
>(event: TEvent, type: TEventType): event is TEvent & { type: TEventType } {
  return event.type === type;
}

import {
  ActionObject,
  AssignAction,
  Assigner,
  PropertyAssigner,
  assign
} from '.';
import { EventObject } from './types';

export interface ContextModel<TC, TE extends EventObject> {
  context: TC;
  actions: {
    [key: string]: ActionObject<TC, TE>;
  };
  assign: <TEventType extends TE['type'] = TE['type']>(
    assigner:
      | Assigner<TC, TE & { type: TEventType }>
      | PropertyAssigner<TC, TE & { type: TEventType }>,
    eventType?: TEventType
  ) => AssignAction<TC, TE & { type: TEventType }>;
}

export function createModel<TContext, TEvent extends EventObject>(
  initialState: TContext
): ContextModel<TContext, TEvent> {
  const model: ContextModel<TContext, TEvent> = {
    context: initialState,
    actions: {},
    assign
  };

  return model;
}

export function assertEvent<TE extends EventObject, TType extends TE['type']>(
  event: TE,
  type: TType
): event is TE & { type: TType } {
  return event.type === type;
}

--- 
+++ 
@@ -1,30 +1,31 @@
-import {
+import { assign } from './actions';
+import type {
   ActionObject,
   AssignAction,
   Assigner,
   PropertyAssigner,
-  assign
-} from '.';
-import { EventObject } from './types';
+  ExtractEvent,
+  EventObject
+} from './types';
 
-export interface ContextModel<TC, TE extends EventObject> {
-  context: TC;
+export interface ContextModel<TContext, TEvent extends EventObject> {
+  initialContext: TContext;
   actions: {
-    [key: string]: ActionObject<TC, TE>;
+    [key: string]: ActionObject<TContext, TEvent>;
   };
-  assign: <TEventType extends TE['type'] = TE['type']>(
+  assign: <TEventType extends TEvent['type'] = TEvent['type']>(
     assigner:
-      | Assigner<TC, TE & { type: TEventType }>
-      | PropertyAssigner<TC, TE & { type: TEventType }>,
+      | Assigner<TContext, ExtractEvent<TEvent, TEventType>>
+      | PropertyAssigner<TContext, ExtractEvent<TEvent, TEventType>>,
     eventType?: TEventType
-  ) => AssignAction<TC, TE & { type: TEventType }>;
+  ) => AssignAction<TContext, ExtractEvent<TEvent, TEventType>>;
 }
 
 export function createModel<TContext, TEvent extends EventObject>(
-  initialState: TContext
+  initialContext: TContext
 ): ContextModel<TContext, TEvent> {
   const model: ContextModel<TContext, TEvent> = {
-    context: initialState,
+    initialContext,
     actions: {},
     assign
   };
@@ -32,9 +33,9 @@
   return model;
 }
 
-export function assertEvent<TE extends EventObject, TType extends TE['type']>(
-  event: TE,
-  type: TType
-): event is TE & { type: TType } {
+export function assertEvent<
+  TEvent extends EventObject,
+  TEventType extends TEvent['type']
+>(event: TEvent, type: TEventType): event is TEvent & { type: TEventType } {
   return event.type === type;
 }
Rename context -> initialContext, use ExtractEvent
