cogsbox-state 0.5.476 → 0.5.479

package/src/plugins.ts

@@ -1,14 +1,71 @@
 import { z } from 'zod';
 import type React from 'react';
 import { StateObject, UpdateTypeDetail } from './CogsState';
-import { getGlobalStore } from './store';
+import {
+  getGlobalStore,
+  getAllFieldElements,
+  setAllFieldsDisabled,
+} from './store';
 import { ClientActivityEvent } from './pluginStore';
+import { RefObject } from 'react';
 
 type Prettify<T> = { [K in keyof T]: T[K] } & {};
 
+export type ChainMethodTarget =
+  | 'any'
+  | 'array'
+  | 'object'
+  | 'primitive'
+  | 'boolean';
+
+export type ChainMethodContext<TOptions = any, THookReturn = any> = {
+  stateKey: string;
+  path: string[];
+  pluginName: string;
+  options: TOptions | undefined;
+  hookData?: THookReturn;
+  $get: () => any;
+  $update: (payload: any) => { synced: () => void };
+  $applyOperation: (operation: any, metaData?: Record<string, any>) => void;
+  getFieldMetaData: () => any;
+  setFieldMetaData: (data: Record<string, any>) => void;
+  removeFieldMetaData: () => void;
+  getFieldRefs: () => RefObject<any>[];
+  getFieldElements: () => HTMLElement[];
+  setFieldDisabled: (disabled: boolean) => void;
+};
+
+export type ChainMethodHandler = (
+  ctx: ChainMethodContext<any, any>,
+  ...args: any[]
+) => any;
+
+export type ChainMethodDefinition<
+  THandler extends ChainMethodHandler = ChainMethodHandler,
+> = {
+  target: ChainMethodTarget;
+  pathPattern?: string[];
+  handler: THandler;
+};
+
+export type ChainMethodDefinitions = Record<string, ChainMethodDefinition<any>>;
+
+type ChainMethodCallable<THandler> = THandler extends (
+  ctx: any,
+  ...args: infer TArgs
+) => infer TReturn
+  ? (...args: TArgs) => TReturn
+  : never;
+
+export type ChainMethodCallables<TMethods> = {
+  [K in keyof TMethods]: TMethods[K] extends ChainMethodDefinition<infer TFn>
+    ? ChainMethodCallable<TFn>
+    : never;
+};
+
 export type KeyedTypes<TMap extends Record<string, any>> = {
   __key: 'keyed';
-  map: { [K in keyof TMap]: TMap[K] };
+  map: TMap;
 };
 
 export const keyedSchema = <TMap extends Record<string, any>>() =>
@@ -58,6 +115,9 @@ type ScopedMetadataMethods<TFieldMetaData> = {
   getFieldMetaData: () => TFieldMetaData | undefined;
   setFieldMetaData: (data: Partial<TFieldMetaData>) => void;
   removeFieldMetaData: () => void;
+  getFieldRefs: () => RefObject<any>[];
+  getFieldElements: () => HTMLElement[];
+  setFieldDisabled: (disabled: boolean) => void;
 };
 
 // These are the existing global methods that still require a path.
@@ -65,6 +125,11 @@ type GlobalMetadataMethods<TFieldMetaData> = {
   getFieldMetaData: (path: string[]) => TFieldMetaData | undefined;
   setFieldMetaData: (path: string[], data: Partial<TFieldMetaData>) => void;
   removeFieldMetaData: (path: string[]) => void;
+  getFieldRefs: (path: string[]) => RefObject<any>[];
+  getFieldElements: (path: string[]) => HTMLElement[];
+  setFieldDisabled: (path: string[], disabled: boolean) => void;
+  getAllFieldElements: () => HTMLElement[];
+  setAllFieldsDisabled: (disabled: boolean) => void;
 };
 // Simplified: All params use the same TFieldMetaData type
 export type UseHookParams<
@@ -193,9 +258,13 @@ export type CogsPlugin<
   THookReturn,
   TPluginMetaData,
   TFieldMetaData,
