import { Emitter } from "@noya-app/emitter";
import {
  Asset,
  EncodedSchema,
  Input,
  OutputTransform,
  Resource,
  decodeSchema,
  encodeSchema,
  findAllDefs,
} from "@noya-app/noya-schemas";
import { Cache, hash, isDeepEqual, uuid } from "@noya-app/noya-utils";
import { GetAtPath, Observable, PathKey, get, set } from "@noya-app/observable";
import { TaskSnapshot } from "@noya-app/task-runner";
import { TSchema } from "@sinclair/typebox";
import { Draft } from "mutative";
import { CreateParams } from "./IOManager";
import { LogManager, type ServerScriptLogEntry } from "./LogManager";
import { Secret } from "./SecretManager";
import { checkType } from "./checkType";
import { diff } from "./diff";
import { ExtendedPatch, ExtendedPathKey, applyExtendedPatch } from "./extended";
import { HistoryEntries } from "./historyEntries";
import { SandboxState } from "./SandboxManager";
import {
  PolicyAuthContext,
  PolicyViolation,
  applyServerComputedMetadata,
  enforceMultiplayerPolicies,
} from "./multiplayerPolicies";
import {
  DrainInputQueueOptions,
  InputQueueItem,
  ServerScriptManager,
  ServerScriptManagerOptions,
  serializeServerScriptLogValue,
} from "./serverScripts";
import {
  HistoryEntry,
  OperationManager,
  StateManager,
  StateManagerOptions,
  createEditHistoryEntry,
  createPatchHistoryEntry,
} from "./stateManager";
import { MultiplayerUser } from "./sync";
import { Task } from "./tasks";

export const createHash: typeof hash = (value, options) =>
  hash(value, { ...options, ignoreUndefinedProperties: true });

type SetStateAction<S> = S | ((prevState: S) => S);

/**
 * TODO:
 *
 * Pending commits may need to be updated when a patch is accepted or rejected?
 * Since right now we're using optimisticState based on pending commits - but
 * instead we could use a memoization approach
 *
 * Server may want to say whether a commit is from self or not, so that we can
 * more easily ignore patches that we sent (e.g. in case we change a pending patch
 * after sending? but it already has an id so maybe not needed)
 *
 * Prune old patches from the server
 *
 * MSM can just have a sm StateManager, mostly separate from any multiplayer logic,
 * that does its usual patch merging etc, but can be used to send undo/redo patches
 */

type ServerInvocation = {
  scriptId?: string;
  invocation: number;
  timestamp: number;
};

export type MultiplayerPatchMetadata = {
  id: string;
  name?: string;
  timestamp: number;
  serverInvocation?: ServerInvocation & {
    source: "server" | "client-sim";
  };
};

export type ClientCommit<S, M> = {
  state: S;
  baseHash: string;
  historyEntry: Required<HistoryEntry<S, M & MultiplayerPatchMetadata>>;
};

export type ClientConnectionStatus = "connected" | "disconnected";

export type ClientToServerMessage<S> =
  | {
      type: "init";
      object: S;
      force?: boolean;
      schema?: EncodedSchema;
      inputs?: CreateParams<Input>[];
      outputTransforms?: CreateParams<OutputTransform>[];
    }
  | {
      type: "patch";
      id: string;
      name?: string;
      patches: ExtendedPatch[];
      baseHash: string;
      serverInvocation?: ServerInvocation;
    }
  | {
      type: "resourceEdit.open";
      resourceId: string;
    }
  | {
      type: "resourceEdit.patch";
      resourceId: string;
      id: string;
      patches: ExtendedPatch[];
      baseHash: string;
    }
  | {
      type: "resourceEdit.close";
      resourceId: string;
    }
  | {
      type: "gitFileEdit.open";
      repoId: string;
      filePath: string;
      ref: string;
    }
  | {
      type: "gitFileEdit.patch";
      repoId: string;
      filePath: string;
      ref: string;
      id: string;
      patches: ExtendedPatch[];
      baseHash: string;
    }
  | {
      type: "gitFileEdit.close";
      repoId: string;
      filePath: string;
      ref: string;
    }
  | { type: "setSharedConnectionData"; connectionId: string; data?: unknown }
  | {
      type: "enqueueInput";
      connectionId?: string;
      queue: string;
      payload: unknown;
    }
  | { type: "ping"; id: string }
  | {
      type: "activityEvents.subscribe";
      streamId: string;
      cursor?: string;
    }
  | {
      type: "activityEvents.unsubscribe";
      streamId: string;
    };

export type RejectInitReason =
  | "invalidDataType"
  | "schemaMismatch"
  | "outdatedSchema"
  | "dataMigrationFailure";

export type RejectPatchReason =
  | "serverStateNotInitialized"
  | "errorApplyingPatch"
  | "clientOutOfSync"
  | "invalidDataType"
  | "outdatedSchema"
  | "policyViolation";

