import { z } from "@bpinternal/zui";
import { chunk } from "lodash-es";
import pLimit from "p-limit";
import { ZaiContext } from "../context";
import { Response } from "../response";
import { getTokenizer } from "../tokenizer";
import { Zai } from "../zai";
import { PROMPT_INPUT_BUFFER, PROMPT_OUTPUT_BUFFER } from "./constants";
const Options = z.object({
  prompt: z.string().describe("What should the text be summarized to?").default("New information, concepts and ideas that are deemed important"),
  format: z.string().describe("How to format the example text").default(
    "A normal text with multiple sentences and paragraphs. Use markdown to format the text into sections. Use headings, lists, and other markdown features to make the text more readable. Do not include links, images, or other non-text elements."
  ),
  length: z.number().min(10).max(1e5).describe("The length of the summary in tokens").default(250),
  intermediateFactor: z.number().min(1).max(10).describe("How many times longer (than final length) are the intermediate summaries generated").default(4),
  maxIterations: z.number().min(1).default(100),
  sliding: z.object({
    window: z.number().min(10).max(1e5),
    overlap: z.number().min(0).max(1e5)
  }).describe("Sliding window options").default({ window: 5e4, overlap: 250 })
});
const START = "\u25A0START\u25A0";
const END = "\u25A0END\u25A0";
const summarize = async (original, options, ctx) => {
  ctx.controller.signal.throwIfAborted();
  const tokenizer = await getTokenizer();
  const model = await ctx.getModel();
  const INPUT_COMPONENT_SIZE = Math.max(100, (model.input.maxTokens - PROMPT_INPUT_BUFFER) / 4);
  options.prompt = tokenizer.truncate(options.prompt, INPUT_COMPONENT_SIZE);
  options.format = tokenizer.truncate(options.format, INPUT_COMPONENT_SIZE);
  const maxOutputSize = model.output.maxTokens - PROMPT_OUTPUT_BUFFER;
  if (options.length > maxOutputSize) {
    throw new Error(
      `The desired output length is ${maxOutputSize} tokens long, which is more than the maximum of ${model.output.maxTokens} tokens for this model (${model.name})`
    );
  }
  options.sliding.window = Math.min(options.sliding.window, model.input.maxTokens - PROMPT_INPUT_BUFFER);
  options.sliding.overlap = Math.min(options.sliding.overlap, options.sliding.window - 3 * options.sliding.overlap);
  const format = (summary, newText) => {
    return `
${START}
${summary.length ? summary : "<summary still empty>"}
${END}

Please amend the summary between the ${START} and ${END} tags to accurately reflect the prompt and the additional text below.

<|start_new_information|>
${newText}
<|new_information|>`.trim();
  };
  const tokens = tokenizer.split(original);
  const parts = Math.ceil(tokens.length / (options.sliding.window - options.sliding.overlap));
  let iteration = 0;
  const N = 2;
  const useMergeSort = parts >= Math.pow(2, N);
  const chunkSize = Math.ceil(tokens.length / (parts * N));
  if (useMergeSort) {
    const limit = pLimit(10);
    const chunks = chunk(tokens, chunkSize).map((x) => x.join(""));
    const allSummaries = (await Promise.allSettled(chunks.map((chunk2) => limit(() => summarize(chunk2, options, ctx))))).filter((x) => x.status === "fulfilled").map((x) => x.value);
    return summarize(allSummaries.join("\n\n============\n\n"), options, ctx);
  }
  const summaries = [];
  let currentSummary = "";
  for (let i = 0; i < tokens.length; i += options.sliding.window) {
    const from = Math.max(0, i - options.sliding.overlap);
    const to = Math.min(tokens.length, i + options.sliding.window + options.sliding.overlap);
    const isFirst = i === 0;
    const isLast = to >= tokens.length;
    const slice = tokens.slice(from, to).join("");
    if (iteration++ >= options.maxIterations) {
      break;
    }
    const instructions = [
      `At each step, you will receive a part of the text to summarize. Make sure to reply with the new summary in the tags ${START} and ${END}.`,
      "Summarize the text and make sure that the main points are included.",
      "Ignore any unnecessary details and focus on the main points.",
      "Use short and concise sentences to increase readability and information density.",
      "When looking at the new information, focus on: " + options.prompt
    ];
    if (isFirst) {
      instructions.push(
        "The current summary is empty. You need to generate a summary that covers the main points of the text."
      );
    }
    let generationLength = options.length;
    if (!isLast) {
      generationLength = Math.min(
        tokenizer.count(currentSummary) + options.length * options.intermediateFactor,
        maxOutputSize
      );
      instructions.push(
        "You need to amend the summary to include the new information. Make sure the summary is complete and covers all the main points."
      );
      instructions.push(`The current summary is ${currentSummary.length} tokens long.`);
      instructions.push(`You can amend the summary to be up to ${generationLength} tokens long.`);
    }
    if (isLast) {
      instructions.push(
        "This is the last part you will have to summarize. Make sure the summary is complete and covers all the main points."
      );
      instructions.push(
        `The current summary is ${currentSummary.length} tokens long. You need to make sure it is ${options.length} tokens or less.`
      );
      if (currentSummary.length > options.length) {
        instructions.push(
          `The current summary is already too long, so you need to shorten it to ${options.length} tokens while also including the new information.`
        );
      }
    }
    let { extracted: result } = await ctx.generateContent({
      systemPrompt: `
You are summarizing a text. The text is split into ${parts} parts, and you are currently working on part ${iteration}.
At every step, you will receive the current summary and a new part of the text. You need to amend the summary to include the new information (if needed).
The summary needs to cover the main points of the text and must be concise.

IMPORTANT INSTRUCTIONS:
${instructions.map((x) => `- ${x.trim()}`).join("\n")}

FORMAT OF THE SUMMARY:
${options.format}
`.trim(),
      messages: [{ type: "text", content: format(currentSummary, slice), role: "user" }],
      maxTokens: generationLength,
      stopSequences: [END],
      transform: (text) => {
        if (!text.trim().length) {
          throw new Error("The model did not return a valid summary. The response was empty.");
        }
        return text;
      }
    });
    if (result.includes(START)) {
      result = result.slice(result.indexOf(START) + START.length);
    }
    if (result.includes("\u25A0")) {
      result = result.slice(0, result.indexOf("\u25A0"));
    }
    summaries.push(result);
    currentSummary = result;
  }
  return currentSummary.trim();
};
Zai.prototype.summarize = function(original, _options) {
  const options = Options.parse(_options ?? {});
  const context = new ZaiContext({
    client: this.client,
    modelId: this.Model,
    taskId: this.taskId,
    taskType: "summarize",
    adapter: this.adapter
  });
  return new Response(context, summarize(original, options, context), (value) => value);
};
