/**
 * AgentRewrite
 * ============
 *
 * One canonical Tiptap command to perform an agent-driven replacement
 * of a range in the document - optionally with a typewriter animation
 * so the user sees the agent "typing" the new content in place of the
 * old one.
 *
 * Called by:
 *   - the real agent's tool calls (applyDiff / rewrite / rephrase),
 *   - the Playwright showcase demo script via the `__demo-chat` event.
 *
 * Why this lives in an extension (and not in App.tsx or the demo)
 * ---------------------------------------------------------------
 * The previous design had an ad-hoc `setTimeout` loop in App.tsx that
 * re-inserted characters at `startPos + i` - a fixed numeric origin.
 * Under Yjs collaboration any concurrent remote edit that shifted
 * `startPos` (e.g. another user inserting a character before it) would
 * desync the animation and land subsequent chunks inside the wrong
 * paragraph, producing a visible "write-over-itself" artifact.
 *
 * Here we track the insertion `cursor` in PM plugin state and remap it
 * through `tr.mapping` on every transaction - local OR remote. This
 * guarantees the typewriter resumes at the correct absolute position
 * no matter what Bob or Carol are doing elsewhere in the doc.
 *
 * Cancellation
 * ------------
 * Each call to `agentRewriteRange` bumps a cancel token stored in
 * plugin state. Any still-scheduled animation step that observes a
 * different token aborts. That way, if the agent fires two
 * overlapping rewrites, the first one stops immediately instead of
 * fighting the second for control of the doc.
 *
 * Undo-friendliness
 * -----------------
 * Every animation step is a small atomic `insertText` transaction,
 * so y-undo sees a tidy sequence of ReplaceSteps instead of one huge
 * debug-hostile mega-op. Undo works as expected (one step per chunk).
 */
import { Extension } from "@tiptap/core";
import { Plugin, PluginKey } from "@tiptap/pm/state";

export type AgentRewriteAnimation = "none" | "typewriter";

export interface AgentRewriteArgs {
  /** Absolute PM position of the start of the range to replace. */
  from: number;
  /** Absolute PM position of the end of the range to replace. */
  to: number;
  /** New text to insert in place of the deleted range. */
  text: string;
  /**
   * "none"  -> replace the range atomically in one transaction.
   * "typewriter" -> stream the replacement chunk-by-chunk with small
   *                 delays so the UI shows the agent "typing".
   * Default: "none" (safer for programmatic / non-demo callers).
   */
  animation?: AgentRewriteAnimation;
  /**
   * Characters per animation tick. A random int in [min,max] is chosen
   * each step, which makes the typing feel organic. Default [1, 2].
   */
  chunkSize?: [number, number];
  /**
   * Milliseconds between animation ticks. Default [22, 48].
   */
  chunkMs?: [number, number];
}

interface AgentRewriteState {
  /**
   * Current insertion position for the in-flight typewriter animation.
   * Mapped through every transaction so local + remote edits never
   * shift the cursor out of sync with the actual doc content.
   */
  cursor: number | null;
  /**
   * Monotonically-increasing token. Any call to `agentRewriteRange`
   * bumps this, which retroactively cancels any animation step still
   * scheduled from a previous call.
   */
  cancelToken: number;
}

interface AgentRewriteMeta {
  cancel?: boolean;
  setCursor?: number | null;
}

const agentRewriteKey = new PluginKey<AgentRewriteState>("agent-rewrite");

declare module "@tiptap/core" {
  // eslint-disable-next-line @typescript-eslint/no-unused-vars
  interface Commands<ReturnType> {
    agentRewrite: {
      /**
       * Replace `[from, to]` with `text`. When `animation === "typewriter"`
       * the insertion is streamed chunk-by-chunk; otherwise it's one
       * atomic transaction.
       */
      agentRewriteRange: (args: AgentRewriteArgs) => ReturnType;
      /** Abort any in-flight typewriter animation immediately. */
      cancelAgentRewrite: () => ReturnType;
    };
  }
}

