import { query } from './lib/mysql';
import { generateImage } from './lib/image-generator';
import { logger } from './lib/logger';
import dotenv from 'dotenv';

dotenv.config();

interface VideoRecord {
    id: number;
    genre: string;
    sub_genre: string;
    image_prompt: string;
    // Add other fields if necessary
}

const servers = [
    {
        baseUrl: process.env.SERVER1_COMFY_BASE_URL,
        outputDir: process.env.SERVER1_COMFY_OUTPUT_DIR,
    },

    /*
    {
        baseUrl: process.env.SERVER2_COMFY_BASE_URL,
        outputDir: process.env.SERVER2_COMFY_OUTPUT_DIR,
    },
    */
];

const sleep = (ms: number) => new Promise(resolve => setTimeout(resolve, ms));

async function processVideo(server: any, video: any) {
    const MAX_RETRIES = 3;
    for (let i = 0; i < MAX_RETRIES; i++) {
        try {
            const fileName = `${video.id}_${video.genre}_${video.sub_genre}.png`.replace(/\s/g, '_');
            const imagePath = await generateImage(
                video.image_prompt,
                fileName,
                server.baseUrl,
                server.outputDir,
                'flux',
                { width: 720, height: 1280 }
            );

            await query('UPDATE video SET image_path = ? WHERE id = ?', [imagePath, video.id]);
            logger.info(`Generated and saved image for video ${video.id} at ${imagePath}`);
            return; // Success
        } catch (error) {
            logger.error(`Failed to generate image for video ${video.id} on server ${server.baseUrl}:`, error);
            if (i < MAX_RETRIES - 1) {
                logger.info(`Retrying in 1 minute... (${i + 1}/${MAX_RETRIES})`);
                await sleep(60000);
            } else {
                logger.error(`All retries failed for video ${video.id} on server ${server.baseUrl}.`);
                await query("UPDATE video SET image_path = '' WHERE id = ?", [video.id]);
            }
        }
    }
}

async function worker(server: any) {
    while (true) {
        await sleep(Math.random() * 3000); // Random delay
        const videosToProcess = (await query(
            "SELECT * FROM video WHERE image_prompt IS NOT NULL AND (image_path IS NULL OR image_path = '') LIMIT 1"
        )) as any[];

        if (videosToProcess.length === 0) {
            logger.info('No more images to generate.');
            await sleep(10000); // Wait for 10 seconds before checking again
            continue;
        }

        const video = videosToProcess[0];
        // Mark the video as being processed to avoid other workers picking it up
        await query("UPDATE video SET image_path = 'processing' WHERE id = ?", [video.id]);

        await processVideo(server, video);
    }
}

async function main() {
    try {
        const promises = servers.map(server => {
            if (!server.baseUrl || !server.outputDir) {
                logger.warn(`Server is not configured. Skipping.`);
                return Promise.resolve();
            }
            return worker(server);
        });

        await Promise.all(promises);
        logger.info('Finished generating all images.');
    } catch (error) {
        logger.error('An error occurred during image generation:', error);
    } finally {
        process.exit();
    }
}

main();
