YiffyPaintings/scripts/common.ts

import { randomBytes } from "crypto";
import { fileTypeFromBuffer } from "file-type";
import Ffmpeg from "fluent-ffmpeg";
import { rm } from "fs/promises";
import { readdir } from "fs/promises";
import { mkdtemp } from "fs/promises";
import { tmpdir } from "os";
import { basename } from "path";
import sharp, { type OverlayOptions, type Sharp } from "sharp";

export async function formatImage(name: string, image: string | string[], frameColors: number[], framePercent: number, width: number, height: number, resizeWidth?: number, resizeHeight?: number, frames?: [start: number, end: number], saveKz?: boolean) {
    saveKz ??= false; 
    const originalWidth = width, originalHeight = height;
    const originalFrameVH = Math.floor(Math.min(originalWidth, originalHeight) * framePercent);
    const originalInputWidth = originalWidth - (originalFrameVH * 2), originalInputHeight = originalHeight - (originalFrameVH * 2);
    if (resizeWidth && resizeHeight) {
        width = resizeWidth;
        height = resizeHeight;
    }
    const frameVH = Math.floor(Math.min(width, height) * framePercent);
    const inputWidth = width - (frameVH * 2), inputHeight = height - (frameVH * 2);
    const isImage = !Array.isArray(image) && ["image/png", "image/jpeg"].includes((await fileTypeFromBuffer(Buffer.isBuffer(image) ? image : await Bun.file(image).arrayBuffer()))!.mime);
    let result: Sharp, kzFile: string | null | true = null;
    const frame = await createFrame(frameColors, framePercent, width, height);
    if (isImage) {
        result = await sharp(frame)
            .composite([{
                input: await sharp(image).resize(inputWidth, inputHeight, { fit: "fill" }).toBuffer(),
                top: frameVH,
                left: frameVH
            }])
            .resize(width, height, { fit: "fill" });
            kzFile = true as const;
    } else {
        let tmpDir: string | undefined;

        let files: string[] = [];
            tmpDir = await mkdtemp(`${tmpdir()}/split-frames-`);
        if (!Array.isArray(image)) {
            const [start = 0, end = null] = frames ?? [];
            console.debug("Input file %s is not an image, assuming we need to extract frames.", name);
             await new Promise<void>((resolve) => {
                Ffmpeg(image)
                .videoFilter(`scale=${inputWidth}:${inputHeight}${end !== null ? `,select='gte(n\\, ${start})*lte(n\\,${end})'` : ""}`)
                .addOption("-vsync vfr")
                .output(`${tmpDir}/frame%04d.png`)
                .on("end", async() => {
                    files = (await readdir(tmpDir!)).filter(f => /frame\d+\.png/.test(f)).sort().map(f => `${tmpDir}/${f}`);
                    if (files.length === 0) {
                        console.error("No frames extracted for %s.", name);
                        if (tmpDir) {
                            await rm(tmpDir, { recursive: true });
                        }
                        process.exit(1);
                    }
                    console.log("Frames extracted to %s, %d total for %s", tmpDir, files.length, name);
                    resolve();
                })
                .on("error", async() => {
                    if (tmpDir) {
                        await rm(tmpDir, { recursive: true });
                    }
                    console.error("Failed to extract frames for %s.", name);
                    process.exit(1);
                })
                .run()
            });
        } else {
            for (const img of image) {
                await sharp(img).resize(inputWidth, inputHeight, { fit: "fill" }).toFile(`${tmpDir}/${basename(img)}`);
            }
            files = image.map(f => `${tmpDir}/${basename(f)}`).sort();
        }
        
        const parts: OverlayOptions[] = [];
        for (const img of files) {
            const i = files.indexOf(img);
            console.log("Processing %s (%d/%d) for %s", img, i + 1, files.length, name);
            const imgFile = await sharp(img).resize(inputWidth, inputHeight, { fit: "fill" }).toBuffer();
            if (saveKz && !kzFile && (files.length < 5 || i >= 5)) {
            const ogFile = await sharp(img).resize(originalInputWidth, originalInputHeight, { fit: "fill" }).toBuffer();
                kzFile = `${tmpdir()}/${randomBytes(8).toString("hex")}-kz.png`;
                await sharp(await sharp(frame).resize(originalWidth, originalHeight, { fit: "fill" }).toBuffer()).composite([{ input: ogFile, top: originalFrameVH, left: originalFrameVH }]).resize(originalWidth, originalHeight, { fit: "fill" }).toFile(kzFile);
            }
            parts.push(
                { input: frame, top: i * height, left: 0 },
                { input: imgFile, top: (i * height) + frameVH, left: frameVH }
            );
        }

        result = await sharp({
            create: {
                width,
                height: height * files.length,
                channels: 4,
                background: { r: 0, g: 0, b: 0, alpha: 0 }
            },
            limitInputPixels: false
        })
        .composite(parts)
        if (tmpDir) {
            await rm(tmpDir, { recursive: true });
        }
    }
    return { isImage, result, kzFile };
}

const frameSizePercent = 0.1;
async function createFrame(colors: number[], percent: number, width: number, height: number) {
    let frameW = width, frameH = height, i = 0;
    const parts: OverlayOptions[] = [];
    const frameVH = Math.floor(Math.min(width, height) * percent);
    while (frameW > 0) {
        if (i > colors.length - 1) {
            i = 0;
        }
        const color = colors[i], r = (color >> 16) & 0xFF, g = (color >> 8) & 0xFF, b = color & 0xFF,
        w = width * frameSizePercent, h = frameVH;
        parts.push({
            input: {
                create: {
                    width: w,
                    height: h,
                    channels: 4,
                    background: { r, g, b, alpha: 255 }
                }
            },
            top: 0,
            // left to right
            left: i * w
        },
        {
            input: {
                create: {
                    width: w,
                    height: h,
                    channels: 4,
                    background: { r, g, b, alpha: 255 }
                }
            },
            top: frameH - h,
            // right to left
            left: width - (i + 1) * w
        });
        frameW -= w;
        i += 1;
    }
    i = 0;
    colors = colors.toReversed();
    while (frameH > 0) {
        if (i > colors.length - 1) {
            i = 0;
        }
        const color = colors[i], r = (color >> 16) & 0xFF, g = (color >> 8) & 0xFF, b = color & 0xFF,
        w = frameVH, h = height * frameSizePercent;
        parts.push({
            input: {
                create: {
                    width: w,
                    height: h,
                    channels: 4,
                    background: { r, g, b, alpha: 255 }
                }
            },
            top: i * h,
            // top to bottom
            left: 0
        },
        {
            input: {
                create: {
                    width: w,
                    height: h,
                    channels: 4,
                    background: { r, g, b, alpha: 255 }
                }
            },
            top: height - (i + 1) * h,
            // bottom to top
            left: width - w
        });
        frameH -= h;
        i += 1;
    }
    return sharp({
        create: {
            width,
            height,
            channels: 4,
            background: { r: 0, g: 0, b: 0, alpha: 0 }
        },
        limitInputPixels: false
    }).composite(parts).png().toBuffer();
}