Agent/frontend/react-template/src/utils/imageExtraction.ts

import type { Message } from "../types/message";

export interface ExtractedImage {
  url: string;
  alt?: string;
}

/**
 * Extracts images from a message's content or result.
 * Handles both JSON array format (MsgResult[]) and Rich Text (Markdown/HTML).
 *
 * @param result - The content or result field from a message
 * @returns Array of extracted images
 */
export const extractImagesFromResult = (result: unknown): ExtractedImage[] => {
  const images: ExtractedImage[] = [];

  if (!result) return images;

  // Case 0: result IS the message content which might be an array directly
  if (Array.isArray(result)) {
    result.forEach((item) => {
      if (typeof item === "object" && item !== null) {
        // eslint-disable-next-line @typescript-eslint/no-explicit-any
        const msgItem = item as any;

        // 1. Check for standard OpenAI-like image_url
        if (msgItem.image_url && typeof msgItem.image_url === "object" && msgItem.image_url.url) {
          images.push({
            url: msgItem.image_url.url,
            alt: "Attached Image",
          });
        }

        // 2. Check for type="image" with source
        if (msgItem.type === "image" && msgItem.source && typeof msgItem.source === "object") {
          const source = msgItem.source;
          if (source.data) {
            const mimeType = source.media_type || "image/png";
            images.push({
              url: `data:${mimeType};base64,${source.data}`,
              alt: "Base64 Image",
            });
          } else if (source.url) {
            images.push({
              url: source.url,
              alt: "Image URL",
            });
          }
        }
      }
    });
  }

  // Case 2: result is a string (Rich Text / Markdown)
  if (typeof result === "string") {
    // 1. Match Markdown images: ![alt](url)
    const markdownRegex = /!\[([^\]]*)\]\(([^)]+)\)/g;
    let match;
    while ((match = markdownRegex.exec(result)) !== null) {
      images.push({
        alt: match[1] || "Markdown Image",
        url: match[2],
      });
    }

    // 2. Match HTML img tags: <img ... src="..." ...>
    const htmlRegex = /<img\s+[^>]*src=["']([^"']+)["'][^>]*>/g;
    while ((match = htmlRegex.exec(result)) !== null) {
      images.push({
        alt: "HTML Image",
        url: match[1],
      });
    }

    // 3. Match JSON "image_url": "..." patterns embedded in text
    // Matches "image_url"\s*:\s*"([^"]+)"
    const jsonRegex = /"image_url"\s*:\s*"([^"]+)"/g;
    while ((match = jsonRegex.exec(result)) !== null) {
      // Basic filtering to avoid matching non-URL strings if the key is reused
      if (match[1].startsWith("http") || match[1].startsWith("data:")) {
        images.push({
          alt: "JSON Image",
          url: match[1],
        });
      }
    }
  }
  return images;
};

/**
 * Helper to extract images from a Message object
 */
export const extractImagesFromMessage = (message: Message): ExtractedImage[] => {
  if (!message.content) return [];

  // If content is a string, treat it as result
  if (typeof message.content === "string") {
    return extractImagesFromResult(message.content);
  }

  // If content is an object (MessageContent)
  if (typeof message.content === "object") {
    // Check 'result' field
    if ("result" in message.content && message.content.result) {
      return extractImagesFromResult(message.content.result);
    }
    // Also check if content itself is an array (e.g. standard MessageContent array)
    if (Array.isArray(message.content)) {
      return extractImagesFromResult(message.content);
    }
  }

  return [];
};