export type PipelineStatus = "idle" | "running" | "success" | "error";

export type ServerToClientMessage<S> =
  | { type: "object"; object: S; schema?: EncodedSchema }
  | { type: "rejectInit"; reason: RejectInitReason; error?: string }
  | {
      type: "acceptPatch";
      id: string;
      patches: ExtendedPatch[];
      hash?: string;
      serverInvocation?: ServerInvocation;
    }
  | {
      type: "rejectPatch";
      id: string;
      reason: RejectPatchReason;
      error?: string;
    }
  | {
      type: "resourceEdit.object";
      resourceId: string;
      content: string;
      hash?: string;
    }
  | {
      type: "resourceEdit.acceptPatch";
      resourceId: string;
      id: string;
      patches: ExtendedPatch[];
      hash?: string;
    }
  | {
      type: "resourceEdit.rejectPatch";
      resourceId: string;
      id: string;
      reason: RejectPatchReason;
      error?: string;
    }
  | {
      type: "gitFileEdit.object";
      repoId: string;
      filePath: string;
      ref: string;
      content: string;
      hash?: string;
    }
  | {
      type: "gitFileEdit.acceptPatch";
      repoId: string;
      filePath: string;
      ref: string;
      id: string;
      patches: ExtendedPatch[];
      hash?: string;
    }
  | {
      type: "gitFileEdit.rejectPatch";
      repoId: string;
      filePath: string;
      ref: string;
      id: string;
      reason: RejectPatchReason;
      error?: string;
    }
  | {
      type: "gitFileEdit.committed";
      repoId: string;
      ref: string;
      oid: string;
    }
  | { type: "connectedUsers"; users: MultiplayerUser[] }
  | { type: "schemaMigration"; schema: EncodedSchema }
  | { type: "sharedConnectionData"; data: Record<string, unknown> }
  | {
      type: "workflow";
      workflow: {
        status: PipelineStatus;
        tasks: TaskSnapshot<any>[];
        result: any;
      };
    }
  | { type: "tasks"; tasks: Task[] }
  | { type: "assets"; assets: Asset[] }
  | { type: "resources"; resources: Resource[] }
  | { type: "inputs"; inputs: Input[] }
  | { type: "outputTransforms"; outputTransforms: OutputTransform[] }
  | { type: "pong"; id: string }
  | { type: "secrets"; secrets: Secret[] }
  | { type: "fileName"; name: string; changeId?: string }
  | { type: "activityEvents"; activityEvents: unknown[] }
  | {
      type: "activityEvents.update";
      streamId: string;
      activityEvents: unknown[];
    }
  | { type: "serverScript.log"; log: ServerScriptLogEntry }
  | { type: "sandbox"; state: SandboxState };

export type MultiplayerStateManagerErrorType =
  | RejectInitReason
  | "schemaMigration"
  | RejectPatchReason;

export class MultiplayerStateManagerError extends Error {
  constructor(
    public readonly reason: MultiplayerStateManagerErrorType,
    message?: string
  ) {
    super(message);
  }
}

export type MultiplayerStateManagerMode = "standalone" | "synchronized";

export type MultiplayerStateManagerOptions<S, M> = StateManagerOptions<
  S,
  Partial<M> & MultiplayerPatchMetadata
> & {
  randomId?: () => string;
  getNow?: () => number;
  serverScripts?: {
    simulateLocally?: boolean;
    getQuickJSModule: ServerScriptManagerOptions["getQuickJSModule"];
    randomUUID?: ServerScriptManagerOptions["randomUUID"];
    getNow?: () => number;
  };
  getSharedConnectionData?: () => Record<string, unknown>;
  drainInputQueue?: (options?: DrainInputQueueOptions) => InputQueueItem[];
  overrideExistingState?: boolean;
  outputTransforms?: CreateParams<OutputTransform>[];
  inputs?: CreateParams<Input>[];
  debug?: boolean;
  mode?: MultiplayerStateManagerMode;
};

export class MultiplayerStateManager<
  TState,
  TMetadata extends object = object,
