// Usage:
const aiExtension = createAIExtension(opts: AIExtensionOptions);

// Definitions:
function createAIExtension(options: AIExtensionOptions): (editor: BlockNoteEditor) => AIExtension;

type AIExtensionOptions = AIRequestHelpers & {
  /**
   * The name and color of the agent cursor when the AI is writing
   * @default { name: "AI", color: "#8bc6ff" }
   */
  agentCursor?: { name: string; color: string };
};

type AIRequestHelpers = {
  /**
   * Transport used by the AI SDK to send requests to your backend/LLM.
   * Implement to customize backend URLs or use a different transport (e.g. websockets).
   */
  transport?: ChatTransport<UIMessage>;

  /**
   * Use the ChatProvider to customize how the AI SDK Chat instance is created.
   * For example, when you want to reuse an existing Chat instance used in the rest of your application.
   *
   * @note you cannot use both `chatProvider` and `transport` together.
   */
  chatProvider?: () => Chat<UIMessage>;

  /**
   * Customize which stream tools are available to the LLM.
   */
  streamToolsProvider?: StreamToolsProvider<any, any>;
  // Provide `streamToolsProvider` in createAIExtension(options) or override per call via InvokeAIOptions.
  // If omitted, defaults to using `aiDocumentFormats.html.getStreamToolsProvider()`.

  /**
   * Extra options (headers/body/metadata) forwarded to the AI SDK request.
   */
  chatRequestOptions?: ChatRequestOptions;

  /**
   * Build the serializable document state that will be forwarded to the backend.
   *
   * @default aiDocumentFormats.html.defaultDocumentStateBuilder
   */
  documentStateBuilder?: DocumentStateBuilder<any>;
};

function getAIExtension(editor: BlockNoteEditor): AIExtension;

class AIExtension {
  /**
   * Execute a call to an LLM and apply the result to the editor
   */
  invokeAI(opts: InvokeAIOptions): Promise<void>;

  /**
   * Returns a read-only zustand store with the state of the AI Menu
   */
  get store(): ReadonlyStoreApi<{
    aiMenuState:
      | ({
          /**
           * The ID of the block that the AI menu is opened at.
           * This changes as the AI is making changes to the document
           */
          blockId: string;
        } & (
          | {
              status: "error";
              error: any;
            }
          | {
              status:
                | "user-input"
                | "thinking"
                | "ai-writing"
                | "user-reviewing";
            }
        ))
      | "closed";
  }>;

  /**
   * Returns a zustand store with the global configuration of the AI Extension.
   * These options are used by default across all LLM calls when calling {@link invokeAI}
   */
  readonly options: StoreApi<AIRequestHelpers>;

  /** Open the AI menu at a specific block */
  openAIMenuAtBlock(blockID: string): void;
  /** Close the AI menu */
  closeAIMenu(): void;
  /** Accept the changes made by the LLM */
  acceptChanges(): void;
  /** Reject the changes made by the LLM */
  rejectChanges(): void;
  /** Retry the previous LLM call (only valid when status is "error") */
  retry(): Promise<void>;
  /** Advanced: manually update the status shown by the AI menu */
  setAIResponseStatus(
    status:
      | "user-input"
      | "thinking"
      | "ai-writing"
      | "user-reviewing"
      | { status: "error"; error: any },
  ): void;
}

type InvokeAIOptions = {
  /** The user prompt */
  userPrompt: string;

  /** Whether to use the editor selection for the LLM call (default: true) */
  useSelection?: boolean;

  /**
   * If the user's cursor is in an empty paragraph, automatically delete it when the AI starts writing.
   * Used when typing `/ai` in an empty block. (default: true)
   */
  deleteEmptyCursorBlock?: boolean;
} & AIRequestHelpers; // Optionally override helpers per request

/** Return a provider for the stream tools available to the LLM */
type getStreamToolsProvider = (
  // Whether to add artificial delays between document edits
  // or apply them immediately as they're streamed in from the LLM without delays
  // (default: true)
  withDelays: boolean,
  // The stream tools to use, there are separate tools for adding, updating and deleting blocks
  // (default: { add: true, update: true, delete: true })
  defaultStreamTools?: {
    add?: boolean;
    update?: boolean;
    delete?: boolean;
  },
) => StreamToolsProvider;

type DocumentStateBuilder<T> = (
  aiRequest: Omit<AIRequest, "documentState">,
) => Promise<
  | {
      selection: false;
      blocks: BlocksWithCursor<T>[];
      isEmptyDocument: boolean;
    }
  | {
      selection: true;
      selectedBlocks: { id: string; block: T }[];
      blocks: { block: T }[];
      isEmptyDocument: boolean;
    }
>;

type AIRequest = {
  editor: BlockNoteEditor;
  selectedBlocks?: Block[];
  emptyCursorBlockToDelete?: string;
  streamTools: StreamTool<any>[];
  documentState: DocumentState<any>;
  onStart: () => void;
};

async function sendMessageWithAIRequest(
  chat: Chat<UIMessage>,
  aiRequest: AIRequest,
  message?: Parameters<Chat<UIMessage>["sendMessage"]>[0],
  options?: Parameters<Chat<UIMessage>["sendMessage"]>[1],
): Promise<Result<void>>;

async function buildAIRequest(opts: {
  editor: BlockNoteEditor;
  useSelection?: boolean;
  deleteEmptyCursorBlock?: boolean;
  streamToolsProvider?: StreamToolsProvider<any, any>;
  documentStateBuilder?: DocumentStateBuilder<any>;
  onBlockUpdated?: (blockId: string) => void;
  onStart?: () => void;
}): Promise<AIRequest>;