ai-chatbot/lib/canvas/agent-context.ts

import type {
  CanvasAgentAnnotationRef,
  CanvasAgentContext,
  CanvasAgentFocusPolicy,
  CanvasAgentNodeRef,
  CanvasAgentNodeType,
  CanvasAgentSlotRef,
  CanvasAgentSource,
  CanvasViewportContext,
} from "@/types/canvas/agent-protocol";

type LegacyCanvasImageNode = {
  nodeId?: unknown;
  assetId?: unknown;
  url?: unknown;
  title?: unknown;
  prompt?: unknown;
  x?: unknown;
  y?: unknown;
  width?: unknown;
  height?: unknown;
  isSelected?: unknown;
};

export type CanvasAgentContextPayloadLike = {
  source?: unknown;
  projectId?: unknown;
  selectedNodeIds?: unknown;
  selectedAssetIds?: unknown;
  focusPolicy?: unknown;
  canvasImageNodes?: unknown;
  canvasAgentContext?: unknown;
};

const CANVAS_AGENT_NODE_TYPES = new Set<string>([
  "message",
  "image",
  "video",
  "text",
  "ui-screen",
  "workflow",
  "scene-fusion",
  "model3d",
  "world3d",
  "agent-slot",
]);

const CANVAS_AGENT_FOCUS_POLICIES = new Set<string>([
  "auto",
  "selected",
  "latest_result",
  "original",
  "slot",
]);

function normalizeString(value: unknown, maxLength = 512) {
  if (typeof value !== "string") return undefined;
  const normalized = value.trim();
  return normalized ? normalized.slice(0, maxLength) : undefined;
}

function normalizeStringArray(value: unknown, maxItems: number) {
  if (!Array.isArray(value)) return [];
  return value
    .map((item) => normalizeString(item, 128))
    .filter((item): item is string => Boolean(item))
    .slice(0, maxItems);
}

function normalizeNumber(value: unknown) {
  const parsed = typeof value === "number" ? value : Number(value);
  return Number.isFinite(parsed) ? parsed : undefined;
}

function normalizeBoolean(value: unknown) {
  return typeof value === "boolean" ? value : undefined;
}

function normalizeSource(value: unknown): CanvasAgentSource {
  return value === "realtime-canvas" ? "realtime-canvas" : "canvas";
}

function normalizeFocusPolicy(value: unknown): CanvasAgentFocusPolicy {
  return CANVAS_AGENT_FOCUS_POLICIES.has(String(value))
    ? (value as CanvasAgentFocusPolicy)
    : "auto";
}

function normalizeNodeType(value: unknown): CanvasAgentNodeType | undefined {
  const normalized = normalizeString(value, 64);
  if (!normalized || !CANVAS_AGENT_NODE_TYPES.has(normalized)) return undefined;
  return normalized as CanvasAgentNodeType;
}

function normalizeViewport(value: unknown): CanvasViewportContext | undefined {
  if (!value || typeof value !== "object" || Array.isArray(value)) return undefined;
  const raw = value as Record<string, unknown>;
  const x = normalizeNumber(raw.x);
  const y = normalizeNumber(raw.y);
  const zoom = normalizeNumber(raw.zoom);
  if (x === undefined || y === undefined || zoom === undefined) return undefined;
  return {
    x,
    y,
    zoom,
    width: normalizeNumber(raw.width),
    height: normalizeNumber(raw.height),
  };
}

function normalizeNodeRef(value: unknown): CanvasAgentNodeRef | null {
  if (!value || typeof value !== "object" || Array.isArray(value)) return null;
  const raw = value as Record<string, unknown>;
  const nodeId = normalizeString(raw.nodeId, 128);
  const type = normalizeNodeType(raw.type);
  const x = normalizeNumber(raw.x);
  const y = normalizeNumber(raw.y);
  if (!nodeId || !type || x === undefined || y === undefined) return null;
  return {
    nodeId,
    type,
    assetId: normalizeString(raw.assetId, 128),
    url: normalizeString(raw.url, 2048),
    title: normalizeString(raw.title, 256),
    prompt: normalizeString(raw.prompt, 2000),
    x,
    y,
    width: normalizeNumber(raw.width),
    height: normalizeNumber(raw.height),
    isSelected: normalizeBoolean(raw.isSelected),
    metadata:
      raw.metadata && typeof raw.metadata === "object" && !Array.isArray(raw.metadata)
        ? (raw.metadata as Record<string, unknown>)
        : undefined,
  };
}