> extends Emitter {
  get _metadataType(): Partial<TMetadata> & MultiplayerPatchMetadata {
    throw new Error("Not implemented - this is just a type hint");
  }

  state: TState;
  schema?: TSchema;

  errorEmitter = new Emitter<[MultiplayerStateManagerError]>();

  sm: StateManager<TState, typeof this._metadataType>;

  pendingCommits: ClientCommit<TState, typeof this._metadataType>[] = [];
  connectionStatus$ = new Observable<ClientConnectionStatus>("disconnected");
  isInitialized$ = new Observable<boolean>(false);

  outgoingMessages: ClientToServerMessage<TState>[] = [];
  incomingMessages: ServerToClientMessage<TState>[] = [];

  optimisticState$: Observable<TState>;

  partialHashCache = new Cache<unknown, string>({ capacity: 20000 });
  createHashCached = new Cache<unknown, string>({ capacity: 100 }).wrap(
    (value) => createHash(value, { cache: this.partialHashCache })
  );

  applyCommitCache = new Cache<string, { input: TState; output: TState }>({
    capacity: 25,
  });

  private policyAuthContext: PolicyAuthContext | undefined;
  private enforcePoliciesLocally = true;

  enhanceMetadata(
    metadata: Partial<TMetadata> & Partial<MultiplayerPatchMetadata>
  ): typeof this._metadataType {
    return {
      id: this.getRandomId(),
      timestamp: Date.now(),
      ...metadata,
    };
  }

  getRandomId() {
    return this.options.randomId?.() ?? uuid();
  }

  private initializeServerSimulator() {
    const serverScripts = this.options.serverScripts;
    if (
      !serverScripts?.getQuickJSModule ||
      serverScripts.simulateLocally !== true
    )
      return;

    this.serverScriptSimulator = new ServerScriptManager({
      getQuickJSModule: serverScripts.getQuickJSModule,
      randomUUID: serverScripts.randomUUID,
      environment: {
        target: "client-simulator",
        mode: "simulation",
        label: "simulator",
        clock: serverScripts.getNow ?? (() => this.getServerNowEstimate()),
      },
      getState: () => this.state,
      getSharedConnectionData:
        this.options.getSharedConnectionData ?? (() => ({})),
      setSharedConnectionData: () => {
        // We don't currently mirror shared connection data from predictions
      },
      getInitialInvocation: (scriptId: string) => {
        const key = scriptId ?? "default";
        const tracked = this.serverInvocationTracker.get(key)?.invocation;
        return tracked;
      },
      drainInputQueue: (options?: DrainInputQueueOptions) =>
        this.options.drainInputQueue?.(options) ?? [],
      applyState: async ({ scriptId, nextState, invocation, timestamp }) => {
        this.applyServerPrediction({
          scriptId,
          nextState,
          invocation,
          timestamp,
        });
      },
      log: (entry) => {
        console[entry.level]("[server-sim]", entry.scriptId, ...entry.values);

        const entryValues = entry.values.map((value) =>
          serializeServerScriptLogValue(value)
        );

        this.serverSimulationLogManager.append({
          id: this.getRandomId(),
          scriptId: entry.scriptId,
          level: entry.level,
          values: entryValues,
          timestamp: Date.now(),
          origin: entry.origin,
          target: entry.target,
          mode: entry.mode,
        });
      },
    });

    this.serverScriptSimulator.setIsHot(false);
    void this.serverScriptSimulator.updateSchema(this.schema);
  }

  setPolicyAuthContext(context: PolicyAuthContext | undefined) {
    this.policyAuthContext = context;
  }

  setLocalPolicyEnforcement(enabled: boolean) {
    this.enforcePoliciesLocally = enabled;
  }

  constructor(
    state: TState | (() => TState),
    public options: MultiplayerStateManagerOptions<TState, TMetadata> = {}
  ) {
    super();

    this.state = state instanceof Function ? state() : state;
    this.schema = options.schema;

    this.sm = new StateManager<TState, typeof this._metadataType>(
      this.state,
      options
    );

    this.optimisticState$ = Observable.from(
      {
        get: this.getOptimisticState,
        subscribe: this.addListener,
      },
      {
        isEqual: (a, b) =>
          this.createHashCached(a) === this.createHashCached(b),
      }
    );

    if (this.options.mode === "standalone") {
      this.autoConnect();
    }

    this.initializeServerSimulator();
  }

  autoConnect = () => {
    const handleEmitMessage = (message: ClientToServerMessage<TState>) => {
      switch (message.type) {
        case "init": {
          this.incomingMessages.push({
            type: "object",
            object: message.object,
            schema: message.schema,
          });
          this.processIncomingMessages();
          break;
        }
        case "patch": {
          this.incomingMessages.push({
            type: "acceptPatch",
            id: message.id,
            patches: message.patches,
            serverInvocation: message.serverInvocation,
          });
          this.processIncomingMessages();
          break;
        }
      }
    };

    this.messageEmitter.addListener(handleEmitMessage);
    this.connect();
  };

  getHash() {
    return this.createHashCached(this.state);
  }

  _setState(state: TState, schema: TSchema | undefined) {
    this.state = state;
    this.schema = schema;
    this.emitIfChanged();
    this.sm._setState(state);
  }

  _lastEmittedHash: string | undefined;

  emitIfChanged = () => {
    if (this._lastEmittedHash === this.getOptimisticHash()) return;
    this._lastEmittedHash = this.getOptimisticHash();
    this.emit();
  };

  get isConnected() {
    return this.connectionStatus$.get() === "connected";
  }

  _pushCommit(
    state: TState,
    historyEntry: Required<HistoryEntry<TState, typeof this._metadataType>>,
    {
      addToHistory,
      shouldTypeCheck = true,
      sendMessage = true,
    }: {
      addToHistory: boolean;
      shouldTypeCheck?: boolean;
      sendMessage?: boolean;
    }
  ) {
    // Can't edit if not connected
    if (!this.isConnected) {
      console.warn("Can't edit when not connected");

      return;
    }

    if (shouldTypeCheck) {
      checkType(state, this.schema, findAllDefs(this.schema));
    }

    const optimisticState = this.getOptimisticState();

    const shouldEnforcePolicies =
      this.enforcePoliciesLocally && this.policyAuthContext !== undefined;

    const policyViolation = shouldEnforcePolicies
      ? this.getLocalPolicyViolation({
          previousState: optimisticState,
          nextState: state,
          patches: historyEntry.redoPatches,
          auth: this.policyAuthContext!,
        })
      : undefined;

    if (policyViolation) {
      this.errorEmitter.emit(
        new MultiplayerStateManagerError(
          "policyViolation",
          policyViolation.message ?? "Policy violation"
        )
      );

      return;
    }

    const commit: ClientCommit<TState, typeof this._metadataType> = {
      state,
      historyEntry,
      baseHash: this.getOptimisticHash(),
    };

    this.pendingCommits.push(commit);
    this.emitIfChanged();

    if (addToHistory) {
      this.applyCommitCache.set(historyEntry.metadata.id, {
        input: optimisticState,
        output: state,
      });

      this.sm._setStateWithHistoryEntry(state, historyEntry, {
        shouldTypeCheck: false,
      });
    }

    if (sendMessage) {
      this.sendMessage({
        type: "patch",
        id: commit.historyEntry.metadata.id,
        name: commit.historyEntry.metadata.name,
        patches: commit.historyEntry.redoPatches,
        baseHash: commit.baseHash,
        ...(historyEntry.metadata.serverInvocation && {
          serverInvocation: historyEntry.metadata.serverInvocation,
        }),
      });
    }
  }

  private getLocalPolicyViolation(options: {
    previousState: TState;
    nextState: TState;
    patches: ExtendedPatch[];
    auth: PolicyAuthContext;
  }): PolicyViolation | undefined {
    if (options.patches.length === 0) return undefined;

    const auth = options.auth;

    if (!this.schema) {
      if ((auth.access ?? "write") !== "write") {
        return {
          action: "update",
          patch: options.patches[0],
          message: "Insufficient access level to edit this field",
        };
      }

      return undefined;
    }

    const now = this.options.getNow?.() ?? Date.now();

    let evaluatedState = options.nextState;

    try {
      const clonedPatches = options.patches.map(cloneExtendedPatch);
      const serverPatches = applyServerComputedMetadata({
        schema: this.schema,
        patches: clonedPatches,
        auth,
        now,
      });
      evaluatedState = applyExtendedPatch(options.previousState, [
        ...clonedPatches,
        ...serverPatches,
      ]);
    } catch {
      evaluatedState = options.nextState;
    }

    return enforceMultiplayerPolicies({
      schema: this.schema,
      patches: options.patches,
      previousState: options.previousState,
      nextState: evaluatedState,
      auth,
      now,
    });
  }

  canUndo = () => this.sm.canUndo();

  canRedo = () => this.sm.canRedo();

  undo = () => {
    if (!this.sm.canUndo()) return;

    const handler = this.sm.history[this.sm.historyIndex - 1];

    this.sm.historyIndex--;

    const result = createPatchHistoryEntry(
      this.getOptimisticState(),
      {},
      {
        ...handler.metadata,
        // we need a unique id for the commit, distinct from the original id
        id: this.getRandomId(),
      },
      handler.undoPatches!,
      handler.redoPatches!
    );

    if (!result) return;

    this._pushCommit(result.state, result.historyEntry, {
      addToHistory: false,
    });

    return this.sm.history[this.sm.historyIndex];
  };

  redo = () => {
    if (!this.sm.canRedo()) return;

    const handler = this.sm.history[this.sm.historyIndex];

    this.sm.historyIndex++;

    const result = createPatchHistoryEntry(
      this.getOptimisticState(),
      {},
      {
        ...handler.metadata,
        // we need a unique id for the commit, distinct from the original id
        id: this.getRandomId(),
      },
      handler.redoPatches!,
      handler.undoPatches!
    );

    if (!result) return;

    this._pushCommit(result.state, result.historyEntry, {
      addToHistory: false,
    });

    return this.sm.history[this.sm.historyIndex - 1];
  };

  /**
   * Mutate the state in-place.
   *
   * This will create a patch based on the changes made to the state and create
   * a history entry with the patch.
   */
  editDraft = (
    ...args:
      | [metadata: Partial<TMetadata>, callback: (draft: Draft<TState>) => void]
      | [callback: (draft: Draft<TState>) => void]
  ) => {
    const [metadata, callback] = args.length === 1 ? [{}, args[0]] : args;

    const optimisticState = this.getOptimisticState();

    const result = createEditHistoryEntry(
      optimisticState,
      {},
      this.enhanceMetadata(metadata),
      callback
    );

    if (!result) return;

    this._pushCommit(result.state, result.historyEntry, { addToHistory: true });
  };

  /**
   * Apply a patch optimistically.
   *
   * Creates a history entry with the patch. If inversePatches are not provided,
   * this will perform a diff between the current state and the new state.
   */
  applyPatch = (
    ...args:
      | [
          metadata: Partial<TMetadata>,
          { patches: ExtendedPatch[]; inversePatches?: ExtendedPatch[] },
        ]
      | [{ patches: ExtendedPatch[]; inversePatches?: ExtendedPatch[] }]
  ) => {
    let [metadata, { patches, inversePatches }] =
      args.length === 1 ? [{}, args[0]] : args;

    const optimisticState = this.getOptimisticState();

    const newState = applyExtendedPatch(optimisticState, patches);

    checkType(newState, this.schema, findAllDefs(this.schema));

    if (!inversePatches) {
      const [, inv] = diff(optimisticState, newState);

      inversePatches = inv;
    }

    const result = createPatchHistoryEntry(
      newState,
      {},
      this.enhanceMetadata(metadata),
      patches,
      inversePatches
    );

    if (!result) return;

    this._pushCommit(result.state, result.historyEntry, { addToHistory: true });
  };

  /**
   * Set the current state optimistically.
   *
   * This will perform a diff between the current state and the new state, and
   * create a history entry with the patch.
   */
  setState = (
    ...args:
      | [metadata: Partial<TMetadata>, action: SetStateAction<TState>]
      | [action: SetStateAction<TState>]
  ) => {
    const [metadata, action] = args.length === 1 ? [{}, args[0]] : args;

    const optimisticState = this.getOptimisticState();

    const newState =
      action instanceof Function ? action(optimisticState) : action;

    checkType(newState, this.schema, findAllDefs(this.schema));

    const [patches, inversePatches] = diff(optimisticState, newState, {
      useExtendedPatches: true,
    });

    const result = createPatchHistoryEntry(
      newState,
      {},
      this.enhanceMetadata(metadata),
      patches,
      inversePatches
    );

    if (!result) return;

    this._pushCommit(result.state, result.historyEntry, { addToHistory: true });
  };

  setStateAtPath = <P extends PathKey[] | string>(
    ...args:
      | [
          metadata: Partial<TMetadata>,
          path: P,
          value: SetStateAction<GetAtPath<TState, P>>,
        ]
      | [path: P, value: SetStateAction<GetAtPath<TState, P>>]
  ) => {
    const [metadata, path, value] =
      args.length === 2 ? [{}, args[0], args[1]] : args;

    const setState = this.setState as (
      metadata: Partial<TMetadata>,
      action: SetStateAction<GetAtPath<TState, P>>
    ) => void;

    if (!path || path.length === 0) {
      setState(metadata, value);
    } else {
      const optimisticState = this.getOptimisticState();

      const newValueAtPath =
        value instanceof Function
          ? value(get(optimisticState, path) as GetAtPath<TState, P>)
          : value;

      const newState = set(optimisticState, path, newValueAtPath);

      this.setState(metadata, newState);
    }
  };

  operate = (
    ...args:
      | [
          metadata: Partial<TMetadata>,
          operator: (operationManager: OperationManager<TState>) => void,
        ]
      | [operator: (operationManager: OperationManager<TState>) => void]
  ) => {
    const [metadata, operator] = args.length === 1 ? [{}, args[0]] : args;

    const operationManager = new OperationManager(this.state);

    operator(operationManager);

    const patches = HistoryEntries.compactPatches(operationManager.patches);
    const inversePatches = HistoryEntries.compactPatches(
      operationManager.inversePatches
    );
    const afterState = applyExtendedPatch(this.state, patches);

    if (this.options.allowEmptyEdits !== true) {
      // Ignore edits that don't change the state
      if (patches.length === 0) return;
    }

    const result = createPatchHistoryEntry(
      afterState,
      {},
      this.enhanceMetadata(metadata),
      patches,
      inversePatches
    );

    if (!result) return;

    this._pushCommit(result.state, result.historyEntry, { addToHistory: true });
  };

  getConfirmedState = () => {
    return this.state;
  };

  getOptimisticState = () => {
    return this.pendingCommits.at(-1)?.state ?? this.state;
  };

  getOptimisticHash() {
    return this.createHashCached(this.getOptimisticState());
  }

  reinitialize(state: TState) {
    this.pendingCommits = [];
    this.serverPredictionCommits.clear();
    this.emitIfChanged();

    this.sendMessage({
      type: "init",
      object: state,
      ...(this.schema && { schema: encodeSchema(this.schema) }),
    });
  }

  _handledMessages: Record<string, boolean> = {};

  processIncomingMessages() {
    while (this.incomingMessages.length > 0) {
      const message = this.incomingMessages.shift()!;

      if ("id" in message) {
        if (this._handledMessages[message.id]) {
          continue;
        }

        this._handledMessages[message.id] = true;
      }

      switch (message.type) {
        case "object": {
          this._setState(
            message.object,
            message.schema ? decodeSchema(message.schema) : undefined
          );

          void this.serverScriptSimulator?.updateSchema(this.schema);
          this.isInitialized$.set(true);
          break;
        }
        case "acceptPatch": {
          const patches = message.patches;
          this.handleServerInvocation(message.serverInvocation);

          if (!message.serverInvocation) {
            this.clearServerPredictions();
          }

          const oldState = this.state;

          let newState: TState;

          const existing = this.applyCommitCache.get(message.id);
          const matchingCommit = this.pendingCommits.find(
            (entry) => entry.historyEntry.metadata.id === message.id
          );
          const canReuseCachedState =
            existing &&
            existing.input === oldState &&
            patchesEqual(matchingCommit?.historyEntry.redoPatches, patches);

          try {
            if (canReuseCachedState) {
              newState = existing.output;
            } else {
              newState = applyExtendedPatch(oldState, patches);
            }
          } catch (e) {
            this.reinitialize(oldState);
            break;
          }

          // Remove from pending entries
          this.pendingCommits = this.pendingCommits.filter(
            (entry) => entry.historyEntry.metadata.id !== message.id
          );

          // Verify that the hash is the same as the one we received
          if (
            message.hash !== undefined &&
            message.hash !== this.createHashCached(newState)
          ) {
            this.reinitialize(oldState);
            break;
          }

          this._setState(newState, this.schema);
          this.serverScriptSimulator?.handleStateChange(patches);
          break;
        }
        case "rejectPatch": {
          switch (message.reason) {
            default: {
              this.reinitialize(this.state);

              this.errorEmitter.emit(
                new MultiplayerStateManagerError(message.reason, message.error)
              );

              break;
            }
          }

          break;
        }
        case "rejectInit": {
          console.error(
            "State initialized with invalid data type.",
            message.error
          );

          this.errorEmitter.emit(
            new MultiplayerStateManagerError(message.reason, message.error)
          );

          break;
        }
        case "schemaMigration": {
          this.errorEmitter.emit(
            new MultiplayerStateManagerError(
              "schemaMigration",
              "Schema migration requires a page reload."
            )
          );

          break;
        }
      }
    }
  }

  connect() {
    this.connectionStatus$.set("connected");
    this.serverScriptSimulator?.setIsHot(true);
    this.sendInit();
  }

  sendInit({
    force = this.options.overrideExistingState,
    state = this.state,
    schema = this.schema,
    outputTransforms = this.options.outputTransforms,
    inputs = this.options.inputs,
  }: {
    force?: boolean;
    state?: TState;
    schema?: TSchema;
    outputTransforms?: CreateParams<OutputTransform>[];
    inputs?: CreateParams<Input>[];
  } = {}) {
    this.sendMessage({
      type: "init",
      object: state,
      ...(schema && { schema: encodeSchema(schema) }),
      ...(outputTransforms && { outputTransforms }),
      ...(inputs && { inputs }),
      ...(force && { force }),
    });
  }

  sendMessage(message: ClientToServerMessage<TState>) {
    this.outgoingMessages.push(message);
    this.messageEmitter.emit(message);
  }

  disconnect() {
    this.connectionStatus$.set("disconnected");
    this.serverScriptSimulator?.setIsHot(false);
  }

  hasPendingMessages() {
    return this.incomingMessages.length > 0 || this.outgoingMessages.length > 0;
  }

  messageEmitter = new Emitter<[ClientToServerMessage<TState>]>();

  /** SERVER PREDICTIONS */

  serverSimulationLogManager = new LogManager();
  serverSimulationLogs$ = this.serverSimulationLogManager.logs$;

  private serverPredictionCommits = new Map<string, string>();
  private serverInvocationTracker = new Map<
    string,
    { invocation: number; timestamp: number }
  >();
  private serverClockSkewMs = 0;
  private serverScriptSimulator?: ServerScriptManager;

  private getNow() {
    return this.options.getNow?.() ?? Date.now();
  }

  private getServerNowEstimate() {
    return this.getNow() + this.serverClockSkewMs;
  }

  getServerSimulationDiagnostics() {
    const invocations = Array.from(this.serverInvocationTracker.entries()).map(
      ([scriptId, info]) => ({
        scriptId: scriptId === "default" ? undefined : scriptId,
        invocation: info.invocation,
        timestamp: info.timestamp,
      })
    );

    const pendingPredictions = Array.from(
      this.serverPredictionCommits.entries()
    ).map(([key, commitId]) => {
      const [scriptId, invocationString] = key.split(":");
      const invocation = Number(invocationString);
      return {
        scriptId: scriptId === "default" ? undefined : scriptId,
        invocation: Number.isNaN(invocation) ? undefined : invocation,
        commitId,
      };
    });

    return {
      simulateLocally: this.options.serverScripts?.simulateLocally === true,
      simulatorActive: !!this.serverScriptSimulator,
      clockSkewMs: this.serverClockSkewMs,
      invocations,
      pendingPredictions,
    };
  }

  getLatestServerInvocation(scriptId?: string) {
    const key = scriptId ?? "default";
    return this.serverInvocationTracker.get(key);
  }

  getLatestPrediction() {
    const latestPrediction = [...this.pendingCommits]
      .reverse()
      .find(
        (commit) =>
          commit.historyEntry.metadata.serverInvocation?.source === "client-sim"
      );

    if (!latestPrediction) return undefined;

    return {
      state: latestPrediction.state,
      serverInvocation: latestPrediction.historyEntry.metadata.serverInvocation,
    };
  }

  private clearServerPredictions() {
    if (this.serverPredictionCommits.size === 0) return;
    const toRemove = new Set(this.serverPredictionCommits.values());
    this.pendingCommits = this.pendingCommits.filter(
      (entry) => !toRemove.has(entry.historyEntry.metadata.id)
    );
    this.serverPredictionCommits.clear();
    this.emitIfChanged();
  }

  applyServerPrediction(options: {
    scriptId?: string;
    nextState: TState;
    invocation?: number;
    timestamp?: number;
  }) {
    if (!this.isConnected) return;

    const key = options.scriptId ?? "default";
    const trackedInvocation = this.serverInvocationTracker.get(key)?.invocation;

    // If we already have this invocation (or newer) from the real server, skip simulation.
    if (
      trackedInvocation !== undefined &&
      options.invocation !== undefined &&
      options.invocation <= trackedInvocation
    ) {
      return;
    }

    const baseState = this.state;
    const [patches, inversePatches] = diff(baseState, options.nextState, {
      useExtendedPatches: true,
    });

    if (patches.length === 0) return;

    const nextInvocation =
      trackedInvocation !== undefined
        ? trackedInvocation + 1
        : options.invocation ?? 1;
    const timestamp = options.timestamp ?? this.getServerNowEstimate();

    // Drop any existing prediction for this invocation so we don't chain stale ones
    const previousId = this.serverPredictionCommits.get(
      `${key}:${nextInvocation}`
    );
    if (previousId) {
      this.pendingCommits = this.pendingCommits.filter(
        (entry) => entry.historyEntry.metadata.id !== previousId
      );
      this.serverPredictionCommits.delete(`${key}:${nextInvocation}`);
    }

    const predictionMetadata = {
      name: `predicted ${key}#${nextInvocation}`,
      serverInvocation: {
        scriptId: options.scriptId,
        invocation: nextInvocation,
        timestamp,
        source: "client-sim",
      },
    } as unknown as Partial<TMetadata> & Partial<MultiplayerPatchMetadata>;

    const result = createPatchHistoryEntry(
      options.nextState,
      {},
      this.enhanceMetadata(predictionMetadata),
      patches,
      inversePatches
    );

    if (!result) return;

    this.serverInvocationTracker.set(key, {
      invocation: nextInvocation,
      timestamp,
    });

    this._pushCommit(result.state, result.historyEntry, {
      addToHistory: false,
      sendMessage: false,
    });

    this.serverPredictionCommits.set(
      `${key}:${nextInvocation}`,
      result.historyEntry.metadata.id
    );

    this.emitIfChanged();
  }

  private handleServerInvocation(info?: {
    scriptId?: string;
    invocation: number;
    timestamp: number;
  }) {
    if (!info) return;

    const key = info.scriptId ?? "default";
    const observedSkew = info.timestamp - this.getNow();
    this.serverClockSkewMs =
      this.serverClockSkewMs === 0
        ? observedSkew
        : this.serverClockSkewMs * 0.9 + observedSkew * 0.1;

    const predictionKey = `${key}:${info.invocation}`;
    const predictedId = this.serverPredictionCommits.get(predictionKey);

    let removed = false;

    if (predictedId) {
      this.pendingCommits = this.pendingCommits.filter(
        (entry) => entry.historyEntry.metadata.id !== predictedId
      );
      this.serverPredictionCommits.delete(predictionKey);
      removed = true;
    }

    for (const [key, id] of Array.from(
      this.serverPredictionCommits.entries()
    )) {
      if (!key.startsWith(`${info.scriptId ?? "default"}:`)) continue;
      const [, invocationString] = key.split(":");
      const predictedInvocation = Number(invocationString);
      if (Number.isNaN(predictedInvocation)) continue;
      if (predictedInvocation <= info.invocation) {
        this.pendingCommits = this.pendingCommits.filter(
          (entry) => entry.historyEntry.metadata.id !== id
        );
        this.serverPredictionCommits.delete(key);
        removed = true;
      }
    }

    if (removed) {
      this.emitIfChanged();
    }

    this.serverInvocationTracker.set(key, {
      invocation: info.invocation,
      timestamp: info.timestamp,
    });

    // Keep the local simulator aligned with confirmed server invocations
    this.serverScriptSimulator?.syncInvocation(
      info.scriptId,
      info.invocation,
      info.timestamp
    );
  }
}

