ai-chatbot/app/(chat)/api/chat/_lib/message-preparation.ts

import { createLogger } from "@/lib/observability/logger";
import {
  analyzeOptimizedPromptMerge,
  buildOptimizedPromptPreview,
} from "@/lib/ai/runtime/optimized-prompt";
import type { ChatMessage } from "@/lib/types";

const chatApiLogger = createLogger("chat.messagePreparation");


/**
 * 将图片 URL 转换为 base64 格式
 * 用于将 S3 URL 的图片转换为 base64 发送给 Gemini API
 */
export async function convertImageUrlToBase64(url: string): Promise<string> {
  try {
    chatApiLogger.info("[Image Convert] Fetching image from:", url);
    const response = await fetch(url);
    if (!response.ok) {
      throw new Error(`Failed to fetch image: ${response.statusText}`);
    }

    const arrayBuffer = await response.arrayBuffer();
    const buffer = Buffer.from(arrayBuffer);
    const base64 = buffer.toString("base64");

    // 获取 MIME 类型
    const contentType = response.headers.get("content-type") || "image/jpeg";
    const dataUrl = `data:${contentType};base64,${base64}`;

    chatApiLogger.info("[Image Convert] Converted to base64, size:", base64.length);
    return dataUrl;
  } catch (error) {
    chatApiLogger.error("[Image Convert] Failed:", error);
    throw error;
  }
}

/**
 * 检测消息中是否包含用户标注过的图片（文件名以 annotated- 开头）
 */
export function hasAnnotatedImage(messages: ChatMessage[]): boolean {
  for (const message of messages) {
    for (const part of message.parts) {
      if (part.type === "file" && "name" in part && typeof part.name === "string") {
        if (part.name.startsWith("annotated-")) {
          return true;
        }
      }
    }
  }
  return false;
}

/**
 * 标注图片的系统提示词（用户不可见，后台自动添加）
 */
export const ANNOTATION_SYSTEM_PROMPT =
  "【重要】用户提供的图片上有标注线条（如箭头、圆圈、画笔痕迹或文字说明），这些标注是用户用来指示需要修改的区域或方向。请仔细理解标注的含义，按照标注指示进行编辑，但在生成的新图片中不要保留这些标注线条，输出干净的编辑结果。";

export function applyHiddenPromptToMessages(
  messages: ChatMessage[],
  hiddenPrompt?: string,
): ChatMessage[] {
  if (!hiddenPrompt || !hiddenPrompt.trim()) return messages;
  const cloned = messages.map((msg) => ({
    ...msg,
    parts: msg.parts.map((part) => ({ ...part })),
  }));
  for (let i = cloned.length - 1; i >= 0; i--) {
    const msg = cloned[i];
    if (msg.role !== "user") continue;
    const textPart = msg.parts.find((part) => part.type === "text") as
      | { type: "text"; text: string }
      | undefined;
    if (textPart) {
      textPart.text = `${hiddenPrompt}\n\n${textPart.text}`.trim();
    } else {
      msg.parts.unshift({ type: "text", text: hiddenPrompt });
    }
    break;
  }
  return cloned;
}

export function applyOptimizedPromptToLatestUserMessage(
  messages: ChatMessage[],
  optimizedPrompt?: string,
  hiddenPromptPrefix = "",
): ChatMessage[] {
  const prompt = optimizedPrompt?.trim();
  if (!prompt) return messages;

  const cloned = messages.map((msg) => ({
    ...msg,
    parts: msg.parts.map((part) => ({ ...part })),
  }));

  for (let i = cloned.length - 1; i >= 0; i -= 1) {
    const msg = cloned[i];
    if (msg.role !== "user") continue;

    const textPart = msg.parts.find(
      (part): part is { type: "text"; text: string } => part.type === "text",
    );
    const originalText = textPart?.text?.trim() || "";
    const { normalizedOriginal, hasOriginalDemand, hasOriginalInOptimized } =
      analyzeOptimizedPromptMerge(originalText, prompt, hiddenPromptPrefix);
    const finalPrompt =
      hasOriginalDemand && !hasOriginalInOptimized
        ? [
            hiddenPromptPrefix.trim(),
            `【核心目标（必须优先满足）】${normalizedOriginal}`,
            `【执行优化细节】${prompt}`,
          ]
            .filter(Boolean)
            .join("\n\n")
            .trim()
        : `${hiddenPromptPrefix}${prompt}`.trim();
    chatApiLogger.info("[Prompt Merge] Final image prompt prepared:", {
      original: normalizedOriginal.substring(0, 120),
      optimized: prompt.substring(0, 120),
      merged: finalPrompt.substring(0, 180),
      forcePreserveOriginal: hasOriginalDemand && !hasOriginalInOptimized,
    });

    if (textPart) {
      textPart.text = finalPrompt;
    } else {
      msg.parts.unshift({ type: "text", text: finalPrompt });
    }
    break;
  }

  return cloned;
}

export function injectOptimizedPromptSection(params: {
  message?: string;
  originalText: string;
  optimizedPrompt?: string;
  hiddenPromptPrefix?: string;
}): string | undefined {
  const { message, originalText, optimizedPrompt, hiddenPromptPrefix } = params;
  if (!message?.trim()) return message;

  const preview = buildOptimizedPromptPreview({
    originalText,
    optimizedPrompt,
    hiddenPromptPrefix,
  });
  if (!preview) return message;

  const optimizedPromptSection = `【优化后提示词】\n${preview}`;
  const optimizedSectionRegex =
    /(?:###\s*优化后提示词|【\s*优化后提示词\s*】)[\s\S]*?(?=\n(?:###\s*|【\s*[^】]+\s*】)|$)/;
  if (optimizedSectionRegex.test(message)) {
    return message.replace(optimizedSectionRegex, optimizedPromptSection);
  }

  const executionSectionRegex =
    /(?:###\s*执行计划|【\s*执行计划\s*】)[\s\S]*?(?=\n(?:###\s*|【\s*[^】]+\s*】)|$)/;
  if (executionSectionRegex.test(message)) {
    return message.replace(executionSectionRegex, optimizedPromptSection);
  }

  const outputSectionRegex = /(?:###\s*输出设置|【\s*输出设置\s*】)/;
  if (outputSectionRegex.test(message)) {
    return message.replace(outputSectionRegex, `${optimizedPromptSection}\n$&`);
  }

  return `${message.trim()}\n\n${optimizedPromptSection}`;
}