import {
  addStoppingStrings,
  getLastChatMessage,
  popLastAssistantMessage,
  popLastChatMessages,
  replaceTemplates,
} from "../src/utils.mjs";

export default ({ user, assistant, messages, config, generationConfig }) => {
  addStoppingStrings(config, ["\nUSER:", "\nASSISTANT:"]);

  const ooc = "";
  const systemPrompt =
    "";
  const newConversation = `USER: ${ooc}Let's start a new roleplay.`;
  const newExample = `USER: ${ooc}Here are some old roleplays that we did before.`;
  const context = `USER: `;//`USER: ${ooc}I'm roleplaying as ${user} and you're roleplaying as ${assistant}. Here's the context for this roleplay:\n`;
  const contextResponse = `${ooc}Okay. I will take that info into account to roleplay as ${assistant}.`;
  const characterBias = replaceTemplates(config.characterBias, config);
  const impersonationPrompt = replaceTemplates(
    config.impersonationPrompt,
    config
  );
  const silentMessage = replaceTemplates(config.silentMessage, config);

  let impersonationPromptFound = false;

  const addNames = false;
  const userName = () => (addNames ? `${user}: ` : "");
  const assistantName = () => (addNames ? `${assistant}: ` : "");

  const beforeSystem = "\n\n";
  const afterSystem = "\n";
  const beforeUser = "\n\nUSER: ";
  const afterUser = "\n";
  const beforeAssistant = "\n\nASSISTANT: ";
  const afterAssistant = config.backendType === "koboldcpp" ? "\n" : "</s>\n";

  const addReplyInstruction = true;
  const addFinalInstruction = false;
  const replyInstruction = ({
    you,
  }) => `${beforeUser}${ooc}Write a continuation for this roleplay, follow these rules:
- The plot is developed slowly.
- Flesh out every scene with copious detail, narration, and dialogue.
- Never speak or act for ${user}.${afterUser}${beforeAssistant}${ooc}Okay. I will follow these rules and the description above. The most engaging, descriptive and creative continuation for this roleplay is this:${afterAssistant}`;
  const finalInstruction = ({ you }) =>
    `${beforeUser}${ooc}Continue the roleplay. Stay in character and write at least two paragraphs.${afterUser}`;

  let prompt = [];
  if (systemPrompt) {
    prompt.push({
      role: "system",
      metadata: { type: "system-prompt" },
      prunable: false,
      content: `${beforeSystem}${systemPrompt}${afterSystem}`,
    });
  }

  for (const msg of messages) {
    const { metadata } = msg;
    let content = msg.content.trim();

    if (metadata.type === "new-conversation") {
      if (newConversation) {
        prompt.push({
          ...msg,
          prunable: false,
          content: `${beforeSystem}${newConversation}${afterSystem}`,
        });
      }
    } else if (metadata.type === "new-example-dialogue") {
      if (newExample && metadata.chatIndex === 0) {
        prompt.push({
          ...msg,
          prunable: false,
          content: `${beforeSystem}${newExample}${afterSystem}`,
        });
      }
    } else if (metadata.type === "context") {
      prompt.push({
        ...msg,
        prunable: false,
        content: `${beforeSystem}${context}${content}${afterSystem}`,
      });
      if (contextResponse) {
        prompt.push({
          role: "assistant",
          metadata: { type: "context-response" },
          prunable: false,
          content: `${beforeAssistant}${contextResponse}${afterAssistant}`,
        });
      }
    } else if (metadata.type === "example-assistant") {
      prompt.push({
        ...msg,
        prunable: !(
          config.keepExampleMessagesInPrompt ||
          metadata.exampleAssistantMsgIndex === 0
        ),
        content: `${beforeAssistant}${assistantName()}${content}${afterAssistant}`,
      });
    } else if (metadata.type === "example-user") {
      prompt.push({
        ...msg,
        prunable: !config.keepExampleMessagesInPrompt,
        content: `${beforeUser}${userName()}${content}${afterUser}`,
      });
    } else if (metadata.type === "other" || metadata.type === "jailbreak") {
      prompt.push({
        ...msg,
        prunable: false,
        content: `${beforeSystem}${content}${afterSystem}`,
      });
    } else if (metadata.type === "impersonation-prompt") {
      impersonationPromptFound = true;
    } else if (metadata.type === "assistant-msg") {
      prompt.push({
        ...msg,
        prunable: true,
        content: `${beforeAssistant}${assistantName()}${content}${afterAssistant}`,
      });
    } else if (metadata.type === "user-msg") {
      prompt.push({
        ...msg,
        prunable: true,
        content: `${beforeUser}${userName()}${content}${afterUser}`,
      });
    }
  }

  const last = getLastChatMessage(prompt);
  const lastMessages = popLastChatMessages(prompt, 2);

  const you = impersonationPromptFound ? user : assistant;

  prompt.push({
    role: "system",
    metadata: { type: "superbig-injection-point" },
    prunable: true,
    content: "",
  });

  if (addReplyInstruction) {
    prompt.push({
      role: "system",
      metadata: { type: "reply-instruction" },
      prunable: false,
      content: replyInstruction({ you }),
    });
  }

  for (const msg of lastMessages) {
    prompt.push(msg);
  }

  if (impersonationPromptFound || last?.role === "user" || silentMessage) {
    if (last?.role === "assistant" && silentMessage) {
      prompt.push({
        role: "user",
        metadata: { type: "silent-message" },
        prunable: false,
        content: `${beforeUser}${userName()}${silentMessage}${afterUser}`,
      });
    }

    if (impersonationPromptFound) {
      prompt.push({
        role: "system",
        metadata: { type: "impersonation-prompt" },
        prunable: false,
        content: `${beforeSystem}${impersonationPrompt}${afterSystem}`,
      });
    }

    if (addFinalInstruction) {
      prompt.push({
        role: "system",
        metadata: { type: "reply-instruction" },
        prunable: false,
        content: finalInstruction({ you }),
      });
    }

    prompt.push({
      role: impersonationPromptFound ? "user" : "assistant",
      metadata: { type: "reply-to-complete" },
      prunable: false,
      content:
        `${impersonationPromptFound ? beforeUser : beforeAssistant}${
          impersonationPromptFound ? userName() : assistantName()
        }`.trimRight() + characterBias,
    });
  } else {
    const msg = popLastAssistantMessage(prompt);
    const end = msg.content.length - afterAssistant.length;
    msg.content = msg.content.substring(0, end);
    prompt.push(msg);
  }

  if (impersonationPromptFound) {
    generationConfig.max_new_tokens = config.impersonationMaxNewTokens;
  }

  return prompt;
};