ai-chatbot/lib/ai/replacement-layout-lock.ts

export type NormalizedReplacementBox = [number, number, number, number];

export type ReplacementLayoutLockRegions = {
  subjectBoundingBox: NormalizedReplacementBox | null;
  textBoundingBoxes: NormalizedReplacementBox[];
  analyzed: boolean;
};

type PixelBox = {
  left: number;
  top: number;
  right: number;
  bottom: number;
};

function extractJsonRecord(text: string): Record<string, unknown> | null {
  const trimmed = String(text || "")
    .trim()
    .replace(/^```(?:json)?\s*/i, "")
    .replace(/\s*```$/i, "");
  const start = trimmed.indexOf("{");
  const end = trimmed.lastIndexOf("}");
  if (start < 0 || end <= start) {
    return null;
  }

  try {
    const parsed = JSON.parse(trimmed.slice(start, end + 1));
    return parsed && typeof parsed === "object" && !Array.isArray(parsed)
      ? (parsed as Record<string, unknown>)
      : null;
  } catch {
    return null;
  }
}

function clampNormalizedCoordinate(value: unknown) {
  const numberValue = Number(value);
  if (!Number.isFinite(numberValue)) {
    return null;
  }
  return Math.max(0, Math.min(1000, numberValue));
}

export function normalizeReplacementBoundingBox(
  value: unknown,
): NormalizedReplacementBox | null {
  let coordinates: unknown[] | null = null;

  if (Array.isArray(value)) {
    coordinates = value;
  } else if (value && typeof value === "object") {
    const record = value as Record<string, unknown>;
    if (Array.isArray(record.bbox)) {
      coordinates = record.bbox;
    } else if (
      record.x1 !== undefined &&
      record.y1 !== undefined &&
      record.x2 !== undefined &&
      record.y2 !== undefined
    ) {
      coordinates = [record.x1, record.y1, record.x2, record.y2];
    } else if (
      record.left !== undefined &&
      record.top !== undefined &&
      record.right !== undefined &&
      record.bottom !== undefined
    ) {
      coordinates = [record.left, record.top, record.right, record.bottom];
    }
  }

  if (!coordinates || coordinates.length !== 4) {
    return null;
  }

  const normalized = coordinates.map(clampNormalizedCoordinate);
  if (normalized.some((coordinate) => coordinate === null)) {
    return null;
  }

  const [x1, y1, x2, y2] = normalized as NormalizedReplacementBox;
  if (x2 - x1 < 2 || y2 - y1 < 2) {
    return null;
  }

  return [x1, y1, x2, y2];
}

export function expandReplacementSubjectBoundingBoxForGeneration(
  value: unknown,
): NormalizedReplacementBox | null {
  const box = normalizeReplacementBoundingBox(value);
  if (!box) return null;

  const [x1, y1, x2, y2] = box;
  const horizontalPadding = (x2 - x1) * 0.1;
  const verticalPadding = (y2 - y1) * 0.015;
  return [
    Math.max(0, Math.round(x1 - horizontalPadding)),
    Math.max(0, Math.round(y1 - verticalPadding)),
    Math.min(1000, Math.round(x2 + horizontalPadding)),
    Math.min(1000, Math.round(y2 + verticalPadding)),
  ];
}

export function buildReplacementInputBoundingBoxList(params: {
  imageUrls: string[];
  referenceImageUrl: string;
  subjectBoundingBox: unknown;
}): NormalizedReplacementBox[][] {
  const referenceImageUrl = String(params.referenceImageUrl || "").trim();
  const expandedBox = expandReplacementSubjectBoundingBoxForGeneration(
    params.subjectBoundingBox,
  );

  return params.imageUrls.map((imageUrl) =>
    expandedBox && String(imageUrl || "").trim() === referenceImageUrl
      ? [expandedBox]
      : [],
  );
}

function findFirstOwnValue(record: Record<string, unknown>, keys: string[]) {
  for (const key of keys) {
    if (Object.prototype.hasOwnProperty.call(record, key)) {
      return { found: true, value: record[key] };
    }
  }
  return { found: false, value: undefined };
}

export function parseReplacementLayoutLockRegions(
  textOrRecord: string | Record<string, unknown>,
): ReplacementLayoutLockRegions {
  const record =
    typeof textOrRecord === "string" ? extractJsonRecord(textOrRecord) : textOrRecord;
  if (!record) {
    return {
      subjectBoundingBox: null,
      textBoundingBoxes: [],
      analyzed: false,
    };
  }

  const subjectValue = findFirstOwnValue(record, [
    "subjectBoundingBox",
    "productBoundingBox",
    "replacementBoundingBox",
  ]);
  const textValue = findFirstOwnValue(record, [
    "textBoundingBoxes",
    "protectedTextBoundingBoxes",
    "copyBoundingBoxes",
  ]);
  const rawTextBoxes = Array.isArray(textValue.value) ? textValue.value : [];
  const textBoundingBoxes = rawTextBoxes
    .map(normalizeReplacementBoundingBox)
    .filter((box): box is NormalizedReplacementBox => Boolean(box))
    .slice(0, 60);

  return {
    subjectBoundingBox: normalizeReplacementBoundingBox(subjectValue.value),
    textBoundingBoxes,
    analyzed: subjectValue.found && textValue.found,
  };
}

function normalizedToPixelBox(
  box: NormalizedReplacementBox,
  width: number,
  height: number,
): PixelBox {
  return {
    left: (box[0] / 1000) * width,
    top: (box[1] / 1000) * height,
    right: (box[2] / 1000) * width,
    bottom: (box[3] / 1000) * height,
  };
}

