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);

        // Ensure the output directory exists before every write
        fs.mkdirSync(path.dirname(outputPath), { recursive: true });

        if (fs.existsSync(outputPath)) {
            const fileBuffer = fs.readFileSync(outputPath);
            const metadata = await sharp(fileBuffer).metadata();

            if (metadata.format !== 'png') {
                console.log(`About to write to: ${outputPath}`);
                console.log(`Dir exists: ${fs.existsSync(path.dirname(outputPath))}`);
                console.log(`Dir writable:`, (() => { try { fs.accessSync(path.dirname(outputPath), fs.constants.W_OK); return true; } catch { return false; } })());
                console.log(`File writable:`, (() => { try { fs.accessSync(outputPath, fs.constants.W_OK); return true; } catch { return false; } })());
                const pngData = await sharp(fileBuffer).png().toBuffer();
                console.log(`PNG buffer size: ${pngData.length}`);
                fs.writeFileSync(outputPath, pngData);
                console.log(`Write succeeded`);
            }

            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)) {
            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]();