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

const liveryInputDirectory = "images/liveries_unprocessed"
const liveryOutputDirectory = "images/liveries_processed"
const liveryBackgroundRemovedDirectory = "images/liveries_background_removed"


import Replicate from "replicate";
import sharp from "sharp";

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

    for (const file of pngFiles) {
        const filePath = path.join(outputDirectory, file);
        const fileBuffer = fs.readFileSync(filePath);
        const metadata = await sharp(fileBuffer).metadata();

        if (metadata.format !== 'png') {
            console.log(`Converting to PNG: ${filePath}`);
            const pngData = await sharp(fileBuffer).png().toBuffer();
            fs.writeFileSync(filePath, pngData);
            console.log(`Converted: ${filePath}`);
        }
    }
};


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

        fs.mkdirSync(path.dirname(outputPath), { recursive: true });

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

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 generateLiveries = async () => {
    const replicate = new Replicate({ auth: process.env.REPLICATE_API_TOKEN });

    const prompt = "Generate a 2D, side-view digital technical illustration of the identical plane pictured - with no wheels (MAKE SURE TO REMOVE ANY WHEELS). (correct the perspective of the plane so it is side view) (in Norebbo's style, preserving the livery markings exactly), with a plain white background (no shadows) and with some mild lighting and sheen effects, correcting perspective if the plane is not side-on. Make sure to remove any wheels if they are visible. DO NOT ADD ANYTHING TO THE LIVERY THAT IS NOT IN THE ORIGINAL PHOTO";

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

    const imageFiles = fs
        .readdirSync(liveryInputDirectory)
        .filter((file) => /\.(png|jpg|jpeg|webp)$/i.test(file));

    for (const file of imageFiles) {
        const filePath = path.join(liveryInputDirectory, file);
        const ext = path.extname(file).toLowerCase();
        const outputFile = ext !== ".png" ? file.replace(/\.[^.]+$/, ".png") : file;
        const outputPath = path.join(liveryOutputDirectory, outputFile);

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

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

        const mimeType = ext === ".jpg" || ext === ".jpeg" ? "image/jpeg"
            : ext === ".webp" ? "image/webp"
                : "image/png";

        const fileBuffer = fs.readFileSync(filePath);
        const blob = new Blob([fileBuffer], { type: mimeType });
        const uploadedFile = await replicate.files.create(blob, {
            filename: file,
            contentType: mimeType,
        });

        console.log(`Uploaded: ${uploadedFile.urls.get}`);

        const output = await replicate.run("google/gemini-2.5-flash-image", {
            input: {
                prompt,
                image_input: [uploadedFile.urls.get],
                aspect_ratio: "16:9",
                output_format: "png",
            },
        }) as ReadableStream;

        const reader = output.getReader();
        const chunks: Uint8Array[] = [];
        while (true) {
            const { done, value } = await reader.read();
            if (done) break;
            chunks.push(value);
        }
        const resultBuffer = Buffer.concat(chunks.map(c => Buffer.from(c)));
        fs.writeFileSync(outputPath, resultBuffer);

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


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

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