import { z } from "@bpinternal/zui";
import pLimit from "p-limit";
import { ZaiContext } from "../context";
import { Response } from "../response";
import { getTokenizer } from "../tokenizer";
import { fastHash, stringify } from "../utils";
import { Zai } from "../zai";
import { PROMPT_INPUT_BUFFER, PROMPT_OUTPUT_BUFFER } from "./constants";
const RATING_VALUES = {
  very_bad: 1,
  bad: 2,
  average: 3,
  good: 4,
  very_good: 5
};
const _Options = z.object({
  tokensPerItem: z.number().min(1).max(1e5).optional().describe("The maximum number of tokens per item").default(250),
  maxItemsPerChunk: z.number().min(1).max(100).optional().describe("The maximum number of items to rate per chunk").default(50)
});
const END = "\u25A0END\u25A0";
const rate = async (input, instructions, _options, ctx) => {
  ctx.controller.signal.throwIfAborted();
  const options = _Options.parse(_options ?? {});
  const tokenizer = await getTokenizer();
  const model = await ctx.getModel();
  if (input.length === 0) {
    return [];
  }
  const taskId = ctx.taskId;
  const taskType = "zai.rate";
  const TOKENS_TOTAL_MAX = model.input.maxTokens - PROMPT_INPUT_BUFFER - PROMPT_OUTPUT_BUFFER;
  const isStringInstructions = typeof instructions === "string";
  const criteriaKeys = isStringInstructions ? [] : Object.keys(instructions);
  const generateCriteriaPrompt = isStringInstructions ? `Generate 3-5 evaluation criteria for: "${instructions}"

For each criterion, provide 5 labels (very_bad, bad, average, good, very_good) with brief descriptions.

Output format (JSON):
{
  "criterion1_name": {
    "very_bad": "description",
    "bad": "description",
    "average": "description",
    "good": "description",
    "very_good": "description"
  },
  "criterion2_name": { ... }
}

Keep criterion names short (1-2 words, lowercase, use underscores).
Keep descriptions brief (5-10 words each).` : `For these evaluation criteria, provide 5 labels (very_bad, bad, average, good, very_good) with brief descriptions for each:

${criteriaKeys.map((key) => `- ${key}: ${instructions[key]}`).join("\n")}

Output format (JSON):
{
  "${criteriaKeys[0]}": {
    "very_bad": "description",
    "bad": "description",
    "average": "description",
    "good": "description",
    "very_good": "description"
  }
  ${criteriaKeys.length > 1 ? "..." : ""}
}

Keep descriptions brief (5-10 words each).`;
  const { extracted: evaluationCriteria } = await ctx.generateContent({
    systemPrompt: `You are creating evaluation criteria for rating items on a 1-5 scale.
Each criterion must have exactly 5 labels: very_bad (1), bad (2), average (3), good (4), very_good (5).
Output valid JSON only.`,
    messages: [
      {
        type: "text",
        role: "user",
        content: generateCriteriaPrompt
      }
    ],
    transform: (text) => {
      const jsonMatch = text.match(/```(?:json)?\s*(\{[\s\S]*?\})\s*```/) || text.match(/(\{[\s\S]*\})/);
      if (!jsonMatch) {
        throw new Error("Failed to parse evaluation criteria JSON");
      }
      return JSON.parse(jsonMatch[1]);
    }
  });
  const finalCriteriaKeys = Object.keys(evaluationCriteria);
  if (finalCriteriaKeys.length === 0) {
    throw new Error("No evaluation criteria generated");
  }
  const TOKENS_CRITERIA_MAX = Math.floor(TOKENS_TOTAL_MAX * 0.3);
  const TOKENS_ITEMS_MAX = TOKENS_TOTAL_MAX - TOKENS_CRITERIA_MAX;
  let chunks = [];
  let currentChunk = [];
  let currentChunkTokens = 0;
  for (const element of input) {
    const elementAsString = tokenizer.truncate(stringify(element, false), options.tokensPerItem);
    const elementTokens = tokenizer.count(elementAsString);
    if (currentChunkTokens + elementTokens > TOKENS_ITEMS_MAX || currentChunk.length >= options.maxItemsPerChunk) {
      if (currentChunk.length > 0) {
        chunks.push(currentChunk);
      }
      currentChunk = [];
      currentChunkTokens = 0;
    }
    currentChunk.push(element);
    currentChunkTokens += elementTokens;
  }
  if (currentChunk.length > 0) {
    chunks.push(currentChunk);
  }
  chunks = chunks.filter((x) => x.length > 0);
  const rateChunk = async (chunk) => {
    ctx.controller.signal.throwIfAborted();
    const chunkInputStr = JSON.stringify(chunk);
    const examples = taskId && ctx.adapter ? await ctx.adapter.getExamples({
      input: chunkInputStr.slice(0, 1e3),
      // Limit search string length
      taskType,
      taskId
    }) : [];
    const key = fastHash(
      stringify({
        taskId,
        taskType,
        input: chunkInputStr,
        instructions: stringify(instructions)
      })
    );
    const exactMatch = examples.find((x) => x.key === key);
    if (exactMatch && exactMatch.output) {
      return {
        ratings: exactMatch.output,
        meta: { cost: { input: 0, output: 0 }, latency: 0, tokens: { input: 0, output: 0 } }
      };
    }
    const formatCriteria = () => {
      return finalCriteriaKeys.map((key2) => {
        const labels = evaluationCriteria[key2];
        return `**${key2}**:
  - very_bad (1): ${labels?.very_bad}
  - bad (2): ${labels?.bad}
  - average (3): ${labels?.average}
  - good (4): ${labels?.good}
  - very_good (5): ${labels?.very_good}`;
      }).join("\n\n");
    };
    const formatItems = (items) => {
      return items.map((item, idx) => {
        const itemStr = tokenizer.truncate(stringify(item, false), options.tokensPerItem);
        return `\u25A0${idx}: ${itemStr}\u25A0`;
      }).join("\n");
    };
    const exampleMessages = [];
    for (const example of examples.slice(0, 5)) {
      try {
        const exampleInput = JSON.parse(example.input);
        exampleMessages.push({
          type: "text",
          role: "user",
          content: `Expert Example - Items to rate:
${formatItems(Array.isArray(exampleInput) ? exampleInput : [exampleInput])}

Rate each item on all criteria.`
        });
        const exampleOutput = example.output;
        if (Array.isArray(exampleOutput) && exampleOutput.length > 0) {
          const formattedRatings = exampleOutput.map((rating, idx) => {
            const pairs = finalCriteriaKeys.map((key2) => {
              const value = rating[key2];
              if (typeof value === "number") {
                const labelMap = {
                  1: "very_bad",
                  2: "bad",
                  3: "average",
                  4: "good",
                  5: "very_good"
                };
                return `${key2}=${labelMap[value] || "average"}`;
              }
              return null;
            }).filter(Boolean).join(";");
            return `\u25A0${idx}:${pairs}\u25A0`;
          }).join("\n");
          exampleMessages.push({
            type: "text",
            role: "assistant",
            content: `${formattedRatings}
${END}`
          });
          if (example.explanation) {
            exampleMessages.push({
              type: "text",
              role: "assistant",
              content: `Reasoning: ${example.explanation}`
            });
          }
        }
      } catch {
      }
    }
    const { extracted, meta } = await ctx.generateContent({
      systemPrompt: `You are rating items based on evaluation criteria.

Evaluation Criteria:
${formatCriteria()}

For each item, rate it on EACH criterion using one of these labels:
very_bad, bad, average, good, very_good

Output format:
\u25A00:criterion1=label;criterion2=label;criterion3=label\u25A0
\u25A01:criterion1=label;criterion2=label;criterion3=label\u25A0
${END}

IMPORTANT:
- Rate every item (\u25A00 to \u25A0${chunk.length - 1})
- Use exact criterion names: ${finalCriteriaKeys.join(", ")}
- Use exact label names: very_bad, bad, average, good, very_good
- Use semicolons (;) between criteria
- Use equals (=) between criterion and label`,
      stopSequences: [END],
      messages: [
        ...exampleMessages,
        {
          type: "text",
          role: "user",
          content: `Items to rate (\u25A00 to \u25A0${chunk.length - 1}):
${formatItems(chunk)}

Rate each item on all criteria.
Output format: \u25A0index:criterion1=label;criterion2=label\u25A0
${END}`
        }
      ],
      transform: (text) => {
        const results = [];
        const regex = /■(\d+):([^■]+)■/g;
        let match;
        while ((match = regex.exec(text)) !== null) {
          const idx = parseInt(match[1] ?? "", 10);
          const ratingsStr = match[2] ?? "";
          if (isNaN(idx) || idx < 0 || idx >= chunk.length) {
            continue;
          }
          const itemRatings = {};
          let total = 0;
          const pairs = ratingsStr.split(";").filter((x) => x.trim().length > 0);
          for (const pair of pairs) {
            const [criterion, label] = pair.split("=").map((x) => x.trim());
            if (!criterion || !label) continue;
            const labelLower = label.toLowerCase().replace(/\s+/g, "_");
            const ratingValue = RATING_VALUES[labelLower] ?? 3;
            itemRatings[criterion] = ratingValue;
            total += ratingValue;
          }
          itemRatings.total = total;
          results[idx] = itemRatings;
        }
        for (let i = 0; i < chunk.length; i++) {
          if (!results[i]) {
            const defaultRatings = {};
            let total = 0;
            for (const key2 of finalCriteriaKeys) {
              defaultRatings[key2] = 3;
              total += 3;
            }
            defaultRatings.total = total;
            results[i] = defaultRatings;
          }
        }
        return results;
      }
    });
    return { ratings: extracted, meta };
  };
  const limit = pLimit(10);
  const chunkPromises = chunks.map((chunk) => limit(() => rateChunk(chunk)));
  const ratedChunks = await Promise.all(chunkPromises);
  const allRatings = ratedChunks.flatMap((result) => result.ratings);
  const totalMeta = ratedChunks.reduce(
    (acc, result) => ({
      cost: {
        input: acc.cost.input + result.meta.cost.input,
        output: acc.cost.output + result.meta.cost.output
      },
      latency: Math.max(acc.latency, result.meta.latency),
      // Use max latency
      tokens: {
        input: acc.tokens.input + result.meta.tokens.input,
        output: acc.tokens.output + result.meta.tokens.output
      }
    }),
    {
      cost: { input: 0, output: 0 },
      latency: 0,
      tokens: { input: 0, output: 0 }
    }
  );
  if (taskId && ctx.adapter && !ctx.controller.signal.aborted) {
    const key = fastHash(
      stringify({
        taskId,
        taskType,
        input: JSON.stringify(input),
        instructions: stringify(instructions)
      })
    );
    await ctx.adapter.saveExample({
      key,
      taskType,
      taskId,
      input: JSON.stringify(input),
      output: allRatings,
      instructions: typeof instructions === "string" ? instructions : JSON.stringify(instructions),
      metadata: {
        cost: {
          input: totalMeta.cost.input,
          output: totalMeta.cost.output
        },
        latency: totalMeta.latency,
        model: ctx.modelId,
        tokens: {
          input: totalMeta.tokens.input,
          output: totalMeta.tokens.output
        }
      }
    });
  }
  return allRatings;
};
Zai.prototype.rate = function(input, instructions, _options) {
  const context = new ZaiContext({
    client: this.client,
    modelId: this.Model,
    taskId: this.taskId,
    taskType: "zai.rate",
    adapter: this.adapter
  });
  return new Response(
    context,
    rate(input, instructions, _options, context),
    (result) => {
      if (typeof instructions === "string") {
        return result.map((r) => r.total);
      }
      return result;
    }
  );
};
