ai-chatbot/lib/agents/psd-template-parser.ts

import "server-only";

import { uploadFile } from "@/lib/storage/service";
import {
  isAgentPsdImageLayerSlot,
  type AgentPsdTemplateLayerSlot,
} from "@/lib/agents/psd-template-config";

type PsdLayerLike = {
  type?: string;
  name?: string;
  children?: PsdLayerLike[];
  width?: number;
  height?: number;
  left?: number;
  top?: number;
  isHidden?: boolean;
  text?: string;
  additionalProperties?: Record<string, unknown>;
  composite?: (effect?: boolean, composed?: boolean) => Promise<Uint8ClampedArray>;
};

type PsdRootLike = {
  width: number;
  height: number;
  children: PsdLayerLike[];
  composite: () => Promise<Uint8ClampedArray>;
};

export interface ParsedPsdTemplateLayer {
  id: string;
  layerId: string;
  layerName: string;
  layerPath: string;
  depth: number;
  x: number;
  y: number;
  width: number;
  height: number;
  slotType: "image" | "text";
  defaultText?: string;
}

export interface ParsedPsdTemplatePayload {
  width: number;
  height: number;
  previewBuffer: Buffer;
  layers: ParsedPsdTemplateLayer[];
}

interface PsdLayerNodeEntry {
  node: PsdLayerLike;
  layer: ParsedPsdTemplateLayer;
}

function normalizeLayerName(value: unknown, fallback: string) {
  const text = String(value || "").trim();
  return text || fallback;
}

function sanitizeFileSegment(value: string, fallback: string) {
  const normalized = value
    .trim()
    .replace(/\s+/g, "-")
    .replace(/[^a-zA-Z0-9_-]/g, "-")
    .replace(/-+/g, "-")
    .replace(/^-+|-+$/g, "")
    .toLowerCase();
  return normalized || fallback;
}

async function getPsdConstructor() {
  const module = await import("@webtoon/psd");
  return module.default;
}

async function getSharp() {
  const sharp = await import("sharp");
  return sharp.default;
}

async function rgbaToPngBuffer(rgba: Uint8ClampedArray, width: number, height: number) {
  const sharp = await getSharp();
  return sharp(Buffer.from(rgba), {
    raw: {
      width,
      height,
      channels: 4,
    },
  })
    .png()
    .toBuffer();
}

function extractLayerId(node: PsdLayerLike, fallback: string) {
  const additional = node.additionalProperties;
  const layerIdRecord =
    additional && typeof additional === "object"
      ? (additional.lyid as { value?: unknown } | undefined)
      : undefined;
  const layerId = layerIdRecord?.value;
  if (typeof layerId === "number" || typeof layerId === "string") {
    const normalized = String(layerId).trim();
    if (normalized) return normalized;
  }
  return fallback;
}

function collectPsdLayerEntries(
  nodes: PsdLayerLike[],
  ancestors: string[] = [],
  depth = 0,
  result: PsdLayerNodeEntry[] = [],
  orderRef: { current: number } = { current: 0 }
) {
  for (const node of nodes) {
    if (!node || typeof node !== "object") continue;

    const nodeType = String(node.type || "");
    if (nodeType === "Group") {
      const groupName = normalizeLayerName(node.name, `分组${orderRef.current + 1}`);
      collectPsdLayerEntries(
        Array.isArray(node.children) ? node.children : [],
        [...ancestors, groupName],
        depth + 1,
        result,
        orderRef
      );
      continue;
    }

    if (nodeType !== "Layer") {
      continue;
    }

    const width = Math.max(0, Math.round(Number(node.width || 0)));
    const height = Math.max(0, Math.round(Number(node.height || 0)));
    if (width <= 0 || height <= 0 || node.isHidden) {
      continue;
    }

    const order = orderRef.current;
    orderRef.current += 1;

    const layerName = normalizeLayerName(node.name, `图层${order + 1}`);
    const fallbackLayerId = `${ancestors.join("/")}/${layerName}#${order + 1}`;
    const layerId = extractLayerId(node, fallbackLayerId);
    const defaultText = String(node.text || "").trim();

    result.push({
      node,
      layer: {
        id: `layer-${sanitizeFileSegment(layerId, `layer-${order + 1}`)}`,
        layerId,
        layerName,
        layerPath: [...ancestors, layerName].join(" / "),
        depth,
        x: Math.max(0, Math.round(Number(node.left || 0))),
        y: Math.max(0, Math.round(Number(node.top || 0))),
        width,
        height,
        slotType: defaultText ? "text" : "image",
        defaultText: defaultText || undefined,
      },
    });
  }

  return result;
}

async function parsePsdBuffer(buffer: ArrayBuffer) {
  const Psd = await getPsdConstructor();
  return Psd.parse(buffer) as PsdRootLike;
}

async function fetchRemoteBuffer(url: string) {
  const response = await fetch(url, {
    cache: "no-store",
    signal: AbortSignal.timeout(20_000),
  });

  if (!response.ok) {
    throw new Error(`PSD file download failed: ${response.status}`);
  }

  return await response.arrayBuffer();
}

