at-mintel/packages/payload-ai/src/endpoints/replicateMediaEndpoint.ts

import type { PayloadRequest, PayloadHandler } from "payload";
import Replicate from "replicate";

type Action = "upscale" | "recover";

const replicate = new Replicate({
    auth: process.env.REPLICATE_API_KEY,
});

/**
 * Downloads a remote URL and returns a Buffer.
 */
async function downloadImage(url: string): Promise<Buffer> {
    const res = await fetch(url);
    if (!res.ok) throw new Error(`Failed to download image: ${res.status} ${res.statusText}`);
    const arrayBuffer = await res.arrayBuffer();
    return Buffer.from(arrayBuffer);
}

/**
 * Resolves the public URL for a media document.
 * Handles both S3 and local static files.
 */
function resolveMediaUrl(doc: any): string | null {
    // S3 storage sets `url` directly
    if (doc.url) return doc.url;

    // Local static files: build from NEXT_PUBLIC_BASE_URL + /media/<filename>
    const base = process.env.NEXT_PUBLIC_BASE_URL || "http://localhost:3000";
    if (doc.filename) return `${base}/media/${doc.filename}`;

    return null;
}

export const replicateMediaHandler: PayloadHandler = async (
    req: PayloadRequest,
) => {
    const { id } = req.routeParams as { id: string };
    const payload = req.payload;

    // Parse action from request body
    let action: Action;
    try {
        const body = await req.json?.();
        action = body?.action as Action;
    } catch {
        return Response.json({ error: "Invalid request body" }, { status: 400 });
    }

    if (action !== "upscale" && action !== "recover") {
        return Response.json(
            { error: "Invalid action. Must be 'upscale' or 'recover'." },
            { status: 400 },
        );
    }

    // Fetch the media document
    let mediaDoc: any;
    try {
        mediaDoc = await payload.findByID({ collection: "media", id });
    } catch {
        return Response.json({ error: "Media not found" }, { status: 404 });
    }

    if (!mediaDoc) {
        return Response.json({ error: "Media not found" }, { status: 404 });
    }

    // Check that it's an image
    const mimeType: string = mediaDoc.mimeType || "";
    if (!mimeType.startsWith("image/")) {
        return Response.json(
            { error: "This media file is not an image and cannot be AI-processed." },
            { status: 422 },
        );
    }

    const imageUrl = resolveMediaUrl(mediaDoc);
    if (!imageUrl) {
        return Response.json(
            { error: "Could not resolve a public URL for this media file." },
            { status: 422 },
        );
    }

    // --- Run Replicate ---
    let outputUrl: string;

    try {
        if (action === "upscale") {
            console.log(`[AI Media] Starting upscale for media ${id} – ${imageUrl}`);
            const output = await replicate.run("google/upscaler", {
                input: {
                    image: imageUrl,
                },
            });
            // google/upscaler returns a string URL
            outputUrl = typeof output === "string" ? output : (output as any)?.url ?? String(output);
        } else {
            // recover
            console.log(`[AI Media] Starting photo recovery for media ${id} – ${imageUrl}`);
            const output = await replicate.run(
                "microsoft/bringing-old-photos-back-to-life",
                {
                    input: {
                        image: imageUrl,
                        HR: true,
                    },
                },
            );
            // returns a FileOutput or URL string
            outputUrl = typeof output === "string" ? output : (output as any)?.url ?? String(output);
        }
    } catch (err: any) {
        console.error("[AI Media] Replicate error:", err);
        return Response.json(
            { error: err?.message ?? "Replicate API call failed" },
            { status: 500 },
        );
    }

    // --- Download and re-upload as new media document ---
    let imageBuffer: Buffer;
    try {
        imageBuffer = await downloadImage(outputUrl);
    } catch (err: any) {
        console.error("[AI Media] Download error:", err);
        return Response.json(
            { error: `Failed to download result: ${err?.message}` },
            { status: 500 },
        );
    }

    const suffix = action === "upscale" ? "_upscaled" : "_recovered";
    const originalName: string = mediaDoc.filename || "image.jpg";
    const ext = originalName.includes(".") ? `.${originalName.split(".").pop()}` : ".jpg";
    const baseName = originalName.includes(".")
        ? originalName.slice(0, originalName.lastIndexOf("."))
        : originalName;
    const newFilename = `${baseName}${suffix}${ext}`;
    const originalAlt: string = mediaDoc.alt || originalName;

    let newMedia: any;
    try {
        newMedia = await payload.create({
            collection: "media",
            data: {
                alt: `${originalAlt}${suffix}`,
            },
            file: {
                data: imageBuffer,
                mimetype: mimeType,
                name: newFilename,
                size: imageBuffer.byteLength,
            },
        });
    } catch (err: any) {
        console.error("[AI Media] Upload error:", err);
        return Response.json(
            { error: `Failed to save result: ${err?.message}` },
            { status: 500 },
        );
    }

    console.log(
        `[AI Media] ${action} complete – new media ID: ${newMedia.id}`,
    );

    return Response.json(
        {
            message: `AI ${action} successful. New media document created.`,
            mediaId: newMedia.id,
            url: resolveMediaUrl(newMedia),
        },
        { status: 200 },
    );
};