export const AgentRewrite = Extension.create({
  name: "agentRewrite",

  addProseMirrorPlugins() {
    return [
      new Plugin<AgentRewriteState>({
        key: agentRewriteKey,
        state: {
          init: (): AgentRewriteState => ({ cursor: null, cancelToken: 0 }),
          apply(tr, prev): AgentRewriteState {
            let next = prev;
            // Remap the tracked cursor through any doc change (local or
            // remote). This is the whole point of using a plugin state
            // instead of a closure-captured number: Yjs pushes tr's for
            // remote edits and we ride their mapping.
            if (tr.docChanged && next.cursor !== null) {
              next = { ...next, cursor: tr.mapping.map(next.cursor) };
            }
            const meta = tr.getMeta(agentRewriteKey) as
              | AgentRewriteMeta
              | undefined;
            if (meta) {
              if (meta.cancel) {
                // `cancel` wipes the cursor AND advances the token so
                // in-flight steps observe a mismatch and bail out.
                next = {
                  cursor: null,
                  cancelToken: next.cancelToken + 1,
                };
              }
              if (meta.setCursor !== undefined) {
                // `setCursor` is the primitive used by the command to
                // re-arm the plugin with a fresh starting position
                // (AFTER any cancel has taken effect in the same meta).
                next = { ...next, cursor: meta.setCursor };
              }
            }
            return next;
          },
        },
      }),
    ];
  },

  addCommands() {
    return {
      cancelAgentRewrite:
        () =>
        ({ tr, dispatch }) => {
          if (dispatch) {
            tr.setMeta(agentRewriteKey, { cancel: true });
          }
          return true;
        },

      agentRewriteRange:
        (args) =>
        ({ state, view, tr, dispatch }) => {
          const docSize = state.doc.content.size;
          const a = Math.max(0, Math.min(args.from, docSize));
          const b = Math.max(0, Math.min(args.to, docSize));
          if (b < a) return false;
          if (!dispatch) return true;

          // Capture `myToken` BEFORE the dispatch. Tiptap's
          // CommandManager turns out to batch the `tr` we mutate here
          // into its own dispatch cycle, so reading `view.state`
          // immediately after calling the provided `dispatch(tr)` can
          // return the PRE-commit state (cancelToken still = N,
          // cursor still = null). The only reliable source is the
          // state we started from (`state`): we know this command
          // will bump cancelToken by 1 and seed the cursor at `a`,
          // so we compute those values ourselves instead of reading
          // them back from a potentially stale view.
          const prevToken = agentRewriteKey.getState(state)?.cancelToken ?? 0;
          const myToken = prevToken + 1;

          const text = args.text ?? "";
          const animation = args.animation ?? "none";

          // ----- Atomic path ("none") -----------------------------------
          // Do the whole replace in ONE transaction: cancel any
          // in-flight rewrite, replace [a, b] with `text`, clear the
          // cursor marker. Two separate dispatches produced double
          // content (insert + zombie old text) because Tiptap's
          // CommandManager does NOT update `view.state` between
          // `dispatch(tr)` and the next `view.state.tr.insertText`
          // call inside the same command - so the second tr
          // insertText ran against the PRE-delete state, effectively
          // adding the new phrase in front of the old one without
          // deleting it.
          if (animation === "none") {
            tr.setMeta(agentRewriteKey, {
              cancel: true,
              setCursor: null,
            } as AgentRewriteMeta);
            if (text) {
              tr.insertText(text, a, b);
            } else if (b > a) {
              tr.delete(a, b);
            }
            dispatch(tr);
            return true;
          }

          // ----- Typewriter path -----------------------------------
          // Step 1: single atomic transaction that (a) cancels any
          // in-flight rewrite, (b) deletes the target range, and (c)
          // plants the insertion cursor at `a`. Done in one tr so the
          // deletion and cursor seed land together - avoids a visible
          // gap between delete and first insert.
          tr.setMeta(agentRewriteKey, {
            cancel: true,
            setCursor: a,
          } as AgentRewriteMeta);
          if (b > a) tr.delete(a, b);
          dispatch(tr);

          if (!text) {
            // Nothing to insert, just a deletion. Done.
            view.dispatch(
              view.state.tr.setMeta(agentRewriteKey, {
                setCursor: null,
              } as AgentRewriteMeta),
            );
            return true;
          }

          // Typewriter: async stream of small insertText trs. Each step
          // reads the CURRENT cursor from plugin state (mapped through
          // every intervening tr, including remote Yjs updates), so the
          // animation is immune to concurrent edits.
          const [chunkMin, chunkMax] = args.chunkSize ?? [1, 2];
          const [delayMin, delayMax] = args.chunkMs ?? [22, 48];
          let i = 0;

          // Closure-held cursor is our source of truth: seeded at `a`,
          // then synced with plugin state before each insert so remote
          // Yjs trs between our ticks remap the cursor correctly via
          // the plugin's mapping logic. Using a closure (rather than
          // reading plugin state blindly) guarantees the first step
          // doesn't bail out if Tiptap's command dispatch hasn't yet
          // committed the `setCursor: a` meta to the plugin state that
          // `view.state` exposes.
          let cursor = a;

          const step = () => {
            if (view.isDestroyed) return;
            const ps = agentRewriteKey.getState(view.state);
            // Cancellation only: token mismatch means another rewrite
            // started (or `cancelAgentRewrite` fired).
            if (ps && ps.cancelToken !== myToken) {
              // Another rewrite started or `cancelAgentRewrite` fired.
              return;
            }
            // Pull the mapped cursor from plugin state if available -
            // this accounts for any remote Yjs inserts that happened
            // since the last step. Fall back to the closure value on
            // the very first tick, where plugin state can still be
            // null due to Tiptap's command dispatch batching.
            if (ps && ps.cursor !== null) cursor = ps.cursor;

            if (i >= text.length) {
              // Done: clear the cursor marker so nothing lingers in
              // plugin state. The caller is responsible for any visual
              // cleanup (e.g. `clearAgentFocus`).
              view.dispatch(
                view.state.tr.setMeta(agentRewriteKey, {
                  setCursor: null,
                } as AgentRewriteMeta),
              );
              return;
            }

            const remaining = text.length - i;
            const maxTake = Math.max(1, Math.min(chunkMax, remaining));
            const minTake = Math.max(1, Math.min(chunkMin, maxTake));
            const take =
              minTake +
              Math.floor(Math.random() * (maxTake - minTake + 1));
            const chunk = text.slice(i, i + take);

            // Insert at our closure-tracked cursor. The plugin's
            // apply() will advance the stored cursor via mapping so
            // subsequent remote trs remap it correctly.
            view.dispatch(view.state.tr.insertText(chunk, cursor));
            cursor += chunk.length;
            i += chunk.length;

            const delay =
              delayMin + Math.random() * Math.max(0, delayMax - delayMin);
            setTimeout(step, delay);
          };
          // Kick off on a macrotask so the initial delete tr has fully
          // applied (and the cursor seed is live in plugin state)
          // before the first step reads it.
          setTimeout(step, 0);
          return true;
        },
    };
  },
});