import { createFalClient } from "@fal-ai/client";
import sharp from "sharp";

export const FAL_OUTPAINT_MODEL = "fal-ai/flux-2-pro/outpaint";

export function getOutpaintGeometry(width: number, height: number, xScale: number, yScale: number) {
  const targetWidth = Math.round(width * xScale);
  const targetHeight = Math.round(height * yScale);
  const ratio = Math.min(1, 2048 / Math.max(targetWidth, targetHeight));
  const modelWidth = Math.max(1, Math.floor(targetWidth * ratio));
  const modelHeight = Math.max(1, Math.floor(targetHeight * ratio));
  const sourceWidth = Math.max(1, Math.round(width * modelWidth / targetWidth));
  const sourceHeight = Math.max(1, Math.round(height * modelHeight / targetHeight));
  const left = Math.floor((targetWidth - width) / 2);
  const top = Math.floor((targetHeight - height) / 2);
  const modelLeft = Math.round(left * modelWidth / targetWidth);
  const modelTop = Math.round(top * modelHeight / targetHeight);
  // Align the model canvas to its 16-pixel grid, then remove the extra border.
  const extraWidth = (16 - modelWidth % 16) % 16;
  const extraHeight = (16 - modelHeight % 16) % 16;
  const cropLeft = Math.floor(extraWidth / 2);
  const cropTop = Math.floor(extraHeight / 2);
  return {
    targetWidth, targetHeight, left, top, modelWidth, modelHeight, sourceWidth, sourceHeight,
    paddedWidth: modelWidth + extraWidth,
    paddedHeight: modelHeight + extraHeight,
    cropLeft,
    cropTop,
    expandLeft: modelLeft + cropLeft,
    expandRight: modelWidth - sourceWidth - modelLeft + extraWidth - cropLeft,
    expandTop: modelTop + cropTop,
    expandBottom: modelHeight - sourceHeight - modelTop + extraHeight - cropTop,
  };
}

export async function composeOutpaint(source: Buffer, expanded: Buffer, geometry: ReturnType<typeof getOutpaintGeometry>) {
  const { targetWidth, targetHeight, left, top, modelWidth, modelHeight, paddedWidth, paddedHeight, cropLeft, cropTop } = geometry;
  const modelImage = sharp(expanded);
  const info = await modelImage.metadata();
  if (info.width !== paddedWidth || info.height !== paddedHeight) {
    throw new Error("扩图服务返回的图片尺寸不正确");
  }

  // Copy the source rectangle, including transparent pixels, without blending or resizing it.
  const { data: original, info: sourceInfo } = await sharp(source).toColourspace("srgb").ensureAlpha().raw().toBuffer({ resolveWithObject: true });
  const result = await modelImage
    .extract({ left: cropLeft, top: cropTop, width: modelWidth, height: modelHeight })
    .resize(targetWidth, targetHeight, { fit: "fill" })
    .toColourspace("srgb").ensureAlpha().raw().toBuffer();
  const stride = sourceInfo.width * 4;
  for (let row = 0; row < sourceInfo.height; row++) {
    original.copy(result, ((top + row) * targetWidth + left) * 4, row * stride, (row + 1) * stride);
  }
  return sharp(result, { raw: { width: targetWidth, height: targetHeight, channels: 4 } }).png().toBuffer();
}

export async function expandWithFalOutpaint({ imageUrl, apiKey, xScale, yScale }: {
  imageUrl: string;
  apiKey: string;
  xScale: number;
  yScale: number;
}) {
  const response = await fetch(imageUrl, { signal: AbortSignal.timeout(45_000) });
  if (!response.ok) throw new Error("无法读取待扩展的原图");
  const { data: source, info } = await sharp(Buffer.from(await response.arrayBuffer()))
    .rotate().png().toBuffer({ resolveWithObject: true });
  const geometry = getOutpaintGeometry(info.width, info.height, xScale, yScale);
  if (geometry.targetWidth * geometry.targetHeight > 64_000_000) {
    throw new Error("扩图后的尺寸过大，请降低扩展倍数");
  }
  const fal = createFalClient({ credentials: apiKey });
  const modelSource = geometry.sourceWidth === info.width && geometry.sourceHeight === info.height
    ? source
    : await sharp(source).resize(geometry.sourceWidth, geometry.sourceHeight, { fit: "fill" }).png().toBuffer();
  const image = await fal.storage.upload(new Blob([new Uint8Array(modelSource)], { type: "image/png" }));
  const { data } = await fal.subscribe(FAL_OUTPAINT_MODEL, {
    input: {
      image_url: image,
      expand_left: geometry.expandLeft,
      expand_right: geometry.expandRight,
      expand_top: geometry.expandTop,
      expand_bottom: geometry.expandBottom,
      mode: "high",
      output_format: "png",
    },
    timeout: 90_000,
  });
  const outputUrl = (data as { images?: Array<{ url?: string }> }).images?.[0]?.url;
  if (!outputUrl) throw new Error("扩图服务未返回图片");
  const output = await fetch(outputUrl, { signal: AbortSignal.timeout(25_000) });
  if (!output.ok) throw new Error("无法读取扩图结果");
  return composeOutpaint(source, Buffer.from(await output.arrayBuffer()), geometry);
}