import fs from "fs";
import path from "path";

const inputDirectory = "images/unprocessed";
const outputDirectory = "images/processed";
const watermarkImage = "images/Watermark.png";
import Replicate from "replicate";

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

const prompt = "Carefully remove the watermark in image 1 from image 2, leaving everything else unchanged. Preserve icon details while removing watermark text that says 'Airhex.com'. Save with a white background."

fs.mkdirSync(outputDirectory, { recursive: true });

const pngFiles = fs
    .readdirSync(inputDirectory)
    .filter((file) => file.toLowerCase().endsWith(".png"))

for (const file of pngFiles) {
    const filePath = path.join(inputDirectory, file);
    const outputPath = path.join(outputDirectory, file);

    if (fs.existsSync(outputPath)) {
        continue;
    }

    console.log(`Processing: ${filePath}`);
    const watermarkData = fs.readFileSync(watermarkImage);

    const imageData = fs.readFileSync(filePath);
    const watermarkBase64 = `data:image/png;base64,${watermarkData.toString("base64")}`;

    const imageBase64 = `data:image/png;base64,${imageData.toString("base64")}`;

    const output = await replicate.run("qwen/qwen-image-edit-plus", {
        input: {
            image: [watermarkBase64, imageBase64],
            prompt,
        },
    }) as { url(): string }[];
    const resultBuffer = Buffer.from(await (await fetch(output[0].url())).arrayBuffer());
    fs.writeFileSync(outputPath, resultBuffer);

    console.log(`Saved: ${outputPath}`);
}