ai-chatbot/lib/ai/context-memory/media-input.ts

import "server-only";

import { extractImageUrlsFromText, inferOriginalImageUrl } from "@/lib/image-url";
import type { DBMessage } from "@/lib/db/schema";
import type {
  CanvasImageNodeContext,
  ResolveChatMediaInputParams,
  ResolvedChatMediaInput,
  ResolvedImageRef,
} from "@/lib/ai/context-memory/types";

function normalizeMediaUrl(url: string): string {
  const raw = String(url || "").trim();
  if (!raw) return "";

  if (/^data:image\/[a-z0-9+.-]+;base64,/i.test(raw)) {
    return raw.replace(/\s+/g, "");
  }

  try {
    const parsed = new URL(raw);
    return `${parsed.origin}${parsed.pathname}`.replace(/\/+$/, "").toLowerCase();
  } catch {
    return raw
      .split("#")[0]
      .split("?")[0]
      .replace(/\/+$/, "")
      .toLowerCase();
  }
}

function normalizeImageInputRef(rawValue: unknown, requestOrigin?: string): string | null {
  if (typeof rawValue !== "string") return null;
  const trimmed = rawValue.trim();
  if (!trimmed) return null;

  if (/^data:image\/[a-z0-9+.-]+;base64,/i.test(trimmed)) {
    return trimmed.replace(/\s+/g, "");
  }

  if (/^https?:\/\//i.test(trimmed)) {
    return inferOriginalImageUrl(trimmed);
  }

  if (requestOrigin && trimmed.startsWith("/")) {
    try {
      return new URL(trimmed, requestOrigin).toString();
    } catch {
      return null;
    }
  }

  return null;
}

function collectImageRefsFromMessageParts(
  parts: Array<Record<string, any>> | undefined,
  requestOrigin?: string,
): Array<{ url: string; assetId?: string }> {
  if (!Array.isArray(parts) || parts.length === 0) return [];

  const result: Array<{ url: string; assetId?: string }> = [];
  const seen = new Set<string>();
  const pushInput = (rawValue: unknown, assetId?: unknown) => {
    const normalized = normalizeImageInputRef(rawValue, requestOrigin);
    if (!normalized) return;
    const key = normalizeMediaUrl(normalized);
    if (!key || seen.has(key)) return;
    seen.add(key);
    result.push({
      url: normalized,
      assetId: typeof assetId === "string" && assetId.trim() ? assetId.trim() : undefined,
    });
  };

  for (const part of parts) {
    if (!part) continue;

    if (part.type === "file") {
      pushInput(part.url, part.assetId);
      pushInput(part.data, part.assetId);
      continue;
    }

    if (part.type === "image") {
      pushInput(part.url, part.assetId);
      pushInput(part.image, part.assetId);
      pushInput(part.data, part.assetId);
      continue;
    }

    if (part.type === "image_url") {
      pushInput(part.image_url?.url, part.assetId);
      pushInput(part.url, part.assetId);
      continue;
    }

    if (part.type === "text" && typeof part.text === "string") {
      for (const parsedUrl of extractImageUrlsFromText(part.text)) {
        pushInput(parsedUrl);
      }
    }
  }

  return result;
}

function dedupeRefs(refs: ResolvedImageRef[], maxImages: number): ResolvedImageRef[] {
  const seen = new Set<string>();
  const result: ResolvedImageRef[] = [];

  for (const ref of refs) {
    const key = ref.normalizedUrl || normalizeMediaUrl(ref.url);
    if (!key || seen.has(key)) continue;
    seen.add(key);
    result.push({ ...ref, normalizedUrl: key });
    if (result.length >= maxImages) break;
  }

  return result;
}

function toCanvasImageRefs(nodes: CanvasImageNodeContext[] | undefined, maxImages: number) {
  if (!Array.isArray(nodes) || nodes.length === 0) return [] as ResolvedImageRef[];

  return dedupeRefs(
    nodes
      .filter((node) => typeof node?.url === "string" && node.url.trim())
      .map((node) => ({
        url: inferOriginalImageUrl(node.url),
        normalizedUrl: normalizeMediaUrl(inferOriginalImageUrl(node.url)),
        assetId: typeof node.assetId === "string" ? node.assetId.trim() || undefined : undefined,
        nodeId: node.nodeId,
        title: node.title,
        prompt: node.prompt,
        x: node.x,
        y: node.y,
        width: node.width,
        height: node.height,
        isSelected: Boolean(node.isSelected),
        source: node.isSelected ? "canvas_selection" : "canvas_context",
        role: node.isSelected ? "selected" : "reference",
      })),
    maxImages,
  );
}