+  TChainMethods extends ChainMethodDefinitions = {},
+  TInitialState extends Record<string, unknown> = {},
 > = {
   name: TName;
 
+  initialState?: () => TInitialState;
+
   useHook?: (
     params: UseHookParams<TOptions, TPluginMetaData, TFieldMetaData, any>
   ) => THookReturn;
@@ -239,6 +308,8 @@ export type CogsPlugin<
       any
     >
   ) => React.ReactNode;
+
+  chainMethods?: TChainMethods;
 };
 
 // Metadata helpers
@@ -276,6 +347,45 @@ export function createMetadataContext<TPluginMetaData, TFieldMetaData>(
       getGlobalStore
         .getState()
         .removePluginMetaData(stateKey, path, pluginName),
+    getFieldRefs: (path: string[]): RefObject<any>[] => {
+      const meta = getGlobalStore.getState().getShadowMetadata(stateKey, path);
+      if (!meta?.clientActivityState?.elements) return [];
+      const refs: RefObject<any>[] = [];
+      meta.clientActivityState.elements.forEach((entry: any) => {
+        if (entry.domRef?.current) refs.push(entry.domRef);
+      });
+      return refs;
+    },
+
+    getFieldElements: (path: string[]): HTMLElement[] => {
+      const meta = getGlobalStore.getState().getShadowMetadata(stateKey, path);
+      if (!meta?.clientActivityState?.elements) return [];
+      const elements: HTMLElement[] = [];
+      meta.clientActivityState.elements.forEach((entry: any) => {
+        if (entry.domRef?.current) elements.push(entry.domRef.current);
+      });
+      return elements;
+    },
+
+    setFieldDisabled: (path: string[], disabled: boolean) => {
+      const meta = getGlobalStore.getState().getShadowMetadata(stateKey, path);
+      if (!meta?.clientActivityState?.elements) return;
+      meta.clientActivityState.elements.forEach((entry: any) => {
+        const el = entry.domRef?.current;
+        if (!el) return;
+        if ('disabled' in el) el.disabled = disabled;
+        else {
+          el.style.pointerEvents = disabled ? 'none' : '';
+          el.setAttribute('aria-disabled', String(disabled));
+        }
+      });
+    },
+    getAllFieldElements: (): HTMLElement[] => {
+      return getAllFieldElements(stateKey);
+    },
+
+    setAllFieldsDisabled: (disabled: boolean): void =>
+      setAllFieldsDisabled(stateKey, disabled),
   };
 }
 
@@ -290,19 +400,54 @@ export function createScopedMetadataContext<TPluginMetaData, TFieldMetaData>(
   );
 
   return {
-    // Return the global methods for plugin metadata
     ...globalContext,
-    // Override the field methods with new, path-scoped versions
     getFieldMetaData: (): TFieldMetaData | undefined =>
       globalContext.getFieldMetaData(path),
-
     setFieldMetaData: (data: Partial<TFieldMetaData>) =>
       globalContext.setFieldMetaData(path, data),
-
     removeFieldMetaData: () => globalContext.removeFieldMetaData(path),
+
+    // NEW: Direct access to the DOM refs for this field
+    getFieldRefs: (): RefObject<any>[] => {
+      const meta = getGlobalStore.getState().getShadowMetadata(stateKey, path);
+      if (!meta?.clientActivityState?.elements) return [];
+      const refs: RefObject<any>[] = [];
+      meta.clientActivityState.elements.forEach((entry: any) => {
+        if (entry.domRef?.current) {
+          refs.push(entry.domRef);
+        }
+      });
+      return refs;
+    },
+
+    getFieldElements: (): HTMLElement[] => {
+      const meta = getGlobalStore.getState().getShadowMetadata(stateKey, path);
+      if (!meta?.clientActivityState?.elements) return [];
+      const elements: HTMLElement[] = [];
+      meta.clientActivityState.elements.forEach((entry: any) => {
+        if (entry.domRef?.current) {
+          elements.push(entry.domRef.current);
+        }
+      });
+      return elements;
+    },
+
+    setFieldDisabled: (disabled: boolean) => {
+      const meta = getGlobalStore.getState().getShadowMetadata(stateKey, path);
+      if (!meta?.clientActivityState?.elements) return;
+      meta.clientActivityState.elements.forEach((entry: any) => {
+        const el = entry.domRef?.current;
+        if (!el) return;
+        if ('disabled' in el) {
+          el.disabled = disabled;
+        } else {
+          el.style.pointerEvents = disabled ? 'none' : '';
+          el.setAttribute('aria-disabled', String(disabled));
+        }
+      });
+    },
   };
 }
