ai-chatbot/lib/video-canvas/prompt-structure.ts

export type VideoPromptMention = {
  id: string;
  index: number;
  label: string;
  mediaType: "audio" | "image" | "video";
  offset: number;
  sourceNodeId: string;
  url: string;
};

export type VideoPromptToolMark = {
  id: string;
  label: string;
};

export type VideoPromptToolReference = Pick<
  VideoPromptMention,
  "id" | "index" | "label" | "mediaType" | "offset"
>;

const STRUCTURE_TOKEN_PATTERN = /\[\[(REF|MARK):([A-Za-z0-9_-]{1,100})\]\]/g;

function normalizeTokenId(value: unknown) {
  return String(value || "")
    .trim()
    .replace(/[^A-Za-z0-9_-]+/g, "-")
    .slice(0, 100);
}

export function normalizeVideoPromptMentions(value: unknown): VideoPromptMention[] {
  if (!Array.isArray(value)) return [];

  return value.slice(0, 20).flatMap((item, position) => {
    if (!item || typeof item !== "object") return [];
    const record = item as Record<string, unknown>;
    const sourceNodeId = String(record.sourceNodeId || "").trim().slice(0, 120);
    if (!sourceNodeId) return [];
    const index = Math.max(1, Math.min(10, Math.round(Number(record.index) || position + 1)));
    const mediaType = ["audio", "image", "video"].includes(String(record.mediaType))
      ? (String(record.mediaType) as VideoPromptMention["mediaType"])
      : "image";
    const id = normalizeTokenId(record.id) || `mention-${sourceNodeId}-${position + 1}`;

    return [{
      id,
      index,
      label: String(record.label || `图片 ${index}`).trim().slice(0, 100) || `图片 ${index}`,
      mediaType,
      offset: Math.max(0, Math.round(Number(record.offset) || 0)),
      sourceNodeId,
      url: String(record.url || "").trim().slice(0, 2048),
    }];
  });
}

function referenceToken(id: string) {
  return `[[REF:${normalizeTokenId(id)}]]`;
}

function markToken(id: string) {
  return `[[MARK:${normalizeTokenId(id)}]]`;
}

export function injectVideoPromptStructure(
  prompt: string,
  references: VideoPromptToolReference[],
  marks: VideoPromptToolMark[],
) {
  const normalizedPrompt = String(prompt || "");
  const insertions = references
    .map((reference, order) => ({
      order,
      offset: Math.max(0, Math.min(normalizedPrompt.length, Math.round(reference.offset || 0))),
      token: referenceToken(reference.id),
    }))
    .sort((left, right) => left.offset - right.offset || left.order - right.order);
  let cursor = 0;
  let result = "";

  for (const insertion of insertions) {
    result += normalizedPrompt.slice(cursor, insertion.offset);
    result += insertion.token;
    cursor = insertion.offset;
  }
  result += normalizedPrompt.slice(cursor);

  const markTokens = marks
    .map((mark) => markToken(mark.id))
    .filter((token) => token !== "[[MARK:]]");
  return markTokens.length > 0 ? `${result} ${markTokens.join(" ")}`.trim() : result;
}

export function buildVideoPromptStructureInstructions(
  references: VideoPromptToolReference[],
  marks: VideoPromptToolMark[],
) {
  const lines = [
    ...references.map(
      (reference) =>
        `${referenceToken(reference.id)} = @${reference.index}「${reference.label}」${reference.mediaType}素材`,
    ),
    ...marks.map((mark) => `${markToken(mark.id)} = 画面元素标记「${mark.label}」`),
  ];
  if (lines.length === 0) return "";

  return [
    "结构化占位符协议：下列占位符代表编辑器中的不可编辑芯片。理解其语义，并在输出中逐字保留每个占位符一次；不要翻译、改写、删除或重复占位符。",
    ...lines,
  ].join("\n");
}

export function resolveVideoPromptStructure(
  generatedText: string,
  references: VideoPromptToolReference[],
  marks: VideoPromptToolMark[],
  options: { flatten: boolean },
) {
  const referenceById = new Map(references.map((item) => [normalizeTokenId(item.id), item]));
  const markById = new Map(marks.map((item) => [normalizeTokenId(item.id), item]));
  const mentionOffsets: Record<string, number> = {};
  const preservedMarkIds: string[] = [];
  let cursor = 0;
  let text = "";

  for (const match of String(generatedText || "").matchAll(STRUCTURE_TOKEN_PATTERN)) {
    const start = match.index ?? 0;
    text += String(generatedText || "").slice(cursor, start);
    const kind = match[1];
    const id = normalizeTokenId(match[2]);
    if (kind === "REF") {
      const reference = referenceById.get(id);
      if (reference) {
        if (options.flatten) text += reference.label;
        else mentionOffsets[reference.id] = text.length;
      }
    } else {
      const mark = markById.get(id);
      if (mark) {
        if (options.flatten) text += mark.label;
        else preservedMarkIds.push(mark.id);
      }
    }
    cursor = start + match[0].length;
  }
  text += String(generatedText || "").slice(cursor);

  const leadingWhitespace = text.length - text.trimStart().length;
  if (leadingWhitespace > 0) {
    for (const key of Object.keys(mentionOffsets)) {
      mentionOffsets[key] = Math.max(0, mentionOffsets[key] - leadingWhitespace);
    }
  }

  return {
    mentionOffsets,
    preservedMarkIds,
    text: text.trim(),
  };
}