export async function parsePsdTemplateFile(
  buffer: ArrayBuffer
): Promise<ParsedPsdTemplatePayload> {
  const psd = await parsePsdBuffer(buffer);
  const width = Math.max(1, Math.round(Number(psd.width || 0)));
  const height = Math.max(1, Math.round(Number(psd.height || 0)));
  const previewRgba = await psd.composite();
  const previewBuffer = await rgbaToPngBuffer(previewRgba, width, height);
  const layers = collectPsdLayerEntries(Array.isArray(psd.children) ? psd.children : []).map(
    (entry) => entry.layer
  );

  return {
    width,
    height,
    previewBuffer,
    layers,
  };
}

async function createTransparentCanvasBuffer(width: number, height: number) {
  const sharp = await getSharp();
  const canvasWidth = Math.max(1, Math.round(Number(width || 0)));
  const canvasHeight = Math.max(1, Math.round(Number(height || 0)));
  return sharp({
    create: {
      width: canvasWidth,
      height: canvasHeight,
      channels: 4,
      background: {
        r: 0,
        g: 0,
        b: 0,
        alpha: 0,
      },
    },
  })
    .png()
    .toBuffer();
}

export async function preparePsdTemplateRenderAssets(params: {
  templateId: string;
  fileUrl: string;
  layers: AgentPsdTemplateLayerSlot[];
}) {
  const psdBuffer = await fetchRemoteBuffer(params.fileUrl);
  const psd = await parsePsdBuffer(psdBuffer);
  const layerEntries = collectPsdLayerEntries(Array.isArray(psd.children) ? psd.children : []);
  const nodeMap = new Map(layerEntries.map((entry) => [entry.layer.layerId, entry]));
  const editableLayerIdSet = new Set(
    params.layers.filter((layer) => layer.isEditable).map((layer) => layer.layerId)
  );

  const nextLayers = [...params.layers];

  for (const targetLayer of nextLayers) {
    if (!isAgentPsdImageLayerSlot(targetLayer) || !targetLayer.isEditable || targetLayer.maskUrl) {
      continue;
    }

    const entry = nodeMap.get(targetLayer.layerId);
    if (!entry?.node || typeof entry.node.composite !== "function") {
      continue;
    }

    try {
      const rgba = await entry.node.composite(false, true);
      const pixelCount = entry.layer.width * entry.layer.height;
      const maskRgba = new Uint8ClampedArray(pixelCount * 4);
      let hasVisiblePixel = false;

      for (let index = 0; index < pixelCount; index += 1) {
        const alpha = rgba[index * 4 + 3] || 0;
        maskRgba[index * 4] = 255;
        maskRgba[index * 4 + 1] = 255;
        maskRgba[index * 4 + 2] = 255;
        maskRgba[index * 4 + 3] = alpha;
        if (alpha > 0) {
          hasVisiblePixel = true;
        }
      }

      if (!hasVisiblePixel) {
        continue;
      }

      const maskBuffer = await rgbaToPngBuffer(maskRgba, entry.layer.width, entry.layer.height);
      const uploaded = await uploadFile(
        `agent-psd-mask-${params.templateId}-${sanitizeFileSegment(targetLayer.layerId, targetLayer.id)}.png`,
        maskBuffer.buffer.slice(
          maskBuffer.byteOffset,
          maskBuffer.byteOffset + maskBuffer.byteLength
        ) as ArrayBuffer
      );

      const layerIndex = nextLayers.findIndex((layer) => layer.id === targetLayer.id);
      if (layerIndex >= 0 && isAgentPsdImageLayerSlot(nextLayers[layerIndex])) {
        nextLayers[layerIndex] = {
          ...nextLayers[layerIndex],
          maskUrl: uploaded.url,
          maskPath: uploaded.pathname,
        };
      }
    } catch (error) {
      console.warn("[PSD Template] Failed to generate layer mask:", {
        layerId: targetLayer.layerId,
        error,
      });
    }
  }

  const baseCompositeInputs: Array<{ input: Buffer; left: number; top: number }> = [];

  for (const entry of layerEntries) {
    if (editableLayerIdSet.has(entry.layer.layerId)) {
      continue;
    }

    if (!entry.node || typeof entry.node.composite !== "function") {
      continue;
    }

    try {
      const rgba = await entry.node.composite();
      const overlayBuffer = await rgbaToPngBuffer(rgba, entry.layer.width, entry.layer.height);
      baseCompositeInputs.push({
        input: overlayBuffer,
        left: entry.layer.x,
        top: entry.layer.y,
      });
    } catch (error) {
      console.warn("[PSD Template] Failed to compose base layer:", {
        layerId: entry.layer.layerId,
        error,
      });
    }
  }

  const sharp = await getSharp();
  const transparentBaseBuffer = await createTransparentCanvasBuffer(psd.width, psd.height);
  const renderBaseBuffer = await sharp(transparentBaseBuffer)
    .composite(baseCompositeInputs)
    .png()
    .toBuffer();

  return {
    layers: nextLayers,
    renderBaseBuffer,
  };
}