ai-chatbot/lib/canvas/agent-direct-writeback.ts

import {
  getBestAnchorNode,
  getSelectedMediaNodes,
  getSelectedSlots,
} from "@/lib/canvas/agent-context";
import { hasCanvasAgentAnnotationEditIntent } from "@/lib/canvas/agent-annotation-prompt";
import type {
  CanvasAgentAnnotationRef,
  CanvasAgentContext,
  CanvasAgentNodeRef,
  CanvasAgentSlotRef,
  CanvasMediaItem,
  CanvasOperationKind,
  CanvasOperationPlan,
} from "@/types/canvas/agent-protocol";

export type CanvasAgentDirectWritebackMode = "generate" | "edit" | "video";
export type CanvasAgentDirectWritebackResultType = "image" | "video";

export interface CanvasAgentDirectWritebackTargetRef {
  nodeId?: string | null;
}

export interface BuildCanvasAgentDirectWritebackPlanParams {
  projectId: string;
  chatId: string;
  assistantMessageId: string;
  operationLabel: string;
  resultType: CanvasAgentDirectWritebackResultType;
  resultUrls: string[];
  prompt: string;
  mode: CanvasAgentDirectWritebackMode;
  context: CanvasAgentContext;
  requestTextContent?: string | null;
  targetRefs?: CanvasAgentDirectWritebackTargetRef[];
  sourceImageUrls?: string[];
  assetIdsByUrl?: Record<string, string | undefined>;
  model?: string | null;
  targetSlotId?: string | null;
}

export interface CanvasAgentDirectWritebackPlanResult {
  plan: CanvasOperationPlan;
  mediaItems: CanvasMediaItem[];
  mediaItemIds: string[];
  operationId: string;
  operation: CanvasOperationKind;
  targetSlotId: string;
  shouldFillSlot: boolean;
  selectedMediaNodes: ReturnType<typeof getSelectedMediaNodes>;
}

function buildKnownNodeMap(context: CanvasAgentContext) {
  return new Map<string, CanvasAgentNodeRef>(
    [...context.candidateNodes, ...context.selectedNodes].map((node) => [
      node.nodeId,
      node,
    ])
  );
}

function isSlotCompatibleWithResultType(
  slot: CanvasAgentSlotRef,
  resultType: CanvasAgentDirectWritebackResultType
) {
  if (slot.slotType === "video") return resultType === "video";
  return resultType === "image";
}

function resolveTargetMediaNodes(params: {
  context: CanvasAgentContext;
  targetNodeIds: string[];
}) {
  const knownNodes = buildKnownNodeMap(params.context);
  return params.targetNodeIds
    .map((nodeId) => knownNodes.get(nodeId))
    .filter(
      (node): node is CanvasAgentNodeRef =>
        node !== undefined && (node.type === "image" || node.type === "video")
    );
}

function formatAnnotationBrief(annotation: CanvasAgentAnnotationRef) {
  const parts = [
    annotation.id,
    annotation.type,
    annotation.text ? `text=${annotation.text}` : "",
    annotation.bounds
      ? `bounds=${annotation.bounds.x},${annotation.bounds.y},${annotation.bounds.width},${annotation.bounds.height}`
      : "",
    annotation.targetPoint
      ? `target=${annotation.targetPoint.x},${annotation.targetPoint.y}`
      : "",
  ].filter(Boolean);
  return parts.join(";");
}

function buildAnnotationTraceMetadata(params: {
  context: CanvasAgentContext;
  sourceNodeIds: string[];
  userText?: string | null;
}) {
  if (!hasCanvasAgentAnnotationEditIntent(params.userText || "")) return {};
  const sourceNodeIdSet = new Set(params.sourceNodeIds);
  const annotations = params.context.annotations
    .filter((annotation) => sourceNodeIdSet.has(annotation.sourceNodeId))
    .slice(0, 8);
  if (annotations.length === 0) return {};

  return {
    canvasAgentAnnotationIds: annotations.map((annotation) => annotation.id),
    canvasAgentAnnotationCount: annotations.length,
    canvasAgentAnnotationBrief: annotations
      .map(formatAnnotationBrief)
      .join("\n")
      .slice(0, 2000),
  };
}

