import * as fs from 'fs';
import * as path from 'path';
import { convertImageWithFile } from '../../lib/image-converter';
import dotenv from 'dotenv';

dotenv.config();

const girlDir = path.join(__dirname, '../../../input/girl');
const monsterDir = path.join(__dirname, '../../../input/monster');
const outputDir = path.join(__dirname, '../../../generated');
const prompt = "只提取图1中的女生，把她放在图2的怪物之间。";

if (!fs.existsSync(outputDir)) {
    fs.mkdirSync(outputDir, { recursive: true });
}

const comfyBaseUrl = process.env.SERVER2_COMFY_BASE_URL;
const comfyOutputDir = process.env.SERVER2_COMFY_OUTPUT_DIR;

if (!comfyBaseUrl || !comfyOutputDir) {
    console.error("Please define SERVER1_COMFY_BASE_URL and SERVER1_COMFY_OUTPUT_DIR in your .env file");
    process.exit(1);
}

const comfyInputDir = comfyOutputDir.replace("output", "input");
if (!fs.existsSync(comfyInputDir)) {
    fs.mkdirSync(comfyInputDir, { recursive: true });
}

async function combineImages() {
    while (true) {
        try {
            const girlImages = fs.readdirSync(girlDir).filter(file => /\.(jpg|jpeg|png)$/i.test(file));
            const monsterImages = fs.readdirSync(monsterDir).filter(file => /\.(jpg|jpeg|png)$/i.test(file));

            if (girlImages.length === 0 || monsterImages.length === 0) {
                console.log('Input directories are empty. Waiting...');
                await new Promise(resolve => setTimeout(resolve, 5000));
                continue;
            }

            const randomGirlImage = girlImages[Math.floor(Math.random() * girlImages.length)];
            const randomMonsterImage = monsterImages[Math.floor(Math.random() * monsterImages.length)];

            const image1Path = path.join(girlDir, randomGirlImage);
            const image2Path = path.join(monsterDir, randomMonsterImage);

            // Copy files to comfy input directory
            const destImage1Path = path.join(comfyInputDir, randomGirlImage);
            const destImage2Path = path.join(comfyInputDir, randomMonsterImage);
            fs.copyFileSync(image1Path, destImage1Path);
            fs.copyFileSync(image2Path, destImage2Path);

            console.log(`Combining ${randomGirlImage} and ${randomMonsterImage}`);

            const generatedFilePath = await convertImageWithFile(prompt, randomGirlImage, randomMonsterImage, comfyBaseUrl!, comfyOutputDir!);

            if (generatedFilePath && fs.existsSync(generatedFilePath)) {
                const timestamp = new Date().getTime();
                const newFileName = `combined_${timestamp}.png`;
                const newFilePath = path.join(outputDir, newFileName);

                fs.renameSync(generatedFilePath, newFilePath);
                console.log(`Renamed generated file to ${newFilePath}`);
            } else {
                console.log("Failed to generate or find the image file.");
            }

        } catch (error) {
            console.error('An error occurred:', error);
        }

        // Wait for a bit before the next iteration
        await new Promise(resolve => setTimeout(resolve, 5000));
    }
}

combineImages();
