import { COPILOT_CLOUD_PUBLIC_API_KEY_HEADER } from "@copilotkit/shared";
import { CopilotContext } from "@copilotkit/react-core";
import { useCallback, useContext } from "react";
import {
  CopilotRuntimeClient,
  Message,
  Role,
  TextMessage,
  convertGqlOutputToMessages,
  convertMessagesToGqlInput,
  CopilotRequestType,
} from "@copilotkit/runtime-client-gql";
import { retry } from "../../lib/retry";
import {
  EditingEditorState,
  Generator_InsertionOrEditingSuggestion,
} from "../../types/base/autosuggestions-bare-function";
import { InsertionsApiConfig } from "../../types/autosuggestions-config/insertions-api-config";
import { EditingApiConfig } from "../../types/autosuggestions-config/editing-api-config";
import { DocumentPointer } from "@copilotkit/react-core";
import { randomId } from "@copilotkit/shared";

/**
 * Returns a memoized function that sends a request to the specified API endpoint to get an autosuggestion for the user's input.
 * The function takes in the text before and after the cursor, and an abort signal.
 * It sends a POST request to the API endpoint with the messages array containing the system message, few shot messages, and user messages.
 * The function returns the suggestion from the API response.
 *
 * @param textareaPurpose - The purpose of the textarea. This is included in the system message.
 * @param apiEndpoint - The API endpoint to send the autosuggestion request to.
 * @param makeSystemMessage - A function that takes in a context string and returns a system message to include in the autosuggestion request.
 * @param fewShotMessages - An array of few shot messages to include in the autosuggestion request.
 * @param contextCategories - The categories of context strings we want to include. By default, we include the (default) "global" context category.
 * @returns A memoized function that sends a request to the specified API endpoint to get an autosuggestion for the user's input.
 */
export function useMakeStandardInsertionOrEditingFunction(
  textareaPurpose: string,
  contextCategories: string[],
  insertionApiConfig: InsertionsApiConfig,
  editingApiConfig: EditingApiConfig,
): Generator_InsertionOrEditingSuggestion {
  const { getContextString, copilotApiConfig } = useContext(CopilotContext);
  const headers = {
    ...(copilotApiConfig.publicApiKey
      ? { [COPILOT_CLOUD_PUBLIC_API_KEY_HEADER]: copilotApiConfig.publicApiKey }
      : {}),
  };

  const runtimeClient = new CopilotRuntimeClient({
    url: copilotApiConfig.chatApiEndpoint,
    publicApiKey: copilotApiConfig.publicApiKey,
    headers,
    credentials: copilotApiConfig.credentials,
  });

  async function runtimeClientResponseToStringStream(
    responsePromise: ReturnType<typeof runtimeClient.generateCopilotResponse>,
  ) {
    const messagesStream = await CopilotRuntimeClient.asStream(responsePromise);

    return new ReadableStream({
      async start(controller) {
        const reader = messagesStream.getReader();
        let sentContent = "";

        while (true) {
          const { done, value } = await reader.read();
          if (done) {
            break;
          }

          const messages = convertGqlOutputToMessages(value.generateCopilotResponse.messages);

          let newContent = "";

          for (const message of messages) {
            if (message instanceof TextMessage) {
              newContent += message.content;
            }
          }
          if (newContent) {
            const contentToSend = newContent.slice(sentContent.length);
            controller.enqueue(contentToSend);
            sentContent += contentToSend;
          }
        }
        controller.close();
      },
    });
  }

  const insertionFunction = useCallback(
    async (
      editorState: EditingEditorState,
      insertionPrompt: string,
      documents: DocumentPointer[],
      abortSignal: AbortSignal,
    ) => {
      const res = await retry(async () => {
        const messages: Message[] = [
          new TextMessage({
            role: Role.System,
            content: insertionApiConfig.makeSystemPrompt(
              textareaPurpose,
              getContextString(documents, contextCategories),
            ),
          }),
          ...insertionApiConfig.fewShotMessages,
          new TextMessage({
            role: Role.User,
            content: `<TextAfterCursor>${editorState.textAfterCursor}</TextAfterCursor>`,
          }),
          new TextMessage({
            role: Role.User,
            content: `<TextBeforeCursor>${editorState.textBeforeCursor}</TextBeforeCursor>`,
          }),
          new TextMessage({
            role: Role.User,
            content: `<InsertionPrompt>${insertionPrompt}</InsertionPrompt>`,
          }),
        ];

        return runtimeClientResponseToStringStream(
          runtimeClient.generateCopilotResponse({
            data: {
              frontend: {
                actions: [],
              },
              messages: convertMessagesToGqlInput(messages),
              metadata: {
                requestType: CopilotRequestType.TextareaCompletion,
              },
            },
            properties: copilotApiConfig.properties,
            signal: abortSignal,
          }),
        );
      });

      return res;
    },
    [insertionApiConfig, getContextString, contextCategories, textareaPurpose],
  );

  const editingFunction = useCallback(
    async (
      editorState: EditingEditorState,
      editingPrompt: string,
      documents: DocumentPointer[],
      abortSignal: AbortSignal,
    ) => {
      const res = await retry(async () => {
        const messages: Message[] = [
          new TextMessage({
            role: Role.System,
            content: editingApiConfig.makeSystemPrompt(
              textareaPurpose,
              getContextString(documents, contextCategories),
            ),
          }),
          ...editingApiConfig.fewShotMessages,
          new TextMessage({
            role: Role.User,
            content: `<TextBeforeCursor>${editorState.textBeforeCursor}</TextBeforeCursor>`,
          }),
          new TextMessage({
            role: Role.User,
            content: `<TextToEdit>${editorState.selectedText}</TextToEdit>`,
          }),
          new TextMessage({
            role: Role.User,
            content: `<TextAfterCursor>${editorState.textAfterCursor}</TextAfterCursor>`,
          }),
          new TextMessage({
            role: Role.User,
            content: `<EditingPrompt>${editingPrompt}</EditingPrompt>`,
          }),
        ];

        const runtimeClient = new CopilotRuntimeClient({
          url: copilotApiConfig.chatApiEndpoint,
          publicApiKey: copilotApiConfig.publicApiKey,
          headers,
          credentials: copilotApiConfig.credentials,
        });

        return runtimeClientResponseToStringStream(
          runtimeClient.generateCopilotResponse({
            data: {
              frontend: {
                actions: [],
              },
              messages: convertMessagesToGqlInput(messages),
              metadata: {
                requestType: CopilotRequestType.TextareaCompletion,
              },
            },
            properties: copilotApiConfig.properties,
            signal: abortSignal,
          }),
        );
      });

      return res;
    },
    [editingApiConfig, getContextString, contextCategories, textareaPurpose],
  );

  const insertionOrEditingFunction = useCallback(
    async (
      editorState: EditingEditorState,
      insertionPrompt: string,
      documents: DocumentPointer[],
      abortSignal: AbortSignal,
    ) => {
      if (editorState.selectedText === "") {
        return await insertionFunction(editorState, insertionPrompt, documents, abortSignal);
      } else {
        return await editingFunction(editorState, insertionPrompt, documents, abortSignal);
      }
    },
    [insertionFunction, editingFunction],
  );

  return insertionOrEditingFunction;
}