export function buildCanvasAgentDirectWritebackPlan(
  params: BuildCanvasAgentDirectWritebackPlanParams
): CanvasAgentDirectWritebackPlanResult | null {
  const projectId = String(params.projectId || "").trim();
  if (!projectId) return null;

  const resultUrls = params.resultUrls
    .map((url) => String(url || "").trim())
    .filter((url) => /^https?:\/\//i.test(url));
  if (resultUrls.length === 0) return null;

  const preferredTargetSlotId = String(params.targetSlotId || "").trim();
  const selectedSlots = getSelectedSlots(params.context);
  const selectedSlot =
    selectedSlots.find((slot) =>
      isSlotCompatibleWithResultType(slot, params.resultType)
    ) || null;
  const selectedSlotAnchor = selectedSlots[0] || null;
  const targetSlotId = preferredTargetSlotId || selectedSlot?.nodeId || "";
  const shouldFillSlot = Boolean(targetSlotId);
  const selectedMediaNodes = getSelectedMediaNodes(params.context);
  const targetRefs = params.targetRefs || [];
  const targetNodeIds = Array.from(
    new Set(
      targetRefs
        .map((ref) => String(ref.nodeId || "").trim())
        .filter((value): value is string => Boolean(value))
    )
  );
  const targetMediaNodes = resolveTargetMediaNodes({
    context: params.context,
    targetNodeIds,
  });
  const writebackMediaNodes =
    targetMediaNodes.length > 0 ? targetMediaNodes : selectedMediaNodes;
  const anchorNode =
    selectedSlot ||
    selectedSlotAnchor ||
    targetMediaNodes[0] ||
    selectedMediaNodes[0] ||
    getBestAnchorNode(params.context);
  const selectedMediaNodeIds = writebackMediaNodes
    .map((node) => node.nodeId)
    .filter((value): value is string => Boolean(value));
  const perResultSourceNodeIds =
    !shouldFillSlot && resultUrls.length > 1
      ? [targetNodeIds, selectedMediaNodeIds].find(
          (nodeIds) => nodeIds.length === resultUrls.length
        ) || []
      : [];
  const sourceNodeIds = Array.from(
    new Set(
      [
        ...targetNodeIds,
        ...selectedMediaNodeIds,
        anchorNode?.nodeId,
      ].filter((value): value is string => Boolean(value))
    )
  );
  const operation: CanvasOperationKind = selectedSlot
    ? "fill_slot"
    : shouldFillSlot
      ? "fill_slot"
      : params.mode === "edit"
        ? "append_revision"
        : "insert_beside";
  const hasWritebackAnchor = Boolean(anchorNode?.nodeId);
  const mediaItems: CanvasMediaItem[] = resultUrls.map((url, index) => {
    const mediaItemId = `${params.assistantMessageId}:${params.operationLabel}:${index + 1}`;
    const itemSourceNodeIds = perResultSourceNodeIds[index]
      ? [perResultSourceNodeIds[index]]
      : sourceNodeIds;
    return {
      mediaItemId,
      type: params.resultType,
      url,
      assetId: params.assetIdsByUrl?.[url],
      title:
        params.resultType === "video" ? "Agent 视频结果" : "Agent 图片结果",
      prompt: params.prompt,
      sourceNodeIds: itemSourceNodeIds,
      metadata: {
        messageId: params.assistantMessageId,
        chatId: params.chatId,
        model: params.model || undefined,
        sourceImageUrls: params.sourceImageUrls || [],
        ...buildAnnotationTraceMetadata({
          context: params.context,
          sourceNodeIds: itemSourceNodeIds,
          userText: params.requestTextContent || params.prompt,
        }),
      },
    };
  });
  const operationId = `${params.assistantMessageId}:${params.operationLabel}`;
  const plan: CanvasOperationPlan = {
    version: 1,
    operationId,
    operation,
    projectId,
    baseCanvasRevision: params.context.canvasRevision,
    baseCanvasUpdatedAt: params.context.canvasUpdatedAt,
    sourceNodeIds,
    targetSlotId: targetSlotId || undefined,
    anchorNodeId: shouldFillSlot ? undefined : anchorNode?.nodeId,
    placement: hasWritebackAnchor ? "right" : "viewport-center",
    placementStrategy:
      perResultSourceNodeIds.length === mediaItems.length
        ? "beside_each_source"
        : mediaItems.length > 1
          ? "grid"
          : "single",
    preserveOriginal: true,
    mediaItems,
    metadata: {
      messageId: params.assistantMessageId,
      chatId: params.chatId,
      mode: params.mode,
      prompt: params.prompt || params.requestTextContent || "",
      canvasAgentViewport: params.context.viewport,
    },
  };

  return {
    plan,
    mediaItems,
    mediaItemIds: mediaItems.map((item) => item.mediaItemId),
    operationId,
    operation,
    targetSlotId,
    shouldFillSlot,
    selectedMediaNodes: writebackMediaNodes,
  };
}