import { Extension } from "@tiptap/core";
import { yCursorPlugin, defaultCursorBuilder, defaultSelectionBuilder } from "@tiptap/y-tiptap";

export interface CursorUser {
  name: string;
  color: string;
  avatarUrl?: string;
}

export interface CollaborationCursorOptions {
  provider: any;
  user: CursorUser;
  render?: (user: CursorUser) => HTMLElement;
  selectionRender?: typeof defaultSelectionBuilder;
}

const awarenessStatesToArray = (states: Map<number, any>) => {
  return Array.from(states.entries()).map(([key, value]) => ({
    clientId: key,
    ...value.user,
  }));
};

export const CollaborationCursorV3 = Extension.create<CollaborationCursorOptions>({
  name: "collaborationCursor",
  priority: 999,

  addOptions() {
    return {
      provider: null,
      user: { name: "Anonymous", color: "#aaaaaa" } as CursorUser,
      render: (user: CursorUser) => {
        const cursor = document.createElement("span");
        cursor.classList.add("collaboration-cursor__caret");
        cursor.setAttribute("style", `border-color: ${user.color}`);
        const label = document.createElement("div");
        label.classList.add("collaboration-cursor__label");
        label.setAttribute("style", `background-color: ${user.color}`);

        if (user.avatarUrl) {
          const avatar = document.createElement("img");
          avatar.src = user.avatarUrl;
          avatar.classList.add("collaboration-cursor__avatar");
          label.appendChild(avatar);
        }

        label.appendChild(document.createTextNode(user.name));
        cursor.appendChild(label);
        return cursor;
      },
      selectionRender: defaultSelectionBuilder,
    };
  },

  addStorage() {
    return { users: [] as any[] };
  },

  onCreate() {
    if (!this.options.provider) {
      throw new Error('The "provider" option is required for CollaborationCursorV3');
    }
  },

  addCommands() {
    // In Tiptap v3, custom commands must be declared via module augmentation of
    // `Commands<ReturnType>` to satisfy `Partial<RawCommands>`. We don't need
    // that here: the command is used internally via `editor.commands.updateUser`
    // with a runtime cast, so we widen the return type to keep typing simple.
    return {
      updateUser:
        (attributes: Record<string, any>) =>
        () => {
          this.options.user = attributes as any;
          this.options.provider.awareness.setLocalStateField("user", this.options.user);
          return true;
        },
    } as any;
  },

  addProseMirrorPlugins() {
    const awareness = this.options.provider.awareness;
    awareness.setLocalStateField("user", this.options.user);
    this.storage.users = awarenessStatesToArray(awareness.states);
    awareness.on("update", () => {
      this.storage.users = awarenessStatesToArray(awareness.states);
    });

    // Wrap the cursor builder so we can silence a peer's caret+label when
    // that same peer already broadcasts an `agentFocus` in awareness. The
    // AgentFocus extension renders its own "<Name> agent" label on the
    // range the AI is working on, and without this guard we end up with
    // two overlapping badges for the same user (e.g. "Bob Mercier" caret
    // label + "Bob Mercier agent" focus label). y-tiptap re-invokes the
    // builder on every awareness change, so this check stays live.
    const userBuilder = this.options.render ?? defaultCursorBuilder;
    const cursorBuilder = (user: CursorUser, clientId: number): HTMLElement => {
      const el = userBuilder(user);
      const peerState = awareness.getStates().get(clientId);
      if (peerState?.agentFocus) {
        el.classList.add("collaboration-cursor--silenced");
      }
      return el;
    };

    return [
      yCursorPlugin(
        awareness,
        {
          cursorBuilder,
          selectionBuilder: this.options.selectionRender ?? defaultSelectionBuilder,
        },
      ),
    ];
  },
});