function normalizeSlotRef(value: unknown): CanvasAgentSlotRef | null {
  if (!value || typeof value !== "object" || Array.isArray(value)) return null;
  const raw = value as Record<string, unknown>;
  const nodeId = normalizeString(raw.nodeId, 128);
  const slotType = normalizeString(raw.slotType, 32);
  const x = normalizeNumber(raw.x);
  const y = normalizeNumber(raw.y);
  const width = normalizeNumber(raw.width);
  const height = normalizeNumber(raw.height);
  if (
    !nodeId ||
    (slotType !== "image" && slotType !== "video" && slotType !== "ui-screen") ||
    x === undefined ||
    y === undefined ||
    width === undefined ||
    height === undefined
  ) {
    return null;
  }
  const fillPolicy = normalizeString(raw.fillPolicy, 32);
  return {
    nodeId,
    slotType,
    label: normalizeString(raw.label, 128),
    x,
    y,
    width,
    height,
    aspectRatio: normalizeString(raw.aspectRatio, 32) || "1:1",
    isSelected: normalizeBoolean(raw.isSelected),
    fillPolicy:
      fillPolicy === "append-version" || fillPolicy === "replace-existing"
        ? fillPolicy
        : "replace-empty",
    acceptedSources: Array.isArray(raw.acceptedSources)
      ? raw.acceptedSources
          .map((item) => normalizeString(item, 32))
          .filter(
            (item): item is "text" | "image" | "selection" | "annotation" =>
              item === "text" ||
              item === "image" ||
              item === "selection" ||
              item === "annotation"
          )
          .slice(0, 8)
      : undefined,
  };
}

function normalizeAnnotationRef(value: unknown): CanvasAgentAnnotationRef | null {
  if (!value || typeof value !== "object" || Array.isArray(value)) return null;
  const raw = value as Record<string, unknown>;
  const id = normalizeString(raw.id, 128);
  const sourceNodeId = normalizeString(raw.sourceNodeId, 128);
  const type = normalizeString(raw.type, 32);
  if (
    !id ||
    !sourceNodeId ||
    (type !== "text-note" &&
      type !== "arrow-note" &&
      type !== "box" &&
      type !== "freehand" &&
      type !== "crop-box")
  ) {
    return null;
  }
  const rawBounds =
    raw.bounds && typeof raw.bounds === "object" && !Array.isArray(raw.bounds)
      ? (raw.bounds as Record<string, unknown>)
      : null;
  const rawTargetPoint =
    raw.targetPoint && typeof raw.targetPoint === "object" && !Array.isArray(raw.targetPoint)
      ? (raw.targetPoint as Record<string, unknown>)
      : null;
  return {
    id,
    sourceNodeId,
    type,
    text: normalizeString(raw.text, 1000),
    bounds: rawBounds
      ? {
          x: normalizeNumber(rawBounds.x) || 0,
          y: normalizeNumber(rawBounds.y) || 0,
          width: normalizeNumber(rawBounds.width) || 0,
          height: normalizeNumber(rawBounds.height) || 0,
        }
      : undefined,
    targetPoint: rawTargetPoint
      ? {
          x: normalizeNumber(rawTargetPoint.x) || 0,
          y: normalizeNumber(rawTargetPoint.y) || 0,
        }
      : undefined,
    color: normalizeString(raw.color, 64),
    createdAt: normalizeString(raw.createdAt, 64),
  };
}

function normalizeNodeRefs(value: unknown, maxItems: number) {
  if (!Array.isArray(value)) return [];
  return value
    .map(normalizeNodeRef)
    .filter((node): node is CanvasAgentNodeRef => Boolean(node))
    .slice(0, maxItems);
}

function normalizeSlotRefs(value: unknown, maxItems: number) {
  if (!Array.isArray(value)) return [];
  return value
    .map(normalizeSlotRef)
    .filter((slot): slot is CanvasAgentSlotRef => Boolean(slot))
    .slice(0, maxItems);
}

function normalizeAnnotationRefs(value: unknown, maxItems: number) {
  if (!Array.isArray(value)) return [];
  return value
    .map(normalizeAnnotationRef)
    .filter((annotation): annotation is CanvasAgentAnnotationRef => Boolean(annotation))
    .slice(0, maxItems);
}

function legacyImageToNodeRef(value: LegacyCanvasImageNode): CanvasAgentNodeRef | null {
  const nodeId = normalizeString(value.nodeId, 128);
  const url = normalizeString(value.url, 2048);
  const x = normalizeNumber(value.x);
  const y = normalizeNumber(value.y);
  if (!nodeId || !url || x === undefined || y === undefined) return null;
  return {
    nodeId,
    type: "image",
    assetId: normalizeString(value.assetId, 128),
    url,
    title: normalizeString(value.title, 256),
    prompt: normalizeString(value.prompt, 2000),
    x,
    y,
    width: normalizeNumber(value.width),
    height: normalizeNumber(value.height),
    isSelected: normalizeBoolean(value.isSelected),
  };
}

