import { z } from "@bpinternal/zui";
import { clamp } from "lodash-es";
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 _InitialGroup = z.object({
  id: z.string().min(1).max(100),
  label: z.string().min(1).max(250),
  elements: z.array(z.any()).optional().default([])
});
const _Options = z.object({
  instructions: z.string().optional(),
  tokensPerElement: z.number().min(1).max(1e5).optional().default(250),
  chunkLength: z.number().min(100).max(1e5).optional().default(16e3),
  initialGroups: z.array(_InitialGroup).optional().default([])
});
const END = "\u25A0END\u25A0";
const normalizeLabel = (label) => {
  return label.trim().toLowerCase().replace(/^(group|new group|new)\s*[-:]\s*/i, "").replace(/^(group|new group|new)\s+/i, "").trim();
};
const group = async (input, _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.group";
  if (input.length === 0) {
    return [];
  }
  const groups = /* @__PURE__ */ new Map();
  const groupElements = /* @__PURE__ */ new Map();
  const elementGroups = /* @__PURE__ */ new Map();
  const labelToGroupId = /* @__PURE__ */ new Map();
  let groupIdCounter = 0;
  options.initialGroups.forEach((ig) => {
    const normalized = normalizeLabel(ig.label);
    groups.set(ig.id, { id: ig.id, label: ig.label, normalizedLabel: normalized });
    groupElements.set(ig.id, /* @__PURE__ */ new Set());
    labelToGroupId.set(normalized, ig.id);
  });
  const elements = input.map((element, idx) => ({
    element,
    index: idx,
    stringified: stringify(element, false)
  }));
  const TOKENS_TOTAL_MAX = model.input.maxTokens - PROMPT_INPUT_BUFFER - PROMPT_OUTPUT_BUFFER;
  const TOKENS_INSTRUCTIONS_MAX = options.instructions ? clamp(tokenizer.count(options.instructions), 100, TOKENS_TOTAL_MAX * 0.2) : 0;
  const TOKENS_AVAILABLE = TOKENS_TOTAL_MAX - TOKENS_INSTRUCTIONS_MAX;
  const TOKENS_FOR_GROUPS_MAX = Math.floor(TOKENS_AVAILABLE * 0.4);
  const TOKENS_FOR_ELEMENTS_MAX = Math.floor(TOKENS_AVAILABLE * 0.6);
  const MAX_ELEMENTS_PER_CHUNK = 50;
  const elementChunks = [];
  let currentChunk = [];
  let currentTokens = 0;
  for (const elem of elements) {
    const truncated = tokenizer.truncate(elem.stringified, options.tokensPerElement);
    const elemTokens = tokenizer.count(truncated);
    if ((currentTokens + elemTokens > TOKENS_FOR_ELEMENTS_MAX || currentChunk.length >= MAX_ELEMENTS_PER_CHUNK) && currentChunk.length > 0) {
      elementChunks.push(currentChunk);
      currentChunk = [];
      currentTokens = 0;
    }
    currentChunk.push(elem.index);
    currentTokens += elemTokens;
  }
  if (currentChunk.length > 0) {
    elementChunks.push(currentChunk);
  }
  const getGroupChunks = () => {
    const allGroupIds2 = Array.from(groups.keys());
    if (allGroupIds2.length === 0) return [[]];
    const chunks = [];
    let currentChunk2 = [];
    let currentTokens2 = 0;
    for (const groupId of allGroupIds2) {
      const group2 = groups.get(groupId);
      const groupTokens = tokenizer.count(`${group2.label}`) + 10;
      if (currentTokens2 + groupTokens > TOKENS_FOR_GROUPS_MAX && currentChunk2.length > 0) {
        chunks.push(currentChunk2);
        currentChunk2 = [];
        currentTokens2 = 0;
      }
      currentChunk2.push(groupId);
      currentTokens2 += groupTokens;
    }
    if (currentChunk2.length > 0) {
      chunks.push(currentChunk2);
    }
    return chunks.length > 0 ? chunks : [[]];
  };
  const processChunk = async (elementIndices, groupIds) => {
    const chunkElements = elementIndices.map((idx) => elements[idx].element);
    const chunkInputStr = JSON.stringify(chunkElements);
    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: options.instructions ?? "",
        groupIds: groupIds.join(",")
      })
    );
    const exactMatch = examples.find((x) => x.key === key);
    if (exactMatch && exactMatch.output) {
      return exactMatch.output;
    }
    const elementsText = elementIndices.map((idx, i) => {
      const elem = elements[idx];
      const truncated = tokenizer.truncate(elem.stringified, options.tokensPerElement);
      return `\u25A0${i}: ${truncated}\u25A0`;
    }).join("\n");
    const groupsList = groupIds.map((gid) => groups.get(gid).label);
    const groupsText = groupsList.length > 0 ? `**Existing Groups (prefer reusing these):**
${groupsList.map((l) => `- ${l}`).join("\n")}

` : "";
    const exampleMessages = [];
    for (const example of examples.slice(0, 5)) {
      try {
        const exampleInput = JSON.parse(example.input);
        const exampleElements = Array.isArray(exampleInput) ? exampleInput : [exampleInput];
        const exampleElementsText = exampleElements.map((el, i) => `\u25A0${i}: ${stringify(el, false).slice(0, 200)}\u25A0`).join("\n");
        exampleMessages.push({
          type: "text",
          role: "user",
          content: `Expert Example - Elements to group:
${exampleElementsText}

Group each element.`
        });
        const exampleOutput = example.output;
        if (Array.isArray(exampleOutput) && exampleOutput.length > 0) {
          const formattedAssignments = exampleOutput.map((assignment) => `\u25A0${assignment.elementIndex}:${assignment.label}\u25A0`).join("\n");
          exampleMessages.push({
            type: "text",
            role: "assistant",
            content: `${formattedAssignments}
${END}`
          });
          if (example.explanation) {
            exampleMessages.push({
              type: "text",
              role: "assistant",
              content: `Reasoning: ${example.explanation}`
            });
          }
        }
      } catch {
      }
    }
    const systemPrompt = `You are grouping elements into cohesive groups.

${options.instructions ? `**Instructions:** ${options.instructions}
` : "**Instructions:** Group similar elements together."}

**Important:**
- Each element gets exactly ONE group label
- Use EXACT SAME label for similar items (case-sensitive)
- Create new descriptive labels when needed

**Output Format:**
One line per element:
\u25A00:Group Label\u25A0
\u25A01:Group Label\u25A0
${END}`.trim();
    const userPrompt = `${groupsText}**Elements (\u25A00 to \u25A0${elementIndices.length - 1}):**
${elementsText}

**Task:** For each element, output one line with its group label.
${END}`.trim();
    const { extracted } = await ctx.generateContent({
      systemPrompt,
      stopSequences: [END],
      messages: [...exampleMessages, { type: "text", role: "user", content: userPrompt }],
      transform: (text) => {
        const assignments = [];
        const regex = /■(\d+):([^■]+)■/g;
        let match;
        while ((match = regex.exec(text)) !== null) {
          const idx = parseInt(match[1] ?? "", 10);
          if (isNaN(idx) || idx < 0 || idx >= elementIndices.length) continue;
          const label = (match[2] ?? "").trim();
          if (!label) continue;
          assignments.push({
            elementIndex: elementIndices[idx],
            label: label.slice(0, 250)
          });
        }
        return assignments;
      }
    });
    return extracted;
  };
  const elementLimit = pLimit(10);
  const groupLimit = pLimit(10);
  const allChunkResults = await Promise.all(
    elementChunks.map(
      (elementChunk) => elementLimit(async () => {
        const groupChunks = getGroupChunks();
        const allAssignments = await Promise.all(
          groupChunks.map((groupChunk) => groupLimit(() => processChunk(elementChunk, groupChunk)))
        );
        return allAssignments.flat();
      })
    )
  );
  for (const assignments of allChunkResults) {
    for (const { elementIndex, label } of assignments) {
      const normalized = normalizeLabel(label);
      let groupId = labelToGroupId.get(normalized);
      if (!groupId) {
        groupId = `group_${groupIdCounter++}`;
        groups.set(groupId, { id: groupId, label, normalizedLabel: normalized });
        groupElements.set(groupId, /* @__PURE__ */ new Set());
        labelToGroupId.set(normalized, groupId);
      }
      groupElements.get(groupId).add(elementIndex);
      if (!elementGroups.has(elementIndex)) {
        elementGroups.set(elementIndex, /* @__PURE__ */ new Set());
      }
      elementGroups.get(elementIndex).add(groupId);
    }
  }
  const allGroupIds = Array.from(groups.keys());
  if (allGroupIds.length > 0) {
    const elementsNeedingReview = [];
    for (const elem of elements) {
      const seenGroups = elementGroups.get(elem.index) ?? /* @__PURE__ */ new Set();
      const unseenCount = allGroupIds.filter((gid) => !seenGroups.has(gid)).length;
      if (unseenCount > 0) {
        elementsNeedingReview.push(elem.index);
      }
    }
    if (elementsNeedingReview.length > 0) {
      const reviewChunks = [];
      let reviewChunk = [];
      let reviewTokens = 0;
      for (const elemIdx of elementsNeedingReview) {
        const elem = elements[elemIdx];
        const truncated = tokenizer.truncate(elem.stringified, options.tokensPerElement);
        const elemTokens = tokenizer.count(truncated);
        const shouldStartNewChunk = (reviewTokens + elemTokens > TOKENS_FOR_ELEMENTS_MAX || reviewChunk.length >= MAX_ELEMENTS_PER_CHUNK) && reviewChunk.length > 0;
        if (shouldStartNewChunk) {
          reviewChunks.push(reviewChunk);
          reviewChunk = [];
          reviewTokens = 0;
        }
        reviewChunk.push(elemIdx);
        reviewTokens += elemTokens;
      }
      if (reviewChunk.length > 0) {
        reviewChunks.push(reviewChunk);
      }
      const reviewResults = await Promise.all(
        reviewChunks.map(
          (chunk) => elementLimit(async () => {
            const groupChunks = getGroupChunks();
            const allAssignments = await Promise.all(
              groupChunks.map((groupChunk) => groupLimit(() => processChunk(chunk, groupChunk)))
            );
            return allAssignments.flat();
          })
        )
      );
      const updateElementGroupAssignment = (elementIndex, label) => {
        const normalized = normalizeLabel(label);
        const groupId = labelToGroupId.get(normalized);
        if (!groupId) return;
        groupElements.get(groupId).add(elementIndex);
        const elemGroups = elementGroups.get(elementIndex) ?? /* @__PURE__ */ new Set();
        if (!elementGroups.has(elementIndex)) {
          elementGroups.set(elementIndex, elemGroups);
        }
        elemGroups.add(groupId);
      };
      for (const assignments of reviewResults) {
        for (const { elementIndex, label } of assignments) {
          updateElementGroupAssignment(elementIndex, label);
        }
      }
    }
  }
  for (const [elementIndex, groupSet] of elementGroups.entries()) {
    if (groupSet.size > 1) {
      const groupIds = Array.from(groupSet);
      for (const gid of groupIds) {
        groupElements.get(gid)?.delete(elementIndex);
      }
      const finalGroupId = groupIds[0];
      groupElements.get(finalGroupId).add(elementIndex);
    }
  }
  const result = [];
  for (const [groupId, elementIndices] of groupElements.entries()) {
    if (elementIndices.size > 0) {
      const groupInfo = groups.get(groupId);
      result.push({
        id: groupInfo.id,
        label: groupInfo.label,
        elements: Array.from(elementIndices).map((idx) => elements[idx].element)
      });
    }
  }
  if (taskId && ctx.adapter && !ctx.controller.signal.aborted) {
    const key = fastHash(
      stringify({
        taskId,
        taskType,
        input: JSON.stringify(input),
        instructions: options.instructions ?? ""
      })
    );
    const outputAssignments = [];
    for (const [groupId, elementIndices] of groupElements.entries()) {
      const groupInfo = groups.get(groupId);
      for (const idx of elementIndices) {
        outputAssignments.push({
          elementIndex: idx,
          label: groupInfo.label
        });
      }
    }
    await ctx.adapter.saveExample({
      key,
      taskType,
      taskId,
      input: JSON.stringify(input),
      output: result,
      instructions: options.instructions ?? "",
      metadata: {
        cost: { input: 0, output: 0 },
        latency: 0,
        model: ctx.modelId,
        tokens: { input: 0, output: 0 }
      }
    });
  }
  return result;
};
Zai.prototype.group = function(input, _options) {
  const context = new ZaiContext({
    client: this.client,
    modelId: this.Model,
    taskId: this.taskId,
    taskType: "zai.group",
    adapter: this.adapter
  });
  return new Response(context, group(input, _options, context), (result) => {
    const merged = {};
    result.forEach((group2) => {
      if (!merged[group2.label]) {
        merged[group2.label] = [];
      }
      merged[group2.label].push(...group2.elements);
    });
    return merged;
  });
};