function patchesEqual(
  a: ExtendedPatch[] | undefined,
  b: ExtendedPatch[] | undefined
) {
  if (!a || !b) return false;
  return isDeepEqual(a, b);
}

function cloneExtendedPatch(patch: ExtendedPatch): ExtendedPatch {
  return {
    ...patch,
    path: patch.path.map(cloneExtendedPathSegment),
    ...(patch.from && { from: patch.from.map(cloneExtendedPathSegment) }),
    ...("value" in patch ? { value: clonePatchValue(patch.value) } : {}),
  };
}

function cloneExtendedPathSegment(segment: ExtendedPathKey): ExtendedPathKey {
  if (typeof segment === "object") {
    return { id: segment.id };
  }

  return segment;
}

function clonePatchValue<T>(value: T): T {
  if (Array.isArray(value)) {
    return value.map((item) => clonePatchValue(item)) as unknown as T;
  }

  if (value && typeof value === "object") {
    const result: Record<string, unknown> = {};

    for (const [key, entry] of Object.entries(value)) {
      result[key] = clonePatchValue(entry);
    }

    return result as T;
  }

  return value;
}

export type SharedConnectionMetadata = {
  updatedAt: number;
};

export type SelectionRange = {
  anchor: PathKey[];
  head?: PathKey[];
};