function normalizeLegacyImageNodes(value: unknown, maxItems: number) {
  if (!Array.isArray(value)) return [];
  return value
    .map((item) => (item && typeof item === "object" ? legacyImageToNodeRef(item) : null))
    .filter((node): node is CanvasAgentNodeRef => Boolean(node))
    .slice(0, maxItems);
}

function hasAgentContextValue(context: CanvasAgentContext) {
  return (
    context.selectedNodes.length > 0 ||
    context.candidateNodes.length > 0 ||
    context.slots.length > 0 ||
    context.annotations.length > 0 ||
    Boolean(context.viewport)
  );
}

function normalizeExplicitCanvasAgentContext(
  value: unknown,
  fallback: CanvasAgentContextPayloadLike
): CanvasAgentContext | null {
  if (!value || typeof value !== "object" || Array.isArray(value)) return null;
  const raw = value as Record<string, unknown>;
  const selectedNodeIds = normalizeStringArray(raw.selectedNodeIds, 20);
  const selectedNodes = normalizeNodeRefs(raw.selectedNodes, 20);
  const candidateNodes = normalizeNodeRefs(raw.candidateNodes, 30);
  const slots = normalizeSlotRefs(raw.slots, 20);
  const annotations = normalizeAnnotationRefs(raw.annotations, 20);
  const context: CanvasAgentContext = {
    version: 1,
    source: normalizeSource(raw.source || fallback.source),
    projectId: normalizeString(raw.projectId || fallback.projectId, 128),
    canvasRevision: normalizeString(raw.canvasRevision, 128),
    canvasUpdatedAt: normalizeString(raw.canvasUpdatedAt, 128),
    viewport: normalizeViewport(raw.viewport),
    selectedNodeIds:
      selectedNodeIds.length > 0
        ? selectedNodeIds
        : selectedNodes.map((node) => node.nodeId).slice(0, 20),
    selectedNodes,
    candidateNodes,
    slots,
    annotations,
    focusPolicy: normalizeFocusPolicy(raw.focusPolicy || fallback.focusPolicy),
  };
  return hasAgentContextValue(context) ? context : null;
}

export function normalizeCanvasAgentContext(
  payload?: CanvasAgentContextPayloadLike | null
): CanvasAgentContext | null {
  if (!payload) return null;

  const explicit = normalizeExplicitCanvasAgentContext(
    payload.canvasAgentContext,
    payload
  );
  if (explicit) return explicit;

  const candidateNodes = normalizeLegacyImageNodes(payload.canvasImageNodes, 20);
  const selectedNodeIdSet = new Set(normalizeStringArray(payload.selectedNodeIds, 20));
  const selectedAssetIdSet = new Set(normalizeStringArray(payload.selectedAssetIds, 20));
  const selectedNodes = candidateNodes.filter((node) => {
    if (node.isSelected) return true;
    if (selectedNodeIdSet.has(node.nodeId)) return true;
    if (node.assetId && selectedAssetIdSet.has(node.assetId)) return true;
    return false;
  });

  const context: CanvasAgentContext = {
    version: 1,
    source: normalizeSource(payload.source),
    projectId: normalizeString(payload.projectId, 128),
    selectedNodeIds:
      selectedNodes.length > 0
        ? selectedNodes.map((node) => node.nodeId).slice(0, 20)
        : Array.from(selectedNodeIdSet).slice(0, 20),
    selectedNodes,
    candidateNodes,
    slots: [],
    annotations: [],
    focusPolicy: normalizeFocusPolicy(payload.focusPolicy),
  };

  return hasAgentContextValue(context) ? context : null;
}

export function getSelectedSlots(context?: CanvasAgentContext | null) {
  return (context?.slots || []).filter(
    (slot) => slot.isSelected || context?.selectedNodeIds.includes(slot.nodeId)
  );
}

export function getSelectedMediaNodes(context?: CanvasAgentContext | null) {
  return (context?.selectedNodes || []).filter(
    (node) => node.type === "image" || node.type === "video"
  );
}

export function getBestAnchorNode(context?: CanvasAgentContext | null) {
  return (
    getSelectedMediaNodes(context)[0] ||
    context?.selectedNodes[0] ||
    context?.candidateNodes[0] ||
    null
  );
}

export function getCanvasAgentFocusPolicy(context?: CanvasAgentContext | null) {
  return context?.focusPolicy || "auto";
}