ai-chatbot/lib/agents/agent-project-save-merge.ts

import { normalizeAgentGenerationProjectPersistedState } from "@/lib/agents/generation-project-state";
import type {
  AgentGenerationProjectPersistedState,
  AgentGenerationProjectReplacementResultState,
  AgentGenerationProjectResultCollections,
  AgentWorkflowBusinessModeKey,
} from "@/types/agent";

const RESULT_COLLECTION_KEYS = [
  "scene_generation",
  "detail_generation",
  "detail_replacement",
  "sku_replacement",
] as const;

const REPLACEMENT_RESULT_KEYS = [
  "detail_replacement",
  "sku_replacement",
] as const;

function createFallbackWorkspaceState() {
  return {
    input: "",
    attachments: [],
    selectedModelId: "auto",
    activeBusinessModeKey: "main_sku" as const,
    imageParams: {
      aspectRatio: "auto",
      imageSize: "2K",
      imageCount: 1,
    },
  };
}

function normalizeProjectResultUrlList(urls: string[] | null | undefined) {
  return Array.from(
    new Set(
      (urls || [])
        .map((url) => String(url || "").trim())
        .filter((url) => /^https?:\/\//i.test(url))
    )
  );
}

function mergeResultCollections(
  existing: AgentGenerationProjectResultCollections | null | undefined,
  incoming: AgentGenerationProjectResultCollections | null | undefined
): AgentGenerationProjectResultCollections {
  return RESULT_COLLECTION_KEYS.reduce<AgentGenerationProjectResultCollections>(
    (next, key) => {
      next[key] = normalizeProjectResultUrlList([
        ...normalizeProjectResultUrlList(existing?.[key]),
        ...normalizeProjectResultUrlList(incoming?.[key]),
      ]);
      return next;
    },
    {}
  );
}

function getResultCollectionUrlSet(
  collections: AgentGenerationProjectResultCollections | null | undefined
) {
  return new Set(
    RESULT_COLLECTION_KEYS.flatMap((key) => normalizeProjectResultUrlList(collections?.[key]))
  );
}

function removeDeletedUrlsFromCollections(
  collections: AgentGenerationProjectResultCollections,
  deletedUrls: Set<string>
): AgentGenerationProjectResultCollections {
  return RESULT_COLLECTION_KEYS.reduce<AgentGenerationProjectResultCollections>(
    (next, key) => {
      next[key] = normalizeProjectResultUrlList(collections[key]).filter(
        (url) => !deletedUrls.has(url)
      );
      return next;
    },
    {}
  );
}

function mergeReplacementResultState(
  existing: AgentGenerationProjectReplacementResultState | null | undefined,
  incoming: AgentGenerationProjectReplacementResultState | null | undefined,
  deletedUrls: Set<string>
): AgentGenerationProjectReplacementResultState {
  return REPLACEMENT_RESULT_KEYS.reduce<AgentGenerationProjectReplacementResultState>(
    (next, key) => {
      const existingModeState = existing?.[key] || {};
      const incomingModeState = incoming?.[key] || {};
      const mergedModeState = { ...existingModeState, ...incomingModeState };
      next[key] = Object.fromEntries(
        Object.entries(mergedModeState).filter(([, record]) => {
          const url = String(record?.url || "").trim();
          return !url || !deletedUrls.has(url);
        })
      );
      return next;
    },
    {}
  );
}

function getReplacementResultUrls(
  state: AgentGenerationProjectReplacementResultState | null | undefined
) {
  return REPLACEMENT_RESULT_KEYS.flatMap((key) =>
    Object.values(state?.[key] || {})
      .map((record) => String(record?.url || "").trim())
      .filter((url) => /^https?:\/\//i.test(url))
  );
}

export function getAllResultUrlsFromAgentProjectState(
  state: AgentGenerationProjectPersistedState
) {
  return normalizeProjectResultUrlList([
    ...normalizeProjectResultUrlList(state.resultImageUrls),
    ...RESULT_COLLECTION_KEYS.flatMap((key) =>
      normalizeProjectResultUrlList(state.resultCollections?.[key])
    ),
    ...getReplacementResultUrls(state.replacementResultState),
  ]);
}

function hasDetailModeSignals(state: AgentGenerationProjectPersistedState) {
  const detailReferenceCount =
    state.workflowState.agentModeState?.detailReplacement?.references?.length || 0;
  const skuReferenceCount =
    state.workflowState.agentModeState?.skuReplacement?.references?.length || 0;
  return (
    detailReferenceCount > 0 ||
    skuReferenceCount > 0 ||
    normalizeProjectResultUrlList(state.resultCollections?.detail_generation).length > 0 ||
    normalizeProjectResultUrlList(state.resultCollections?.detail_replacement).length > 0 ||
    normalizeProjectResultUrlList(state.resultCollections?.sku_replacement).length > 0 ||
    getReplacementResultUrls(state.replacementResultState).length > 0
  );
}

function hasDetailReplacementResultSignals(params: {
  resultCollections: AgentGenerationProjectResultCollections;
  replacementResultState: AgentGenerationProjectReplacementResultState;
}) {
  return (
    normalizeProjectResultUrlList(params.resultCollections.detail_replacement).length > 0 ||
    getReplacementResultUrls({
      detail_replacement: params.replacementResultState.detail_replacement,
    }).length > 0
  );
}

function hasSkuReplacementResultSignals(params: {
  resultCollections: AgentGenerationProjectResultCollections;
  replacementResultState: AgentGenerationProjectReplacementResultState;
}) {
  return (
    normalizeProjectResultUrlList(params.resultCollections.sku_replacement).length > 0 ||
    getReplacementResultUrls({
      sku_replacement: params.replacementResultState.sku_replacement,
    }).length > 0
  );
}

function mergeWorkflowStateForSave(params: {
  existingState: AgentGenerationProjectPersistedState;
  incomingState: AgentGenerationProjectPersistedState;
  resultCollections: AgentGenerationProjectResultCollections;
  replacementResultState: AgentGenerationProjectReplacementResultState;
}) {
  const incomingAgentModeState = params.incomingState.workflowState.agentModeState;
  const existingAgentModeState = params.existingState.workflowState.agentModeState;
  const shouldKeepDetailReplacementState =
    hasDetailReplacementResultSignals(params) &&
    (existingAgentModeState?.detailReplacement?.references?.length || 0) > 0 &&
    (incomingAgentModeState?.detailReplacement?.references?.length || 0) === 0;
  const shouldKeepSkuReplacementState =
    hasSkuReplacementResultSignals(params) &&
    (existingAgentModeState?.skuReplacement?.references?.length || 0) > 0 &&
    (incomingAgentModeState?.skuReplacement?.references?.length || 0) === 0;

  if (!shouldKeepDetailReplacementState && !shouldKeepSkuReplacementState) {
    return params.incomingState.workflowState;
  }

  return {
    ...params.incomingState.workflowState,
    agentModeState: {
      ...incomingAgentModeState,
      detailReplacement: shouldKeepDetailReplacementState
        ? existingAgentModeState?.detailReplacement
        : incomingAgentModeState?.detailReplacement,
      skuReplacement: shouldKeepSkuReplacementState
        ? existingAgentModeState?.skuReplacement
        : incomingAgentModeState?.skuReplacement,
    },
  };
}

function hasExplicitIncomingBusinessMode(value: unknown) {
  if (!value || typeof value !== "object" || Array.isArray(value)) {
    return false;
  }
  const raw = value as Record<string, unknown>;
  const workspaceState =
    raw.workspaceState && typeof raw.workspaceState === "object" && !Array.isArray(raw.workspaceState)
      ? (raw.workspaceState as Record<string, unknown>)
      : {};
  return Boolean(String(workspaceState.activeBusinessModeKey || "").trim());
}

function resolveMergedBusinessModeKey(params: {
  existingState: AgentGenerationProjectPersistedState;
  incomingState: AgentGenerationProjectPersistedState;
  rawIncomingState: unknown;
  mergedState: AgentGenerationProjectPersistedState;
}): AgentWorkflowBusinessModeKey {
  if (hasDetailModeSignals(params.mergedState)) {
    const incomingKey = params.incomingState.workspaceState?.activeBusinessModeKey;
    const existingKey = params.existingState.workspaceState?.activeBusinessModeKey;
    const incomingHasExplicitKey = hasExplicitIncomingBusinessMode(params.rawIncomingState);

    if (!incomingHasExplicitKey || existingKey === "detail" || incomingKey === "main_sku") {
      return "detail";
    }
  }

  return params.incomingState.workspaceState?.activeBusinessModeKey ||
    params.existingState.workspaceState?.activeBusinessModeKey ||
    "main_sku";
}

export function mergeAgentProjectStateForSave(params: {
  existingState: AgentGenerationProjectPersistedState;
  incomingState: unknown;
}): AgentGenerationProjectPersistedState {
  const existingState = normalizeAgentGenerationProjectPersistedState(params.existingState);
  const incomingState = normalizeAgentGenerationProjectPersistedState(params.incomingState);
  const deletedResultCollections = mergeResultCollections(
    existingState.deletedResultCollections,
    incomingState.deletedResultCollections
  );
  const deletedUrls = getResultCollectionUrlSet(deletedResultCollections);
  const replacementResultState = mergeReplacementResultState(
    existingState.replacementResultState,
    incomingState.replacementResultState,
    deletedUrls
  );
  const resultCollections = removeDeletedUrlsFromCollections(
    mergeResultCollections(existingState.resultCollections, incomingState.resultCollections),
    deletedUrls
  );
  const workflowState = mergeWorkflowStateForSave({
    existingState,
    incomingState,
    resultCollections,
    replacementResultState,
  });
  const mergedState = normalizeAgentGenerationProjectPersistedState({
    ...incomingState,
    workflowState,
    resultCollections,
    deletedResultCollections,
    replacementResultState,
    resultImageUrls: normalizeProjectResultUrlList([
      ...getAllResultUrlsFromAgentProjectState(existingState),
      ...getAllResultUrlsFromAgentProjectState(incomingState),
      ...RESULT_COLLECTION_KEYS.flatMap((key) => resultCollections[key] || []),
      ...getReplacementResultUrls(replacementResultState),
    ]).filter((url) => !deletedUrls.has(url)),
  });

  const workspaceState = mergedState.workspaceState || createFallbackWorkspaceState();

  return {
    ...mergedState,
    workspaceState: {
      ...workspaceState,
      activeBusinessModeKey: resolveMergedBusinessModeKey({
        existingState,
        incomingState,
        rawIncomingState: params.incomingState,
        mergedState,
      }),
    },
  };
}