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

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

const removeWatermark = async () => {

    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)) {
            console.log(`Converting: ${file} (webp → png)`);
            const webpData = fs.readFileSync(outputPath);
            const pngData = await sharp(webpData).png().toBuffer();
            fs.writeFileSync(outputPath, pngData);
            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}`);
    }
}

const removeBackground = async () => {
    fs.mkdirSync(backgroundRemovedDirectory, { recursive: true });

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

    // Load mask once
    const { data: maskData } = await sharp("images/Mask.png")
        .greyscale()
        .raw()
        .toBuffer({ resolveWithObject: true });

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

        if (fs.existsSync(outputPath)) {
            console.log(`Skipping: ${file} (already processed)`);
            continue;
        }

        console.log(`Processing: ${filePath}`);

        const { data, info } = await sharp(filePath)
            .ensureAlpha()
            .raw()
            .toBuffer({ resolveWithObject: true });

        for (let i = 0; i < info.width * info.height; i++) {
            data[i * 4 + 3] = maskData[i];
        }

        await sharp(data, {
            raw: { width: info.width, height: info.height, channels: 4 }
        }).png().resize(256, 256).toFile(outputPath);

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

const commands: Record<string, () => Promise<void>> = {
    removeWatermark,
    removeBackground,
};

const command = process.argv[2];

if (!command || !commands[command]) {
    console.error(`Usage: npx tsx process-images.ts <command>`);
    console.error(`Available commands: ${Object.keys(commands).join(", ")}`);
    process.exit(1);
}

await commands[command]();