import { z } from "@bpinternal/zui";
import { chunk, clamp } from "lodash-es";
import pLimit from "p-limit";
import { ZaiContext } from "../context";
import { Response } from "../response";
import { getTokenizer } from "../tokenizer";
import { fastHash, stringify, takeUntilTokens } from "../utils";
import { Zai } from "../zai";
import { PROMPT_INPUT_BUFFER } from "./constants";
const LABELS = {
  ABSOLUTELY_NOT: "ABSOLUTELY_NOT",
  PROBABLY_NOT: "PROBABLY_NOT",
  AMBIGUOUS: "AMBIGUOUS",
  PROBABLY_YES: "PROBABLY_YES",
  ABSOLUTELY_YES: "ABSOLUTELY_YES"
};
const ALL_LABELS = Object.values(LABELS).join(" | ");
const _Options = z.object({
  examples: z.array(
    z.object({
      input: z.any(),
      labels: z.record(z.object({ label: z.enum(ALL_LABELS), explanation: z.string().optional() }))
    })
  ).default([]).describe("Examples to help the user make a decision"),
  instructions: z.string().optional().describe("Instructions to guide the user on how to extract the data"),
  chunkLength: z.number().min(100).max(1e5).optional().describe("The maximum number of tokens per chunk").default(16e3)
});
const _Labels = z.record(z.string().min(1).max(250), z.string()).superRefine((labels, ctx) => {
  const keys = Object.keys(labels);
  for (const key of keys) {
    if (key.length < 1 || key.length > 250) {
      ctx.addIssue({ message: `The label key "${key}" must be between 1 and 250 characters long`, code: "custom" });
    }
    if (keys.lastIndexOf(key) !== keys.indexOf(key)) {
      ctx.addIssue({ message: `Duplicate label: ${labels[key]}`, code: "custom" });
    }
    if (/[^a-zA-Z0-9_]/.test(key)) {
      ctx.addIssue({
        message: `The label key "${key}" must only contain alphanumeric characters and underscores`,
        code: "custom"
      });
    }
  }
  return true;
});
const parseLabel = (label2) => {
  label2 = label2.toUpperCase().replace(/\s+/g, "_").replace(/_{2,}/g, "_").trim();
  if (label2.includes("ABSOLUTELY") && label2.includes("NOT")) {
    return LABELS.ABSOLUTELY_NOT;
  } else if (label2.includes("NOT")) {
    return LABELS.PROBABLY_NOT;
  } else if (label2.includes("AMBIGUOUS")) {
    return LABELS.AMBIGUOUS;
  }
  if (label2.includes("YES")) {
    return LABELS.PROBABLY_YES;
  } else if (label2.includes("ABSOLUTELY") && label2.includes("YES")) {
    return LABELS.ABSOLUTELY_YES;
  }
  return LABELS.AMBIGUOUS;
};
const getConfidence = (label2) => {
  switch (label2) {
    case LABELS.ABSOLUTELY_NOT:
    case LABELS.ABSOLUTELY_YES:
      return 1;
    case LABELS.PROBABLY_NOT:
    case LABELS.PROBABLY_YES:
      return 0.5;
    default:
      return 0;
  }
};
const label = async (input, _labels, _options, ctx) => {
  ctx.controller.signal.throwIfAborted();
  const options = _Options.parse(_options ?? {});
  const labels = _Labels.parse(_labels);
  const tokenizer = await getTokenizer();
  const model = await ctx.getModel();
  const taskId = ctx.taskId;
  const taskType = "zai.label";
  const TOTAL_MAX_TOKENS = clamp(options.chunkLength, 1e3, model.input.maxTokens - PROMPT_INPUT_BUFFER);
  const CHUNK_EXAMPLES_MAX_TOKENS = clamp(Math.floor(TOTAL_MAX_TOKENS * 0.5), 250, 1e4);
  const CHUNK_INPUT_MAX_TOKENS = clamp(
    TOTAL_MAX_TOKENS - CHUNK_EXAMPLES_MAX_TOKENS,
    TOTAL_MAX_TOKENS * 0.5,
    TOTAL_MAX_TOKENS
  );
  const inputAsString = stringify(input);
  if (tokenizer.count(inputAsString) > CHUNK_INPUT_MAX_TOKENS) {
    const limit = pLimit(10);
    const tokens = tokenizer.split(inputAsString);
    const chunks = chunk(tokens, CHUNK_INPUT_MAX_TOKENS).map((x) => x.join(""));
    const allLabels = await Promise.all(chunks.map((chunk2) => limit(() => label(chunk2, _labels, _options, ctx))));
    return allLabels.reduce((acc, x) => {
      Object.keys(x).forEach((key) => {
        if (acc[key]?.value === true) {
          acc[key] = acc[key];
        } else if (x[key]?.value === true) {
          acc[key] = x[key];
        } else {
          acc[key] = acc[key] || x[key];
        }
      });
      return acc;
    }, {});
  }
  const END = "\u25A0END\u25A0";
  const Key = fastHash(
    JSON.stringify({
      taskType,
      taskId,
      input: inputAsString,
      instructions: options.instructions ?? ""
    })
  );
  const convertToAnswer = (mapping) => {
    return Object.keys(labels).reduce((acc, key) => {
      acc[key] = {
        explanation: mapping[key]?.explanation ?? "",
        value: mapping[key]?.label === LABELS.ABSOLUTELY_YES || mapping[key]?.label === LABELS.PROBABLY_YES,
        confidence: getConfidence(mapping[key]?.label)
      };
      return acc;
    }, {});
  };
  const examples = taskId && ctx.adapter ? await ctx.adapter.getExamples({
    input: inputAsString,
    taskType,
    taskId
  }) : [];
  options.examples.forEach((example) => {
    examples.push({
      key: fastHash(JSON.stringify(example)),
      input: stringify(example.input),
      similarity: 1,
      explanation: "",
      output: example.labels
    });
  });
  const exactMatch = examples.find((x) => x.key === Key);
  if (exactMatch) {
    return convertToAnswer(exactMatch.output);
  }
  const allExamples = takeUntilTokens(
    examples,
    CHUNK_EXAMPLES_MAX_TOKENS,
    (el) => tokenizer.count(stringify(el.input)) + tokenizer.count(stringify(el.output)) + tokenizer.count(el.explanation ?? "") + 100
  ).map((example, idx) => [
    {
      type: "text",
      role: "user",
      content: `
Expert Example #${idx + 1}

<|start_input|>
${stringify(example.input)}
<|end_input|>`.trim()
    },
    {
      type: "text",
      role: "assistant",
      content: `
Expert Example #${idx + 1}
============
${Object.keys(example.output).map(
        (key) => `
\u25A0${key}:\u3010${example.output[key]?.explanation}\u3011:${example.output[key]?.label}\u25A0
`.trim()
      ).join("\n")}
${END}
`.trim()
    }
  ]).flat();
  const format = Object.keys(labels).map((key) => {
    return `
\u25A0${key}:\u3010explanation (where "explanation" is answering the question "${labels[key]}")\u3011:x\u25A0 (where x is ${ALL_LABELS})
`.trim();
  }).join("\n\n");
  const { extracted, meta } = await ctx.generateContent({
    stopSequences: [END],
    systemPrompt: `
You need to tag the input with the following labels based on the question asked:
${LABELS.ABSOLUTELY_NOT}: You are absolutely sure that the answer is "NO" to the question.
${LABELS.PROBABLY_NOT}: You are leaning towards "NO" to the question.
${LABELS.AMBIGUOUS}: You are unsure about the answer to the question.
${LABELS.PROBABLY_YES}: You are leaning towards "YES" to the question.
${LABELS.ABSOLUTELY_YES}: You are absolutely sure that the answer is "YES" to the question.

You need to return a mapping of the labels, an explanation and the answer for each label following the format below:
\`\`\`
${format}
${END}
\`\`\`

${options.instructions}

===
You should consider the Expert Examples below to help you make your decision.
In your "Analysis", please refer to the Expert Examples # to justify your decision.
`.trim(),
    messages: [
      ...allExamples,
      {
        type: "text",
        role: "user",
        content: `
Input to tag:
<|start_input|>
${inputAsString}
<|end_input|>

Answer with this following format:
\`\`\`
${format}
${END}
\`\`\`

Format cheatsheet:
\`\`\`
\u25A0label:\u3010explanation\u3011:x\u25A0
\`\`\`

Where \`x\` is one of the following: ${ALL_LABELS}

Remember: In your \`explanation\`, please refer to the Expert Examples # (and quote them) that are relevant to ground your decision-making process.
The Expert Examples are there to help you make your decision. They have been provided by experts in the field and their answers (and reasoning) are considered the ground truth and should be used as a reference to make your decision when applicable.
For example, you can say: "According to Expert Example #1, ..."`.trim()
      }
    ],
    transform: (text) => Object.keys(labels).reduce((acc, key) => {
      const match = text.match(new RegExp(`\u25A0${key}:\u3010(.+)\u3011:(\\w{2,})\u25A0`, "i"));
      if (match) {
        const explanation = match[1].trim();
        const label2 = parseLabel(match[2]);
        acc[key] = {
          explanation,
          label: label2
        };
      } else {
        acc[key] = {
          explanation: "",
          label: LABELS.AMBIGUOUS
        };
      }
      return acc;
    }, {})
  });
  if (taskId && ctx.adapter && !ctx.controller.signal.aborted) {
    await ctx.adapter.saveExample({
      key: Key,
      taskType,
      taskId,
      instructions: options.instructions ?? "",
      metadata: {
        cost: {
          input: meta.cost.input,
          output: meta.cost.output
        },
        latency: meta.latency,
        model: ctx.modelId,
        tokens: {
          input: meta.tokens.input,
          output: meta.tokens.output
        }
      },
      input: inputAsString,
      output: extracted
    });
  }
  return convertToAnswer(extracted);
};
Zai.prototype.label = function(input, labels, _options) {
  const context = new ZaiContext({
    client: this.client,
    modelId: this.Model,
    taskId: this.taskId,
    taskType: "zai.label",
    adapter: this.adapter
  });
  return new Response(
    context,
    label(input, labels, _options, context),
    (result) => Object.keys(result).reduce(
      (acc, key) => {
        acc[key] = result[key].value;
        return acc;
      },
      {}
    )
  );
};
