packages/agent-core-v2/src/state/state.ts

import { enableMapSet, type Draft } from 'immer';
import type { z } from 'zod';

import { BugIndicatingError } from '#/_base/errors/errors';
import type { StateKey } from '#/_base/state/stateRegistry';
import { Event2, registerEvent2Class, type Event2Class } from '#/app/event/event2';
import type { PartsTransformer, RecordDehydrator } from '#/wire/record';

import { StateError, StateErrors } from './errors';

enableMapSet();

export type { StateKey } from '#/_base/state/stateRegistry';
export type { PartsTransformer } from '#/wire/record';

export interface StateBlobCodec<S> {
  dehydrate: RecordDehydrator;
  rehydrate(state: S, transform: PartsTransformer): S | Promise<S>;
}

export interface FoldContext {
  readonly silent: boolean;
  checkpoint(): void;
  clearCheckpoints(): void;
  undoToCheckpoint(count: number): void;
  emit(event: Event2): void;
}

export type StateFold<S, E extends Event2<any> = Event2<any>> = (
  state: Draft<S>,
  event: E,
  ctx: FoldContext,
) => S | void;

export interface ReplayableOptions<S> {
  readonly schema: z.ZodType<S>;
  readonly durable?: boolean;
  readonly blobs?: StateBlobCodec<S>;
}

export interface UndoableOptions<S> {
  readonly onUndo?: (state: Draft<S>, count: number) => S | void;
}

export interface ReplayableStateMeta<S> {
  readonly schema: z.ZodType<S>;
  readonly durable: boolean;
  readonly blobs?: StateBlobCodec<S>;
  readonly undoable?: UndoableOptions<S>;
  readonly folds: ReadonlyMap<Event2Class<any, any>, StateFold<S, any>>;
}

export interface ReplayableStateKey<S> extends StateKey<DeepReadonly<S>> {
  readonly replayable: ReplayableStateMeta<S>;
  undoable(opts?: UndoableOptions<S>): ReplayableStateKey<S>;
  on<P, E extends Event2<P>>(cls: Event2Class<P, E>, fold: StateFold<S, E>): ReplayableStateKey<S>;
}

export interface StateKeyBuilder<T> extends StateKey<T> {
  replayable(opts: ReplayableOptions<T>): ReplayableStateKey<T>;
}

class ReplayableStateKeyImpl<S> implements ReplayableStateKey<S> {
  readonly snapshotExcluded = true;
  readonly initial: () => DeepReadonly<S>;

  private readonly meta: {
    readonly schema: z.ZodType<S>;
    readonly durable: boolean;
    readonly blobs?: StateBlobCodec<S>;
    undoable?: UndoableOptions<S>;
    readonly folds: Map<Event2Class<any, any>, StateFold<S, any>>;
  };

  constructor(
    readonly name: string,
    initial: () => S,
    opts: ReplayableOptions<S>,
  ) {
    this.initial = () => Object.freeze(initial()) as DeepReadonly<S>;
    this.meta = {
      schema: opts.schema,
      durable: opts.durable ?? true,
      blobs: opts.blobs,
      folds: new Map(),
    };
  }

  get replayable(): ReplayableStateMeta<S> {
    return this.meta;
  }

  undoable(opts?: UndoableOptions<S>): ReplayableStateKey<S> {
    if (this.meta.undoable !== undefined) {
      throw new BugIndicatingError(`State key '${this.name}' is already undoable`);
    }
    if (!this.meta.durable) {
      throw new BugIndicatingError(`Transient state key '${this.name}' cannot be undoable`);
    }
    this.meta.undoable = opts ?? {};
    return this;
  }

  on<P, E extends Event2<P>>(cls: Event2Class<P, E>, fold: StateFold<S, E>): ReplayableStateKey<S> {
    if (this.meta.folds.has(cls)) {
      throw new StateError(
        StateErrors.codes.STATE_DUPLICATE_FOLD,
        `State '${this.name}' already folds event '${cls.type}'`,
        { details: { state: this.name, type: cls.type } },
      );
    }
    if (!this.meta.durable && cls.durable) {
      throw new StateError(
        StateErrors.codes.STATE_DURABILITY_MISMATCH,
        `Transient state '${this.name}' cannot fold durable event '${cls.type}'`,
        { details: { state: this.name, type: cls.type } },
      );
    }
    registerEvent2Class(cls);
    this.meta.folds.set(cls, fold as StateFold<S, any>);
    return this;
  }
}

class StateKeyBuilderImpl<T> implements StateKeyBuilder<T> {
  constructor(
    readonly name: string,
    readonly initial: () => T,
  ) {}

  replayable(opts: ReplayableOptions<T>): ReplayableStateKey<T> {
    return new ReplayableStateKeyImpl(this.name, this.initial, opts);
  }
}

export function defineState<T>(name: string, initial: () => T): StateKeyBuilder<T> {
  return new StateKeyBuilderImpl(name, initial);
}

export interface UndoableProtocol {
  readonly events: {
    readonly appendMessage: Event2Class<any, any>;
    readonly applyCompaction: Event2Class<any, any>;
    readonly clear: Event2Class<any, any>;
    readonly undo: Event2Class<any, any>;
  };
  readonly isUndoAnchor: (message: unknown) => boolean;
  readonly isValidUndoCount: (count: number) => boolean;
}

let undoableProtocol: UndoableProtocol | undefined;

export function registerUndoableProtocol(protocol: UndoableProtocol): void {
  if (undoableProtocol !== undefined) {
    throw new BugIndicatingError('The undoable protocol is already registered');
  }
  undoableProtocol = protocol;
  for (const cls of Object.values(protocol.events)) {
    registerEvent2Class(cls);
  }
}

