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

export const FIBO_EXPAND_MODEL = "bria/fibo-edit-1.5/expand";
export const FIBO_EXPAND_PROMPT = "Continue the same scene edge to edge, matching the existing lighting, colors, texture, perspective and cut lines. Preserve the original subject and all lettering; continue any caption background that reaches the side edges.";

export function getFiboGeometry(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, Math.sqrt(24_000_000 / (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);
  return {
    targetWidth, targetHeight, left, top, modelWidth, modelHeight, sourceWidth, sourceHeight,
    modelLeft: Math.round(left * modelWidth / targetWidth),
    modelTop: Math.round(top * modelHeight / targetHeight),
  };
}

export async function composeFiboExpansion(source: Buffer, expanded: Buffer, geometry: ReturnType<typeof getFiboGeometry>) {
  const { targetWidth, targetHeight, left, top, modelWidth, modelHeight } = geometry;
  const modelImage = sharp(expanded);
  const info = await modelImage.metadata();
  if (info.width !== modelWidth || info.height !== modelHeight) {
    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.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 expandWithFibo({ 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 = getFiboGeometry(info.width, info.height, xScale, yScale);
  if (geometry.targetWidth * geometry.targetHeight > 64_000_000) {
    throw new Error("扩图后的尺寸过大，请降低扩展倍数");
  }
  const fal = createFalClient({ credentials: apiKey });
  const image = await fal.storage.upload(new Blob([new Uint8Array(source)], { type: "image/png" }));
  const { data } = await fal.subscribe(FIBO_EXPAND_MODEL, {
    input: {
      image_url: image,
      canvas_size: [geometry.modelWidth, geometry.modelHeight],
      original_image_size: [geometry.sourceWidth, geometry.sourceHeight],
      original_image_location: [geometry.modelLeft, geometry.modelTop],
      retain_original_quality: true,
      prompt: FIBO_EXPAND_PROMPT,
    },
    timeout: 90_000,
  });
  const outputUrl = (data as { image?: { url?: string } }).image?.url;
  if (!outputUrl) throw new Error("扩图服务未返回图片");
  const output = await fetch(outputUrl, { signal: AbortSignal.timeout(25_000) });
  if (!output.ok) throw new Error("无法读取扩图结果");
  return composeFiboExpansion(source, Buffer.from(await output.arrayBuffer()), geometry);
}