import * as fs from 'fs/promises';
import * as path from 'path';
import dotenv from 'dotenv';
import { readJsonToPng, embedJsonToPng } from '../lib/util';
import { convertImage } from '../lib/image-converter';

dotenv.config();

const inputDir = './generated/prompts';
const outputDir = './generated/image';

const COMFY_BASE_URL = process.env.SERVER1_COMFY_BASE_URL!;
const COMFY_OUTPUT_DIR = process.env.SERVER1_COMFY_OUTPUT_DIR!;

interface PngMetadata {
    prompts: {
        imagePrompt: string;
        videoPrompt: string;
    }[];
}

async function main() {
    await fs.mkdir(outputDir, { recursive: true });

    const files = await fs.readdir(inputDir);
    let generatedImageIndex = 0;

    for (const file of files) {
        if (path.extname(file).toLowerCase() !== '.png') {
            continue;
        }

        const inputFile = path.join(inputDir, file);
        const metadata = await readJsonToPng(inputFile) as PngMetadata;

        if (metadata && metadata.prompts && Array.isArray(metadata.prompts)) {
            console.log(`Processing ${file} with ${metadata.prompts.length} prompt pairs.`);

            const inputfolderFullpath = COMFY_OUTPUT_DIR.replace("output", "input");
            await fs.copyFile(inputFile, path.join(inputfolderFullpath, file));

            for (const promptPair of metadata.prompts) {
                const { imagePrompt, videoPrompt } = promptPair;
                const newFileName = `cleaned_prompt_generated_${generatedImageIndex}.png`;
                generatedImageIndex++;
                const outputPath = path.join(outputDir, newFileName);

                try {
                    await fs.access(outputPath);
                    console.log(`File ${newFileName} already exists, skipping.`);
                    continue;
                } catch (error) {
                    // File does not exist, proceed with generation
                }

                console.log(`Generating image for prompt: "${imagePrompt}"`);

                try {
                    const generatedFilePath = await convertImage(
                        imagePrompt,
                        file, // Using the same image for both inputs as per interpretation
                        COMFY_BASE_URL,
                        COMFY_OUTPUT_DIR
                    );

                    // The convertImage function saves the file in a generic location.
                    // We need to move it to the correct location with the correct name.
                    await fs.rename(generatedFilePath, outputPath);

                    const newMetadata = {
                        imagePrompt: imagePrompt,
                        videoPrompt: videoPrompt
                    };

                    await embedJsonToPng(outputPath, newMetadata);

                    console.log(`Successfully generated and saved ${newFileName} with metadata.`);
                } catch (error) {
                    console.error(`Error generating image for prompt "${imagePrompt}":`, error);
                }
            }
        } else {
            console.log(`Skipping ${file}, no valid prompts metadata found.`);
        }
    }
}

main().catch(console.error);