function expandPixelBox(
  box: PixelBox,
  expansion: { left: number; top: number; right: number; bottom: number },
  width: number,
  height: number,
): PixelBox {
  return {
    left: Math.max(0, box.left - expansion.left),
    top: Math.max(0, box.top - expansion.top),
    right: Math.min(width, box.right + expansion.right),
    bottom: Math.min(height, box.bottom + expansion.bottom),
  };
}

function paintFeatheredRectangle(
  mask: Buffer,
  width: number,
  height: number,
  box: PixelBox,
  feather: number,
) {
  const startX = Math.max(0, Math.floor(box.left));
  const endX = Math.min(width - 1, Math.ceil(box.right));
  const startY = Math.max(0, Math.floor(box.top));
  const endY = Math.min(height - 1, Math.ceil(box.bottom));

  for (let y = startY; y <= endY; y += 1) {
    const verticalDistance = Math.min(y - box.top, box.bottom - y);
    if (verticalDistance < 0) {
      continue;
    }
    for (let x = startX; x <= endX; x += 1) {
      const horizontalDistance = Math.min(x - box.left, box.right - x);
      if (horizontalDistance < 0) {
        continue;
      }
      const edgeDistance = Math.min(horizontalDistance, verticalDistance);
      const alpha = Math.round(255 * Math.min(1, edgeDistance / Math.max(1, feather)));
      const offset = y * width + x;
      if (alpha > mask[offset]!) {
        mask[offset] = alpha;
      }
    }
  }
}

function buildTextLockMask(params: {
  boxes: NormalizedReplacementBox[];
  width: number;
  height: number;
  feather: number;
}) {
  const mask = Buffer.alloc(params.width * params.height);
  const padding = Math.max(4, Math.round(Math.min(params.width, params.height) * 0.006));

  for (const normalizedBox of params.boxes) {
    const box = normalizedToPixelBox(normalizedBox, params.width, params.height);
    paintFeatheredRectangle(
      mask,
      params.width,
      params.height,
      expandPixelBox(
        box,
        {
          left: padding + params.feather,
          top: padding + params.feather,
          right: padding + params.feather,
          bottom: padding + params.feather,
        },
        params.width,
        params.height,
      ),
      params.feather,
    );
  }

  return mask;
}

async function addAlphaMask(params: {
  buffer: Buffer;
  mask: Buffer;
  width: number;
  height: number;
}) {
  const sharp = (await import("sharp")).default;
  const opaqueBuffer = await sharp(params.buffer, { failOn: "none" })
    .removeAlpha()
    .png()
    .toBuffer();
  return sharp(opaqueBuffer, { failOn: "none" })
    .joinChannel(params.mask, {
      raw: { width: params.width, height: params.height, channels: 1 },
    })
    .png()
    .toBuffer();
}

export async function compositeReplacementWithReferenceLocks(params: {
  generatedBuffer: Buffer;
  referenceBuffer: Buffer;
  width: number;
  height: number;
  preserveReferenceOutsideSubject: boolean;
  subjectBoundingBox?: NormalizedReplacementBox | null;
  textBoundingBoxes?: NormalizedReplacementBox[];
}) {
  const sharp = (await import("sharp")).default;
  const width = Math.max(1, Math.round(params.width));
  const height = Math.max(1, Math.round(params.height));
  const textBoundingBoxes = params.textBoundingBoxes || [];
  const feather = Math.max(3, Math.min(12, Math.round(Math.min(width, height) * 0.006)));
  const [generatedBuffer, referenceBuffer] = await Promise.all([
    sharp(params.generatedBuffer, { failOn: "none" })
      .rotate()
      .resize(width, height, { fit: "fill" })
      .png()
      .toBuffer(),
    sharp(params.referenceBuffer, { failOn: "none" })
      .rotate()
      .resize(width, height, { fit: "fill" })
      .png()
      .toBuffer(),
  ]);

  if (params.preserveReferenceOutsideSubject && params.subjectBoundingBox) {
    const subjectBox = normalizedToPixelBox(params.subjectBoundingBox, width, height);
    const subjectWidth = subjectBox.right - subjectBox.left;
    const subjectHeight = subjectBox.bottom - subjectBox.top;
    const expandedSubjectBox = expandPixelBox(
      subjectBox,
      {
        left: subjectWidth * 0.22,
        right: subjectWidth * 0.22,
        top: subjectHeight * 0.035,
        bottom: subjectHeight * 0.025,
      },
      width,
      height,
    );
    const subjectMask = Buffer.alloc(width * height);
    paintFeatheredRectangle(subjectMask, width, height, expandedSubjectBox, feather);

    if (textBoundingBoxes.length > 0) {
      const textMask = buildTextLockMask({
        boxes: textBoundingBoxes,
        width,
        height,
        feather,
      });
      for (let index = 0; index < subjectMask.length; index += 1) {
        subjectMask[index] = Math.round(
          (subjectMask[index]! * (255 - textMask[index]!)) / 255,
        );
      }
    }

    const generatedLayer = await addAlphaMask({
      buffer: generatedBuffer,
      mask: subjectMask,
      width,
      height,
    });
    return sharp(referenceBuffer, { failOn: "none" })
      .composite([{ input: generatedLayer }])
      .png()
      .toBuffer();
  }

  if (textBoundingBoxes.length === 0) {
    return generatedBuffer;
  }

  const textMask = buildTextLockMask({
    boxes: textBoundingBoxes,
    width,
    height,
    feather,
  });
  const referenceTextLayer = await addAlphaMask({
    buffer: referenceBuffer,
    mask: textMask,
    width,
    height,
  });
  return sharp(generatedBuffer, { failOn: "none" })
    .composite([{ input: referenceTextLayer }])
    .png()
    .toBuffer();
}