-
 // Add this type export - derives from what createScopedMetadataContext returns
 export type ScopedMetadataContext<TFieldMetaData = any> = {
   getFieldMetaData: () => TFieldMetaData | undefined;
@@ -323,200 +468,769 @@ type ZodObjOutput<T extends z.ZodObject<any>> = {
 type OutputOf<T extends z.ZodTypeAny> =
   T extends z.ZodObject<any> ? Prettify<ZodObjOutput<T>> : z.output<T>;
 
+type MethodFactory = <THandler extends ChainMethodHandler>(
+  handler: THandler
+) => ChainMethodDefinition<THandler>;
+
+type PathMethodFactory = MethodFactory & {
+  array: MethodFactory;
+  object: MethodFactory;
+  primitive: MethodFactory;
+  boolean: MethodFactory;
+  field: MethodFactory;
+};
+
+type MethodsBuilderParams = {
+  path: (selector: (state: any) => any) => PathMethodFactory;
+  array: MethodFactory;
+  object: MethodFactory;
+  primitive: MethodFactory;
+  boolean: MethodFactory;
+  field: MethodFactory;
+};
+
+const createMethodDefinition =
+  (target: ChainMethodTarget, pathPattern?: string[]): MethodFactory =>
+  (handler) => ({
+    target,
+    pathPattern,
+    handler,
+  });
+
+const createPathRecorder = () => {
+  const path: string[] = [];
+
+  const recorder = new Proxy(
+    {},
+    {
+      get(_target, prop) {
+        if (typeof prop !== 'string') return recorder;
+        if (prop === 'then') return undefined;
+        path.push(prop === '$' ? '*' : prop);
+        return recorder;
+      },
+    }
+  );
+
+  return { recorder, path };
+};
+
+const createMethodsBuilderParams = (): MethodsBuilderParams => {
+  const path = (selector: (state: any) => any): PathMethodFactory => {
+    const { recorder, path: pathPattern } = createPathRecorder();
+    selector(recorder);
+
+    const base = createMethodDefinition(
+      'any',
+      pathPattern
+    ) as PathMethodFactory;
+    base.array = createMethodDefinition('array', pathPattern);
+    base.object = createMethodDefinition('object', pathPattern);
+    base.primitive = createMethodDefinition('primitive', pathPattern);
+    base.boolean = createMethodDefinition('boolean', pathPattern);
+    base.field = createMethodDefinition('any', pathPattern);
+
+    return base;
+  };
+
+  return {
+    path,
+    array: createMethodDefinition('array'),
+    object: createMethodDefinition('object'),
+    primitive: createMethodDefinition('primitive'),
+    boolean: createMethodDefinition('boolean'),
+    field: createMethodDefinition('any'),
+  };
+};
+// --- EXTRACTED BUILDER TYPES (Solves TS7056 Serialization Error) ---
+
+export type PluginTransformFn<
+  TOptions,
+  THookReturn,
+  TPluginMetaData,
+  TFieldMetaData,
+> = (
+  params: TransformStateParams<
+    TOptions,
+    THookReturn,
+    TPluginMetaData,
+    TFieldMetaData,
+    any
+  >
+) => void;
+
+export type PluginUpdateFn<
+  TOptions,
+  THookReturn,
+  TPluginMetaData,
+  TFieldMetaData,
+> = (
+  params: OnUpdateParams<
+    TOptions,
+    THookReturn,
+    TPluginMetaData,
+    TFieldMetaData,
+    any
+  >
+) => void;
+
+export type PluginFormUpdateFn<
+  TOptions,
+  THookReturn,
+  TPluginMetaData,
+  TFieldMetaData,
+> = (
+  params: OnFormUpdateParams<
+    TOptions,
+    THookReturn,
+    TPluginMetaData,
+    TFieldMetaData,
+    any
+  >
+) => void;
+
+export type PluginFormWrapperFn<
+  TOptions,
+  THookReturn,
+  TPluginMetaData,
+  TFieldMetaData,
+> = (
+  params: FormWrapperParams<
+    TOptions,
+    THookReturn,
+    TPluginMetaData,
+    TFieldMetaData,
+    any
+  >
+) => React.ReactNode;
+
+export type CogsPluginBuilder<
+  TName extends string,
+  TOptions,
+  TPluginMetaData,
+  TFieldMetaData,
+  THookReturn,
+  TChainMethods extends ChainMethodDefinitions,
+  HasTransform extends boolean,
+  HasUpdate extends boolean,
+  HasFormUpdate extends boolean,
+  HasMethods extends boolean,
+  HasWrapper extends boolean,
+  HasInitialState extends boolean,
+  TInitialState extends Record<string, unknown> = {},
+> = Prettify<
+  CogsPlugin<
+    TName,
+    TOptions,
+    THookReturn,
+    TPluginMetaData,
+    TFieldMetaData,
+    TChainMethods,
+    TInitialState
+  >
+> &
+  (HasTransform extends true
+    ? {}
+    : {
+        transformState(
+          fn: PluginTransformFn<
+            TOptions,
+            THookReturn,
+            TPluginMetaData,
+            TFieldMetaData
+          >
+        ): CogsPluginBuilder<
+          TName,
+          TOptions,
+          TPluginMetaData,
+          TFieldMetaData,
+          THookReturn,
+          TChainMethods,
+          true,
+          HasUpdate,
+          HasFormUpdate,
+          HasMethods,
+          HasWrapper,
+          HasInitialState,
+          TInitialState
+        >;
+      }) &
+  (HasUpdate extends true
+    ? {}
+    : {
+        onUpdate(
+          fn: PluginUpdateFn<
+            TOptions,
+            THookReturn,
+            TPluginMetaData,
+            TFieldMetaData
+          >
+        ): CogsPluginBuilder<
+          TName,
+          TOptions,
+          TPluginMetaData,
+          TFieldMetaData,
+          THookReturn,
+          TChainMethods,
+          HasTransform,
+          true,
+          HasFormUpdate,
+          HasMethods,
+          HasWrapper,
+          HasInitialState,
+          TInitialState
+        >;
+      }) &
+  (HasFormUpdate extends true
+    ? {}
+    : {
+        onFormUpdate(
+          fn: PluginFormUpdateFn<
+            TOptions,
+            THookReturn,
+            TPluginMetaData,
+            TFieldMetaData
+          >
+        ): CogsPluginBuilder<
+          TName,
+          TOptions,
+          TPluginMetaData,
+          TFieldMetaData,
+          THookReturn,
+          TChainMethods,
+          HasTransform,
+          HasUpdate,
+          true,
+          HasMethods,
+          HasWrapper,
+          HasInitialState,
+          TInitialState
+        >;
+      }) &
+  (HasMethods extends true
+    ? {}
+    : {
+        methods<TNextMethods extends ChainMethodDefinitions>(
+          fn: (helpers: MethodsBuilderParams) => TNextMethods
+        ): CogsPluginBuilder<
+          TName,
+          TOptions,
+          TPluginMetaData,
+          TFieldMetaData,
+          THookReturn,
+          TNextMethods,
+          HasTransform,
+          HasUpdate,
+          HasFormUpdate,
+          true,
+          HasWrapper,
+          HasInitialState,
+          TInitialState
+        >;
+      }) &
+  (HasWrapper extends true
+    ? {}
+    : {
+        formWrapper(
+          fn: PluginFormWrapperFn<
+            TOptions,
+            THookReturn,
+            TPluginMetaData,
+            TFieldMetaData
+          >
+        ): CogsPluginBuilder<
+          TName,
+          TOptions,
+          TPluginMetaData,
+          TFieldMetaData,
+          THookReturn,
+          TChainMethods,
+          HasTransform,
+          HasUpdate,
+          HasFormUpdate,
+          HasMethods,
+          true,
+          HasInitialState,
+          TInitialState
+        >;
+      }) &
+  (HasInitialState extends true
+    ? {}
+    : {
+        initialState<TNewState extends Record<string, unknown>>(
+          fn: () => TNewState
+        ): CogsPluginBuilder<
+          TName,
+          TOptions,
+          TPluginMetaData,
+          TFieldMetaData,
+          THookReturn,
+          TChainMethods,
+          HasTransform,
+          HasUpdate,
+          HasFormUpdate,
+          HasMethods,
+          HasWrapper,
+          true,
+          TNewState
+        >;
+      });
+
+export type CreatePluginStart<
+  TName extends string,
+  TOptions,
+  TPluginMetaData,
+  TFieldMetaData,
+> = CogsPluginBuilder<
+  TName,
+  TOptions,
+  TPluginMetaData,
+  TFieldMetaData,
+  never,
+  {},
+  false,
+  false,
+  false,
+  false,
+  false,
+  false,
+  {}
+> & {
+  useHook<THookReturn>(
+    hookFn: (
+      params: UseHookParams<TOptions, TPluginMetaData, TFieldMetaData, any>
+    ) => THookReturn
+  ): CogsPluginBuilder<
+    TName,
+    TOptions,
+    TPluginMetaData,
+    TFieldMetaData,
+    THookReturn,
+    {},
+    false,
+    false,
+    false,
+    false,
+    false,
+    false,
+    {}
+  >;
+  initialState<TNewState extends Record<string, unknown>>(
+    fn: () => TNewState
+  ): CogsPluginBuilder<
+    TName,
+    TOptions,
+    TPluginMetaData,
+    TFieldMetaData,
+    never,
+    {},
+    false,
+    false,
+    false,
+    false,
+    false,
+    true,
+    TNewState
+  >;
+};
+
+// --- UPDATED EXPLICIT CONTEXT ---
+
 export function createPluginContext<
-  O extends z.ZodTypeAny,
+  O extends z.ZodTypeAny | undefined = undefined,
   PM extends z.ZodTypeAny | undefined = undefined,
   FM extends z.ZodTypeAny | undefined = undefined,
->(schemas: { options: O; pluginMetaData?: PM; fieldMetaData?: FM }) {
-  // Crucial: compute from the generic params, not from an object-indexed optional type
-  type Options = OutputOf<O>;
+>(schemas?: {
+  options?: O;
+  pluginMetaData?: PM;
+  fieldMetaData?: FM;
+}): {
+  // EXPLICIT RETURN ANNOTATION STOPS TS FROM SERIALIZING
+  createPlugin: <TName extends string>(
+    name: TName
+  ) => CreatePluginStart<
+    TName,
+    O extends z.ZodTypeAny ? OutputOf<O> : undefined,
+    PM extends z.ZodTypeAny ? OutputOf<PM> : unknown,
+    FM extends z.ZodTypeAny ? OutputOf<FM> : unknown
+  >;
+} {
+  type Options = O extends z.ZodTypeAny ? OutputOf<O> : undefined;
   type PluginMetaData = PM extends z.ZodTypeAny ? OutputOf<PM> : unknown;
   type FieldMetaData = FM extends z.ZodTypeAny ? OutputOf<FM> : unknown;
 
   function createPlugin<TName extends string>(name: TName) {
-    type TransformFn<THookReturn> = (
-      params: TransformStateParams<
+    type Builder<
+      THookReturn,
+      TChainMethods extends ChainMethodDefinitions,
+      HT extends boolean,
+      HU extends boolean,
+      HFU extends boolean,
+      HM extends boolean,
+      HW extends boolean,
+      HI extends boolean,
+      TInitialState extends Record<string, unknown> = {},
+    > = CogsPluginBuilder<
+      TName,
+      Options,
+      PluginMetaData,
+      FieldMetaData,
+      THookReturn,
+      TChainMethods,
+      HT,
+      HU,
+      HFU,
+      HM,
+      HW,
+      HI,
+      TInitialState
+    >;
+
+    const createPluginObject = <
+      THookReturn = never,
+      TChainMethods extends ChainMethodDefinitions = {},
+      TInitialState extends Record<string, unknown> = {},
+    >(
+      hookFn?: (
+        params: UseHookParams<Options, PluginMetaData, FieldMetaData, any>
+      ) => THookReturn,
+      transformFn?: PluginTransformFn<
         Options,
         THookReturn,
         PluginMetaData,
         FieldMetaData
-      >
-    ) => void;
-
-    type UpdateFn<THookReturn> = (
-      params: OnUpdateParams<
+      >,
+      updateHandler?: PluginUpdateFn<
         Options,
         THookReturn,
         PluginMetaData,
         FieldMetaData
-      >
-    ) => void;
-
-    type FormUpdateFn<THookReturn> = (
-      params: OnFormUpdateParams<
+      >,
+      formUpdateHandler?: PluginFormUpdateFn<
+        Options,
+        THookReturn,
+        PluginMetaData,
+        FieldMetaData
+      >,
+      chainMethods?: TChainMethods,
+      initialStateFn?: () => TInitialState,
+      wrapperFn?: PluginFormWrapperFn<
         Options,
         THookReturn,
         PluginMetaData,
         FieldMetaData
       >
-    ) => void;
-
-    type Plugin<THookReturn> = Prettify<
-      CogsPlugin<TName, Options, THookReturn, PluginMetaData, FieldMetaData>
-    >;
-
-    const createPluginObject = <THookReturn = never>(
-      hookFn?: (
-        params: UseHookParams<Options, PluginMetaData, FieldMetaData>
-      ) => THookReturn,
-      transformFn?: TransformFn<THookReturn>,
-      updateHandler?: UpdateFn<THookReturn>,
-      formUpdateHandler?: FormUpdateFn<THookReturn>
-    ): Plugin<THookReturn> => {
+    ): Prettify<
+      CogsPlugin<
+        TName,
+        Options,
+        THookReturn,
+        PluginMetaData,
+        FieldMetaData,
+        TChainMethods,
+        TInitialState
+      >
+    > => {
       return {
         name,
+        initialState: initialStateFn,
         useHook: hookFn as any,
         transformState: transformFn as any,
         onUpdate: updateHandler as any,
         onFormUpdate: formUpdateHandler as any,
+        formWrapper: wrapperFn as any,
+        chainMethods,
       };
     };
 
-    type BuildRet<
-      THookReturn,
-      HasTransform extends boolean,
-      HasUpdate extends boolean,
-      HasFormUpdate extends boolean,
-      HasWrapper extends boolean,
-    > = Plugin<THookReturn> &
-      (HasTransform extends true
-        ? {}
-        : {
-            transformState(
-              fn: TransformFn<THookReturn>
-            ): BuildRet<
-              THookReturn,
-              true,
-              HasUpdate,
-              HasFormUpdate,
-              HasWrapper
-            >;
-          }) &
-      (HasUpdate extends true
-        ? {}
-        : {
-            onUpdate(
-              fn: UpdateFn<THookReturn>
-            ): BuildRet<
-              THookReturn,
-              HasTransform,
-              true,
-              HasFormUpdate,
-              HasWrapper
-            >;
-          }) &
-      (HasFormUpdate extends true
-        ? {}
-        : {
-            onFormUpdate(
-              fn: FormUpdateFn<THookReturn>
-            ): BuildRet<THookReturn, HasTransform, HasUpdate, true, HasWrapper>;
-          });
-
     function createBuilder<
       THookReturn = never,
+      TChainMethods extends ChainMethodDefinitions = {},
       HasTransform extends boolean = false,
       HasUpdate extends boolean = false,
       HasFormUpdate extends boolean = false,
+      HasMethods extends boolean = false,
       HasWrapper extends boolean = false,
+      HasInitialState extends boolean = false,
+      TInitialState extends Record<string, unknown> = {},
     >(
       hookFn?: (
-        params: UseHookParams<Options, PluginMetaData, FieldMetaData>
+        params: UseHookParams<Options, PluginMetaData, FieldMetaData, any>
       ) => THookReturn,
-      transformFn?: TransformFn<THookReturn>,
-      updateHandler?: UpdateFn<THookReturn>,
-      formUpdateHandler?: FormUpdateFn<THookReturn>
-    ): BuildRet<
+      transformFn?: PluginTransformFn<
+        Options,
+        THookReturn,
+        PluginMetaData,
+        FieldMetaData
+      >,
+      updateHandler?: PluginUpdateFn<
+        Options,
+        THookReturn,
+        PluginMetaData,
+        FieldMetaData
+      >,
+      formUpdateHandler?: PluginFormUpdateFn<
+        Options,
+        THookReturn,
+        PluginMetaData,
+        FieldMetaData
+      >,
+      chainMethods?: TChainMethods,
+      initialStateFn?: () => TInitialState,
+      wrapperFn?: PluginFormWrapperFn<
+        Options,
+        THookReturn,
+        PluginMetaData,
+        FieldMetaData
+      >
+    ): Builder<
       THookReturn,
+      TChainMethods,
       HasTransform,
       HasUpdate,
       HasFormUpdate,
-      HasWrapper
+      HasMethods,
+      HasWrapper,
+      HasInitialState,
+      TInitialState
     > {
-      const plugin = createPluginObject<THookReturn>(
+      const plugin = createPluginObject<THookReturn, TChainMethods, TInitialState>(
         hookFn,
         transformFn,
         updateHandler,
-        formUpdateHandler
+        formUpdateHandler,
+        chainMethods,
+        initialStateFn,
+        wrapperFn
       );
 
       const methods = {} as Partial<
-        BuildRet<
+        Builder<
           THookReturn,
+          TChainMethods,
           HasTransform,
           HasUpdate,
           HasFormUpdate,
-          HasWrapper
+          HasMethods,
+          HasWrapper,
+          HasInitialState,
+          TInitialState
         >
       >;
 
       if (!transformFn) {
-        (methods as any).transformState = (fn: TransformFn<THookReturn>) =>
+        (methods as any).transformState = (
+          fn: PluginTransformFn<
+            Options,
+            THookReturn,
+            PluginMetaData,
+            FieldMetaData
+          >
+        ) =>
           createBuilder<
             THookReturn,
+            TChainMethods,
             true,
             HasUpdate,
             HasFormUpdate,
-            HasWrapper
-          >(hookFn, fn, updateHandler, formUpdateHandler);
+            HasMethods,
+            HasWrapper,
+            HasInitialState,
+            TInitialState
+          >(
+            hookFn,
+            fn,
+            updateHandler,
+            formUpdateHandler,
+            chainMethods,
+            initialStateFn,
+            wrapperFn
+          );
       }
       if (!updateHandler) {
-        (methods as any).onUpdate = (fn: UpdateFn<THookReturn>) =>
+        (methods as any).onUpdate = (
+          fn: PluginUpdateFn<
+            Options,
+            THookReturn,
+            PluginMetaData,
+            FieldMetaData
+          >
+        ) =>
           createBuilder<
             THookReturn,
+            TChainMethods,
             HasTransform,
             true,
             HasFormUpdate,
-            HasWrapper
-          >(hookFn, transformFn, fn, formUpdateHandler);
+            HasMethods,
+            HasWrapper,
+            HasInitialState,
+            TInitialState
+          >(
+            hookFn,
+            transformFn,
+            fn,
+            formUpdateHandler,
+            chainMethods,
+            initialStateFn,
+            wrapperFn
+          );
       }
       if (!formUpdateHandler) {
-        (methods as any).onFormUpdate = (fn: FormUpdateFn<THookReturn>) =>
-          createBuilder<THookReturn, HasTransform, HasUpdate, true, HasWrapper>(
+        (methods as any).onFormUpdate = (
+          fn: PluginFormUpdateFn<
+            Options,
+            THookReturn,
+            PluginMetaData,
+            FieldMetaData
+          >
+        ) =>
+          createBuilder<
+            THookReturn,
+            TChainMethods,
+            HasTransform,
+            HasUpdate,
+            true,
+            HasMethods,
+            HasWrapper,
+            HasInitialState,
+            TInitialState
+          >(
+            hookFn,
+            transformFn,
+            updateHandler,
+            fn,
+            chainMethods,
+            initialStateFn,
+            wrapperFn
+          );
+      }
+      if (!chainMethods) {
+        (methods as any).methods = <
+          TNextMethods extends ChainMethodDefinitions,
+        >(
+          fn: (helpers: MethodsBuilderParams) => TNextMethods
+        ) =>
+          createBuilder<
+            THookReturn,
+            TNextMethods,
+            HasTransform,
+            HasUpdate,
+            HasFormUpdate,
+            true,
+            HasWrapper,
+            HasInitialState,
+            TInitialState
+          >(
+            hookFn,
+            transformFn,
+            updateHandler,
+            formUpdateHandler,
+            fn(createMethodsBuilderParams()),
+            initialStateFn,
+            wrapperFn
+          );
+      }
+      if (!initialStateFn) {
+        (methods as any).initialState = <TNewState extends Record<string, unknown>>(fn: () => TNewState) =>
+          createBuilder<
+            THookReturn,
+            TChainMethods,
+            HasTransform,
+            HasUpdate,
+            HasFormUpdate,
+            HasMethods,
+            HasWrapper,
+            true,
+            TNewState
+          >(
+            hookFn,
+            transformFn,
+            updateHandler,
+            formUpdateHandler,
+            chainMethods,
+            fn,
+            wrapperFn
+          );
+      }
+      if (!wrapperFn) {
+        (methods as any).formWrapper = (
+          fn: PluginFormWrapperFn<
+            Options,
+            THookReturn,
+            PluginMetaData,
+            FieldMetaData
+          >
+        ) =>
+          createBuilder<
+            THookReturn,
+            TChainMethods,
+            HasTransform,
+            HasUpdate,
+            HasFormUpdate,
+            HasMethods,
+            true,
+            HasInitialState,
+            TInitialState
+          >(
             hookFn,
             transformFn,
             updateHandler,
+            formUpdateHandler,
+            chainMethods,
+            initialStateFn,
             fn
           );
       }
 
-      return Object.assign(plugin, methods) as BuildRet<
+      return Object.assign(plugin, methods) as Builder<
         THookReturn,
+        TChainMethods,
         HasTransform,
         HasUpdate,
         HasFormUpdate,
-        HasWrapper
+        HasMethods,
+        HasWrapper,
+        HasInitialState,
+        TInitialState
       >;
     }
 
     const start = Object.assign(
-      createBuilder<never, false, false, false, false>(),
+      createBuilder<never, {}, false, false, false, false, false, false, {}>(),
       {
         useHook<THookReturn>(
           hookFn: (
-            params: UseHookParams<Options, PluginMetaData, FieldMetaData>
+            params: UseHookParams<Options, PluginMetaData, FieldMetaData, any>
           ) => THookReturn
         ) {
-          return createBuilder<THookReturn, false, false, false, false>(hookFn);
+          return createBuilder<
+            THookReturn,
+            {},
+            false,
+            false,
+            false,
+            false,
+            false,
+            false,
+            {}
+          >(hookFn);
+        },
+        initialState<TNewState extends Record<string, unknown>>(fn: () => TNewState) {
+          return createBuilder<
+            never,
+            {},
+            false,
+            false,
+            false,
+            false,
+            false,
+            true,
+            TNewState
+          >(undefined, undefined, undefined, undefined, undefined, fn);
         },
       }
-    ) as BuildRet<never, false, false, false, false> & {
-      useHook<THookReturn>(
-        hookFn: (
-          params: UseHookParams<Options, PluginMetaData, FieldMetaData>
-        ) => THookReturn
-      ): BuildRet<THookReturn, false, false, false, false>;
-    };
+    ) as CreatePluginStart<TName, Options, PluginMetaData, FieldMetaData>;
 
     return start;
   }