type Point = {
  x: number;
  y: number;
};

export type DefaultSharedConnectionData = {
  pointer?: Point;
  selection?: SelectionRange[];
};

export class SharedConnectionDataManager<E extends object> extends Emitter<
  [
    Record<string, (E & DefaultSharedConnectionData) | undefined>,
    Record<string, { updatedAt: number } | undefined>,
  ]
> {
  data$ = new Observable<
    Record<string, (E & DefaultSharedConnectionData) | undefined>
  >({});

  metadata$ = new Observable<
    Record<string, SharedConnectionMetadata | undefined>
  >({});

  currentConnectionId$ = new Observable<string | undefined>(undefined);

  currentConnectionDataEmitter = new Emitter<
    [(E & DefaultSharedConnectionData) | undefined]
  >();

  setCurrentConnectionId = (currentConnectionId: string | undefined) => {
    this.currentConnectionId$.set(currentConnectionId);
  };

  setForCurrentConnection = (value: E & DefaultSharedConnectionData) => {
    const currentConnectionId = this.currentConnectionId$.get();

    if (!currentConnectionId) {
      // console.log("Shared Connection Data: no current connection");
      return;
    }

    this.set(currentConnectionId, value);

    this.currentConnectionDataEmitter.emit(value);
  };

  set = (connectionId: string, value: E & DefaultSharedConnectionData) => {
    this.metadata$.set({
      ...this.metadata$.get(),
      [connectionId]: { updatedAt: Date.now() },
    });

    this.data$.set({ ...this.data$.get(), [connectionId]: value });

    this.emit(this.data$.get(), this.metadata$.get());
  };

  merge = (partial: Record<string, E & DefaultSharedConnectionData>) => {
    this.data$.set({ ...this.data$.get(), ...partial });

    for (const connectionId in partial) {
      this.metadata$.set({
        ...this.metadata$.get(),
        [connectionId]: { updatedAt: Date.now() },
      });
    }

    this.emit(this.data$.get(), this.metadata$.get());
  };

  delete = (connectionId: string) => {
    const { [connectionId]: _d, ...restData } = this.data$.get();
    this.data$.set(restData);

    const { [connectionId]: _m, ...restMetadata } = this.metadata$.get();
    this.metadata$.set(restMetadata);

    this.emit(this.data$.get(), this.metadata$.get());

    if (connectionId === this.currentConnectionId$.get()) {
      this.currentConnectionDataEmitter.emit(undefined);
    }
  };

  _snapshot:
    | {
        data: Record<string, E | undefined>;
        metadata: Record<string, SharedConnectionMetadata | undefined>;
        currentConnectionId: string | undefined;
      }
    | undefined;

  getSnapshot = () => {
    if (
      !this._snapshot ||
      this._snapshot.data !== this.data$.get() ||
      this._snapshot.metadata !== this.metadata$.get() ||
      this._snapshot.currentConnectionId !== this.currentConnectionId$.get()
    ) {
      this._snapshot = {
        data: this.data$.get(),
        metadata: this.metadata$.get(),
        currentConnectionId: this.currentConnectionId$.get(),
      };
    }

    return this._snapshot;
  };
}