import * as fs from 'fs/promises';
import * as path from 'path';
import dotenv from 'dotenv';
import { callLMStudioAPIWithFile, callLmstudio } from '../lib/lmstudio';
import { convertImage } from '../lib/image-converter';
import { logger } from '../lib/logger';

dotenv.config();

const COMFY_BASE_URL = process.env.SERVER2_COMFY_BASE_URL;
const COMFY_OUTPUT_DIR = process.env.SERVER2_COMFY_OUTPUT_DIR;
const INPUT_DIR = './input';
const GENERATED_DIR = './generated';

async function batchConvert() {
    if (!COMFY_BASE_URL || !COMFY_OUTPUT_DIR) {
        throw new Error('ComfyUI server details are not defined in the .env file');
    }

    try {
        await fs.mkdir(GENERATED_DIR, { recursive: true });
        const files = await fs.readdir(INPUT_DIR);

        for (const file of files) {

            try {
                const inputFile = path.join(INPUT_DIR, file);
                logger.info(`Processing ${inputFile}`);

                const firstPrompt = `
            Read the file in the photo and detect the main subject. Extract the following information:
            - what is the main subject in one word
            - describe character in 5 words
            - describe background in 5 words
            - 3 ideas to make this character's look, 2 or 3 words per idea

            Return the result in this format:
            { "character":"", "characterDescription": "", "background":"", "idea":"idea1, idea2, idea3"}
            `;

                const imageInfo = await callLMStudioAPIWithFile(inputFile, firstPrompt);

                const { character, idea, background } = imageInfo;
                const ideas = idea.split(',').map((i: string) => i.trim());

                const secondPrompt = `
            Generate a prompt to convert the photo to a group dancing photo.
            I need only the prompt in this format,  main subject is ${character}
            {"result":"
            Change the camera angle to a full-body shot, place many ${character} in the scene messily,
            Change lighting to silhouette lighting,
            Change the style for each body using these ideas: ${ideas.join(', ')},
            Change color of each body,
            Change the background in horror movie style of ${background},
            Overall photo should be realistic and spooky,
            "} `;

                const promptResult = await callLmstudio(secondPrompt);
                const finalPrompt = promptResult.result;

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

                logger.info(`Generating image for ${file} with prompt: ${finalPrompt}`);
                const generatedFile = await convertImage(finalPrompt, file, COMFY_BASE_URL, COMFY_OUTPUT_DIR);

                const finalOutputPath = path.join(GENERATED_DIR, path.basename(generatedFile));
                await fs.copyFile(generatedFile, finalOutputPath);
                logger.info(`Saved converted image to ${finalOutputPath}`);
            } catch (e) {
                logger.error('An error occurred during batch conversion:', e);
                continue;
            }

        }
    } catch (error) {
        logger.error('An error occurred during batch conversion:', error);
    }
}

batchConvert();
