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 _Options = z.object({
  tokensPerItem: z.number().min(1).max(1e5).optional().describe("The maximum number of tokens per item").default(250)
});
const END = "\u25A0END\u25A0";
const sort = async (input, instructions, _options, ctx) => {
  ctx.controller.signal.throwIfAborted();
  const options = _Options.parse(_options ?? {});
  const tokenizer = await getTokenizer();
  const model = await ctx.getModel();
  const taskId = ctx.taskId;
  const taskType = "zai.sort";
  if (input.length === 0) {
    return [];
  }
  if (input.length === 1) {
    return input;
  }
  const TOKENS_TOTAL_MAX = model.input.maxTokens - PROMPT_INPUT_BUFFER - PROMPT_OUTPUT_BUFFER;
  const sampleSize = Math.min(5, input.length);
  const sampleItems = input.slice(0, sampleSize);
  const sampleItemsText = sampleItems.map((item, idx) => `\u25A0${idx}: ${stringify(item, false)}`).join("\n");
  const criteriaInputStr = JSON.stringify({ instructions, sampleItems });
  const criteriaExamples = taskId && ctx.adapter ? await ctx.adapter.getExamples({
    input: criteriaInputStr.slice(0, 1e3),
    taskType,
    taskId
  }) : [];
  const criteriaExampleMessages = [];
  for (const example of criteriaExamples.slice(0, 2)) {
    try {
      const exampleInput = JSON.parse(example.input);
      const exampleItems = Array.isArray(exampleInput) ? exampleInput : [exampleInput];
      criteriaExampleMessages.push({
        type: "text",
        role: "user",
        content: `Expert Example - Analyze sorting instruction: "${instructions}"

Sample items:
${exampleItems.slice(0, 3).map((el, i) => `\u25A0${i}: ${stringify(el, false).slice(0, 200)}`).join("\n")}

Generate sorting criteria.`
      });
      if (example.explanation) {
        criteriaExampleMessages.push({
          type: "text",
          role: "assistant",
          content: `${example.explanation}
${END}`
        });
      }
    } catch {
    }
  }
  const generateCriteriaPrompt = `Analyze this sorting instruction: "${instructions}"

Sample items to be sorted:
${sampleItemsText}

Create 1-3 sorting criteria with ordered label arrays (3-10 labels each).

**CRITICAL RULES**:
1. Labels are single words, lowercase, no spaces, use underscores
2. Labels are ordered from FIRST to LAST in sorted result
3. If instruction says "from X to Y": first label represents X, last label represents Y
4. If instruction says "prioritize" or "highest/lowest priority":
   - First label = HIGHEST priority (top of todo list)
   - Last label = LOWEST priority (bottom of todo list)

Examples:

"from slowest to fastest" \u2192 first=slowest, last=fastest
\u25A0speed\u25A0
very_slow;slow;medium;fast;very_fast
\u25A0END\u25A0

"from most dangerous to least dangerous" \u2192 first=most dangerous, last=least dangerous
\u25A0danger\u25A0
extremely_dangerous;very_dangerous;dangerous;moderate;slightly_dangerous;harmless
\u25A0END\u25A0

"from least urgent (spam) to most urgent (bills)" \u2192 first=spam, last=bills
\u25A0urgency\u25A0
spam;promotional;normal;important;urgent;critical
\u25A0END\u25A0

"prioritize: highest priority=open old tickets; lowest priority=closed" \u2192 first=high priority, last=low priority
\u25A0status\u25A0
open_old;open_recent;closed
\u25A0age\u25A0
oldest;old;recent;new
\u25A0END\u25A0

Output format:
\u25A0criterion_name\u25A0
label1;label2;label3;label4
\u25A0END\u25A0

Use 3-10 labels per criterion. Labels should be intuitive and match the domain.
Keep criterion names short (1-2 words, lowercase, underscores).
`;
  const { extracted: sortingCriteria } = await ctx.generateContent({
    systemPrompt: `You are creating sorting criteria with ordered label arrays.

CRITICAL: Output ordered labels from FIRST to LAST position in sorted result.
- Labels are single words, lowercase, underscores only
- 3-10 labels per criterion
- Order matters: first label = appears first, last label = appears last`,
    messages: [
      ...criteriaExampleMessages,
      {
        type: "text",
        role: "user",
        content: generateCriteriaPrompt
      }
    ],
    transform: (text) => {
      const criteria = {};
      const criterionRegex = /■([^■]+)■\s*([^\n■]+)/g;
      let match;
      while ((match = criterionRegex.exec(text)) !== null) {
        const name = (match[1] ?? "").trim().toLowerCase();
        const labelsStr = (match[2] ?? "").trim();
        if (!name || name === "end") continue;
        const labels = labelsStr.split(";").map((l) => l.trim().toLowerCase().replace(/\s+/g, "_")).filter((l) => l.length > 0 && l.length < 50);
        if (labels.length >= 3 && labels.length <= 10) {
          criteria[name] = {
            description: `${labels.length} ordered labels`,
            labels
          };
        }
      }
      if (Object.keys(criteria).length === 0) {
        throw new Error(`Failed to parse sorting criteria. LLM output: ${text.slice(0, 500)}`);
      }
      return criteria;
    }
  });
  const criteriaKeys = Object.keys(sortingCriteria);
  if (criteriaKeys.length === 0) {
    throw new Error("No sorting criteria generated");
  }
  const TOKENS_CRITERIA_MAX = Math.floor(TOKENS_TOTAL_MAX * 0.2);
  const TOKENS_ITEMS_MAX = TOKENS_TOTAL_MAX - TOKENS_CRITERIA_MAX;
  const MAX_ITEMS_PER_CHUNK = 50;
  const elements = input.map((element, idx) => ({
    element,
    index: idx,
    stringified: stringify(element, false)
  }));
  const chunks = [];
  let currentChunk = [];
  let currentTokens = 0;
  for (const elem of elements) {
    const truncated = tokenizer.truncate(elem.stringified, options.tokensPerItem);
    const elemTokens = tokenizer.count(truncated);
    if ((currentTokens + elemTokens > TOKENS_ITEMS_MAX || currentChunk.length >= MAX_ITEMS_PER_CHUNK) && currentChunk.length > 0) {
      chunks.push(currentChunk);
      currentChunk = [];
      currentTokens = 0;
    }
    currentChunk.push(elem);
    currentTokens += elemTokens;
  }
  if (currentChunk.length > 0) {
    chunks.push(currentChunk);
  }
  const scoreChunk = async (chunk) => {
    ctx.controller.signal.throwIfAborted();
    const chunkSize = chunk.length;
    const chunkInputStr = JSON.stringify(chunk.map((c) => c.element));
    const examples = taskId && ctx.adapter ? await ctx.adapter.getExamples({
      input: chunkInputStr.slice(0, 1e3),
      taskType,
      taskId
    }) : [];
    const key = fastHash(
      stringify({
        taskId,
        taskType,
        input: chunkInputStr,
        instructions
      })
    );
    const exactMatch = examples.find((x) => x.key === key);
    if (exactMatch && exactMatch.output) {
      return exactMatch.output;
    }
    const elementsText = chunk.map((elem, i) => {
      const truncated = tokenizer.truncate(elem.stringified, options.tokensPerItem);
      return `\u25A0${i}: ${truncated}\u25A0`;
    }).join("\n");
    const criteriaText = criteriaKeys.map((key2) => {
      const criterion = sortingCriteria[key2];
      const labelsText = criterion.labels.join(";");
      return `**${key2}**: ${labelsText}`;
    }).join("\n");
    const exampleMessages = [];
    for (const example of examples.slice(0, 3)) {
      try {
        const exampleInput = JSON.parse(example.input);
        const exampleItems = Array.isArray(exampleInput) ? exampleInput : [exampleInput];
        exampleMessages.push({
          type: "text",
          role: "user",
          content: `Expert Example - Items to score:
${exampleItems.map((el, i) => `\u25A0${i}: ${stringify(el, false).slice(0, 200)}\u25A0`).join("\n")}

Score each item.`
        });
        const exampleOutput = example.output;
        if (Array.isArray(exampleOutput) && exampleOutput.length > 0) {
          const formattedScores = exampleOutput.map((score) => {
            const pairs = criteriaKeys.map((key2) => `${key2}=${score.scores[key2] ?? 0}`).join(";");
            return `\u25A0${score.elementIndex}:${pairs}\u25A0`;
          }).join("\n");
          exampleMessages.push({
            type: "text",
            role: "assistant",
            content: `${formattedScores}
${END}`
          });
          if (example.explanation) {
            exampleMessages.push({
              type: "text",
              role: "assistant",
              content: `Reasoning: ${example.explanation}`
            });
          }
        }
      } catch {
      }
    }
    const { extracted } = await ctx.generateContent({
      systemPrompt: `You are ranking items for sorting using ordered label arrays.

${criteriaText}

Instructions: "${instructions}"

SCORING RULES:
- For each item and each criterion, assign ONE label from the ordered list
- Labels are ordered: first label = appears FIRST in sorted result, last label = appears LAST
- Choose the label that best describes each item

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

IMPORTANT:
- Rank every item (\u25A00 to \u25A0${chunkSize - 1})
- Use exact criterion names: ${criteriaKeys.join(", ")}
- Use exact labels from the lists above (lowercase, underscores)
- Use semicolons (;) between criteria
- Use equals (=) between criterion and label`,
      stopSequences: [END],
      messages: [
        ...exampleMessages,
        {
          type: "text",
          role: "user",
          content: `Items to rank (\u25A00 to \u25A0${chunkSize - 1}):
${elementsText}

Rank each item using the labeled scales.
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 labelsStr = match[2] ?? "";
          if (isNaN(idx) || idx < 0 || idx >= chunkSize) continue;
          const scores = {};
          let total = 0;
          const pairs = labelsStr.split(";").filter((x) => x.trim().length > 0);
          for (const pair of pairs) {
            const [criterion, labelStr] = pair.split("=").map((x) => x.trim().toLowerCase().replace(/\s+/g, "_"));
            if (!criterion || !labelStr) continue;
            const labels = sortingCriteria[criterion]?.labels ?? [];
            const labelIndex = labels.findIndex((l) => l === labelStr);
            if (labelIndex >= 0) {
              scores[criterion] = labelIndex;
              total += labelIndex;
            } else {
              const middleIndex = labels.length > 0 ? Math.floor(labels.length / 2) : 5;
              scores[criterion] = middleIndex;
              total += middleIndex;
            }
          }
          results[idx] = {
            elementIndex: chunk[idx].index,
            scores,
            totalScore: total
          };
        }
        for (let i = 0; i < chunkSize; i++) {
          if (!results[i]) {
            const scores = {};
            let total = 0;
            for (const key2 of criteriaKeys) {
              const labels = sortingCriteria[key2]?.labels ?? [];
              const middleIndex = labels.length > 0 ? Math.floor(labels.length / 2) : 5;
              scores[key2] = middleIndex;
              total += middleIndex;
            }
            results[i] = {
              elementIndex: chunk[i].index,
              scores,
              totalScore: total
            };
          }
        }
        return results;
      }
    });
    return extracted;
  };
  const limit = pLimit(10);
  const chunkPromises = chunks.map((chunk) => limit(() => scoreChunk(chunk)));
  const allScores = await Promise.all(chunkPromises);
  const scoreMap = /* @__PURE__ */ new Map();
  for (const chunkScores of allScores) {
    for (const itemScore of chunkScores) {
      const existing = scoreMap.get(itemScore.elementIndex);
      if (existing) {
        for (const key of criteriaKeys) {
          existing.scores[key] = (existing.scores[key] + (itemScore.scores[key] ?? 0)) / 2;
        }
        existing.totalScore = (existing.totalScore + itemScore.totalScore) / 2;
      } else {
        scoreMap.set(itemScore.elementIndex, {
          scores: { ...itemScore.scores },
          totalScore: itemScore.totalScore
        });
      }
    }
  }
  if (scoreMap.size !== input.length) {
    throw new Error(`Score map size mismatch: expected ${input.length}, got ${scoreMap.size}`);
  }
  const scoreGroups = /* @__PURE__ */ new Map();
  for (const [index, scoreData] of scoreMap.entries()) {
    const roundedScore = Math.round(scoreData.totalScore * 100);
    const group = scoreGroups.get(roundedScore) ?? [];
    group.push(index);
    scoreGroups.set(roundedScore, group);
  }
  const tiedGroups = Array.from(scoreGroups.values()).filter((group) => group.length > 1);
  if (tiedGroups.length > 0) {
    const tieBreakLimit = pLimit(10);
    await Promise.all(
      tiedGroups.map(
        (tiedIndices) => tieBreakLimit(async () => {
          if (tiedIndices.length <= 1) return;
          const tiedElements = tiedIndices.map((idx) => elements[idx]);
          const tieBreakText = tiedElements.map((elem, i) => {
            const truncated = tokenizer.truncate(elem.stringified, options.tokensPerItem);
            return `\u25A0${i}: ${truncated}\u25A0`;
          }).join("\n");
          const tieBreakInputStr = JSON.stringify(tiedElements.map((e) => e.element));
          const tieBreakExamples = taskId && ctx.adapter ? await ctx.adapter.getExamples({
            input: tieBreakInputStr.slice(0, 1e3),
            taskType,
            taskId
          }) : [];
          const tieBreakExampleMessages = [];
          for (const example of tieBreakExamples.slice(0, 3)) {
            try {
              const exampleInput = JSON.parse(example.input);
              const exampleItems = Array.isArray(exampleInput) ? exampleInput : [exampleInput];
              tieBreakExampleMessages.push({
                type: "text",
                role: "user",
                content: `Expert Example - Items to order:
${exampleItems.map((el, i) => `\u25A0${i}: ${stringify(el, false).slice(0, 200)}\u25A0`).join("\n")}

Order them from first to last.`
              });
              if (Array.isArray(example.output) && example.output.length > 0) {
                const formattedOrder = example.output.map((_, i) => `\u25A0${i}\u25A0`).join("\n");
                tieBreakExampleMessages.push({
                  type: "text",
                  role: "assistant",
                  content: `${formattedOrder}
${END}`
                });
                if (example.explanation) {
                  tieBreakExampleMessages.push({
                    type: "text",
                    role: "assistant",
                    content: `Reasoning: ${example.explanation}`
                  });
                }
              }
            } catch {
            }
          }
          const { extracted: tieBreakOrder } = await ctx.generateContent({
            systemPrompt: `You are breaking a tie between items with identical total scores.

Instructions: ${instructions}

Criteria:
${criteriaKeys.map((key) => {
              const labels = sortingCriteria[key].labels.join(";");
              return `- ${key}: ${labels}`;
            }).join("\n")}

Order these ${tiedElements.length} items from FIRST to LAST based on the instructions.
Earlier labels in each criterion should come FIRST.

Output format:
\u25A0original_index\u25A0
\u25A0original_index\u25A0
${END}

Output the indices in the order they should appear (first item at top).`,
            stopSequences: [END],
            messages: [
              ...tieBreakExampleMessages,
              {
                type: "text",
                role: "user",
                content: `Items with identical scores (need tie-breaking):
${tieBreakText}

Order them from first to last.
Output format: \u25A0index\u25A0 (one per line)
${END}`
              }
            ],
            transform: (text) => {
              const order = [];
              const regex = /■(\d+)■/g;
              let match;
              while ((match = regex.exec(text)) !== null) {
                const idx = parseInt(match[1] ?? "", 10);
                if (!isNaN(idx) && idx >= 0 && idx < tiedElements.length) {
                  order.push(idx);
                }
              }
              for (let i = 0; i < tiedElements.length; i++) {
                if (!order.includes(i)) {
                  order.push(i);
                }
              }
              return order;
            }
          });
          for (let i = 0; i < tieBreakOrder.length; i++) {
            const elementIndex = tiedElements[tieBreakOrder[i]].index;
            const scoreData = scoreMap.get(elementIndex);
            if (scoreData) {
              scoreData.tieBreakOrder = i;
            }
          }
        })
      )
    );
  }
  const sorted = Array.from(scoreMap.entries()).sort((a, b) => {
    const scoreDiff = a[1].totalScore - b[1].totalScore;
    if (scoreDiff !== 0) return scoreDiff;
    const orderA = a[1].tieBreakOrder ?? 0;
    const orderB = b[1].tieBreakOrder ?? 0;
    return orderA - orderB;
  }).map(([index]) => elements[index].element);
  const result = sorted;
  if (taskId && ctx.adapter && !ctx.controller.signal.aborted) {
    const key = fastHash(
      stringify({
        taskId,
        taskType,
        input: JSON.stringify(input),
        instructions
      })
    );
    await ctx.adapter.saveExample({
      key,
      taskType,
      taskId,
      input: JSON.stringify(input),
      output: result,
      instructions,
      metadata: {
        cost: { input: 0, output: 0 },
        latency: 0,
        model: ctx.modelId,
        tokens: { input: 0, output: 0 }
      }
    });
  }
  return result;
};
Zai.prototype.sort = function(input, instructions, _options) {
  const context = new ZaiContext({
    client: this.client,
    modelId: this.Model,
    taskId: this.taskId,
    taskType: "zai.sort",
    adapter: this.adapter
  });
  return new Response(
    context,
    sort(input, instructions, _options, context),
    (result) => result
    // Simplified form is just the sorted array
  );
};