export function expandedStateFolds(
  key: ReplayableStateKey<any>,
): ReadonlyMap<Event2Class<any, any>, StateFold<any, any>> {
  const meta = key.replayable;
  if (meta.undoable === undefined) return meta.folds;
  if (undoableProtocol === undefined) {
    throw new BugIndicatingError(
      `State key '${key.name}' is undoable but no undoable protocol is registered ` +
        '(the contextMemory domain registers it at import time)',
    );
  }
  const protocol = undoableProtocol;
  if (meta.folds.has(protocol.events.undo)) {
    throw new BugIndicatingError(
      `Undoable state key '${key.name}' must not fold the undo event itself; ` +
        'use .undoable({ onUndo }) to customize the rollback',
    );
  }
  const custom = meta.undoable.onUndo !== undefined;
  const folds = new Map<Event2Class<any, any>, StateFold<any, any>>(meta.folds);
  const domainAppend = folds.get(protocol.events.appendMessage);
  folds.set(protocol.events.appendMessage, (state, event, ctx) => {
    if (!custom && protocol.isUndoAnchor(event.message)) {
      ctx.checkpoint();
      return;
    }
    return domainAppend?.(state, event, ctx);
  });
  for (const cls of [protocol.events.applyCompaction, protocol.events.clear]) {
    const domain = folds.get(cls);
    folds.set(cls, (state, event, ctx) => {
      ctx.clearCheckpoints();
      return domain?.(state, event, ctx);
    });
  }
  folds.set(protocol.events.undo, (state, event, ctx) => {
    if (!protocol.isValidUndoCount(event.count)) return;
    if (meta.undoable?.onUndo !== undefined) {
      return meta.undoable.onUndo(state, event.count);
    }
    ctx.undoToCheckpoint(event.count);
  });
  return folds;
}

export type EventApplier = (event: any, ctx: FoldContext) => void;

export function expandedModelAppliers(
  owner: string,
  undoable: boolean,
  appliers: ReadonlyMap<Event2Class<any, any>, EventApplier>,
  onUndo: ((count: number) => void) | undefined,
): ReadonlyMap<Event2Class<any, any>, EventApplier> {
  if (!undoable) return appliers;
  if (undoableProtocol === undefined) {
    throw new BugIndicatingError(
      `Agent model '${owner}' is undoable but no undoable protocol is registered ` +
        '(the contextMemory domain registers it at import time)',
    );
  }
  const protocol = undoableProtocol;
  if (appliers.has(protocol.events.undo)) {
    throw new BugIndicatingError(
      `Undoable agent model '${owner}' must not apply the undo event itself; ` +
        'override onUndo on the model to customize the rollback',
    );
  }
  const custom = onUndo !== undefined;
  const expanded = new Map<Event2Class<any, any>, EventApplier>(appliers);
  const domainAppend = expanded.get(protocol.events.appendMessage);
  expanded.set(protocol.events.appendMessage, (event, ctx) => {
    if (!custom && protocol.isUndoAnchor(event.message)) {
      ctx.checkpoint();
      return;
    }
    domainAppend?.(event, ctx);
  });
  for (const cls of [protocol.events.applyCompaction, protocol.events.clear]) {
    const domain = expanded.get(cls);
    expanded.set(cls, (event, ctx) => {
      ctx.clearCheckpoints();
      domain?.(event, ctx);
    });
  }
  expanded.set(protocol.events.undo, (event, ctx) => {
    if (!protocol.isValidUndoCount(event.count)) return;
    if (onUndo !== undefined) {
      onUndo(event.count);
      return;
    }
    ctx.undoToCheckpoint(event.count);
  });
  return expanded;
}

export function expandedRuntimeFolds(
  owner: string,
  undoable: boolean,
  folds: ReadonlyMap<Event2Class<any, any>, StateFold<any, any>>,
): ReadonlyMap<Event2Class<any, any>, StateFold<any, any>> {
  if (!undoable) return folds;
  if (undoableProtocol === undefined) {
    throw new BugIndicatingError(
      `Agent runtime '${owner}' is undoable but no undoable protocol is registered ` +
        '(the contextMemory domain registers it at import time)',
    );
  }
  const protocol = undoableProtocol;
  if (folds.has(protocol.events.undo)) {
    throw new BugIndicatingError(
      `Undoable agent runtime '${owner}' must not fold the undo event itself`,
    );
  }
  const expanded = new Map(folds);
  const domainAppend = expanded.get(protocol.events.appendMessage);
  expanded.set(protocol.events.appendMessage, (state, event, ctx) => {
    if (protocol.isUndoAnchor(event.message)) {
      ctx.checkpoint();
      return;
    }
    return domainAppend?.(state, event, ctx);
  });
  for (const cls of [protocol.events.applyCompaction, protocol.events.clear]) {
    const domain = expanded.get(cls);
    expanded.set(cls, (state, event, ctx) => {
      ctx.clearCheckpoints();
      return domain?.(state, event, ctx);
    });
  }
  expanded.set(protocol.events.undo, (_state, event, ctx) => {
    if (protocol.isValidUndoCount(event.count)) ctx.undoToCheckpoint(event.count);
  });
  return expanded;
}

export type DeepReadonly<T> = T extends (...args: infer A) => infer R
  ? (...args: A) => R
  : T extends ReadonlyMap<infer K, infer V>
    ? ReadonlyMap<DeepReadonly<K>, DeepReadonly<V>>
    : T extends ReadonlySet<infer V>
      ? ReadonlySet<DeepReadonly<V>>
      : T extends readonly (infer E)[]
        ? ReadonlyArray<DeepReadonly<E>>
        : T extends object
          ? { readonly [K in keyof T]: DeepReadonly<T[K]> }
          : T;