function extractHistoryImageRefs(params: {
  historyMessages: DBMessage[];
  requestOrigin?: string;
  maxImages: number;
}) {
  const latestAssistant: ResolvedImageRef[] = [];
  const latestUser: ResolvedImageRef[] = [];
  const seenAssistant = new Set<string>();
  const seenUser = new Set<string>();

  const pushRef = (
    bucket: ResolvedImageRef[],
    seen: Set<string>,
    url: string,
    options: Pick<ResolvedImageRef, "source" | "role" | "messageId">,
  ) => {
    const normalizedUrl = normalizeMediaUrl(url);
    if (!normalizedUrl || seen.has(normalizedUrl) || bucket.length >= params.maxImages) return;
    seen.add(normalizedUrl);
    bucket.push({
      url,
      normalizedUrl,
      source: options.source,
      role: options.role,
      messageId: options.messageId,
    });
  };

  for (let index = params.historyMessages.length - 1; index >= 0; index -= 1) {
    const historyMessage = params.historyMessages[index];
    const parts = historyMessage.parts as Array<Record<string, any>> | undefined;
    if (!Array.isArray(parts) || parts.length === 0) continue;

    const refs = collectImageRefsFromMessageParts(parts, params.requestOrigin);
    if (!refs.length) continue;

    if (historyMessage.role === "assistant" && latestAssistant.length < params.maxImages) {
      for (const ref of refs) {
        pushRef(latestAssistant, seenAssistant, ref.url, {
          source: "history_assistant_result",
          role: "history_result",
          messageId: historyMessage.id,
        });
      }
    }

    if (historyMessage.role === "user" && latestUser.length < params.maxImages) {
      for (const ref of refs) {
        pushRef(latestUser, seenUser, ref.url, {
          source: "history_user_upload",
          role: "reference",
          messageId: historyMessage.id,
        });
      }
    }

    if (latestAssistant.length >= params.maxImages && latestUser.length >= params.maxImages) {
      break;
    }
  }

  return {
    latestAssistant,
    latestUser,
  };
}

export function resolveChatMediaInput(
  params: ResolveChatMediaInputParams,
): ResolvedChatMediaInput {
  const maxImages = Math.max(1, Math.min(12, Math.round(params.maxImages || 4)));
  const currentMessageImageRefs = dedupeRefs(
    collectImageRefsFromMessageParts(params.messageParts, params.requestOrigin).map((ref) => ({
      url: ref.url,
      normalizedUrl: normalizeMediaUrl(ref.url),
      assetId: ref.assetId,
      source: "current_message",
      role: "upload",
    })),
    maxImages,
  );

  const canvasImageRefs = toCanvasImageRefs(params.contextPayload?.canvasImageNodes, maxImages);
  const selectedNodeIds = new Set(params.contextPayload?.selectedNodeIds || []);
  const selectedImageRefs = dedupeRefs(
    canvasImageRefs.filter((ref) => ref.isSelected || (ref.nodeId && selectedNodeIds.has(ref.nodeId))),
    maxImages,
  );
  const history = extractHistoryImageRefs({
    historyMessages: params.historyMessages,
    requestOrigin: params.requestOrigin,
    maxImages,
  });

  const resolvedImageRefs =
    currentMessageImageRefs.length > 0
      ? currentMessageImageRefs
      : selectedImageRefs.length > 0
        ? selectedImageRefs
        : history.latestAssistant.length > 0
          ? history.latestAssistant
          : history.latestUser;

  const imageRefs = dedupeRefs(
    [
      ...currentMessageImageRefs,
      ...selectedImageRefs,
      ...canvasImageRefs,
      ...history.latestAssistant,
      ...history.latestUser,
    ],
    maxImages * 4,
  );

  const selectedAssetIds = Array.from(
    new Set(
      selectedImageRefs
        .map((ref) => ref.assetId)
        .filter((value): value is string => typeof value === "string" && value.trim().length > 0),
    ),
  );
  const uploadedAssetIds = Array.from(
    new Set(
      currentMessageImageRefs
        .map((ref) => ref.assetId)
        .filter((value): value is string => typeof value === "string" && value.trim().length > 0),
    ),
  );
  const candidateAssetIds = Array.from(
    new Set(
      imageRefs
        .map((ref) => ref.assetId)
        .filter((value): value is string => typeof value === "string" && value.trim().length > 0),
    ),
  );

  const source =
    currentMessageImageRefs.length > 0
      ? "current_message"
      : selectedImageRefs.length > 0
        ? "canvas_selection"
        : resolvedImageRefs.length > 0
          ? "history"
          : "none";

  return {
    imageRefs,
    currentMessageImageRefs,
    selectedImageRefs,
    canvasImageRefs,
    historyUserImageRefs: history.latestUser,
    historyAssistantImageRefs: history.latestAssistant,
    resolvedImageRefs,
    currentMessageImageUrls: currentMessageImageRefs.map((ref) => ref.url),
    selectedImageUrls: selectedImageRefs.map((ref) => ref.url),
    historyUserImageUrls: history.latestUser.map((ref) => ref.url),
    historyAssistantImageUrls: history.latestAssistant.map((ref) => ref.url),
    resolvedImageUrls: resolvedImageRefs.map((ref) => ref.url),
    selectedAssetIds,
    candidateAssetIds,
    uploadedAssetIds,
    hasImage: resolvedImageRefs.length > 0,
    source,
    fromHistory: source === "history" || source === "canvas_selection",
  };
}