import { query } from './lib/mysql';
import { generateVideo } from './lib/video-generator';
import { logger } from './lib/logger';
import dotenv from 'dotenv';
import path from 'path';
import fs from 'fs/promises';

dotenv.config();

interface VideoRecord {
    id: number;
    genre: string;
    sub_genre: string;
    video_prompt: string;
    image_path: string;
}

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: VideoRecord) {
    const MAX_RETRIES = 3;
    for (let i = 0; i < MAX_RETRIES; i++) {
        try {
            const imageFileName = path.basename(video.image_path);
            const generatedPath = path.resolve(__dirname, '..', 'generated');
            const localImagePath = path.join(generatedPath, imageFileName);

            try {
                await fs.access(localImagePath);
            } catch (error) {
                logger.error(`Image file not found for video ${video.id} at ${localImagePath}. Skipping.`);
                await query("UPDATE video SET video_path = 'image not found' WHERE id = ?", [video.id]);
                return;
            }

            const serverImagePath = path.join(server.outputDir.replace('output', 'input'), imageFileName);
            await fs.copyFile(localImagePath, serverImagePath);

            const videoFileName = `${video.id}_${video.genre}_${video.sub_genre}.mp4`.replace(/\s/g, '_');
            const videoPath = await generateVideo(
                video.video_prompt,
                imageFileName,
                videoFileName,
                server.baseUrl,
                server.outputDir,
                { width: 720, height: 1280 }
            );

            const absolutePath = path.resolve(videoPath);
            await query('UPDATE video SET video_path = ? WHERE id = ?', [absolutePath, video.id]);
            logger.info(`Generated and saved video for video ${video.id} at ${absolutePath}`);
            return; // Success
        } catch (error) {
            logger.error(`Failed to generate video 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 video_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 video_prompt IS NOT NULL 
                    AND image_path IS NOT NULL
                    AND image_path != 'processing' 
                    AND (video_path IS NULL OR video_path = '') 
                ORDER BY RAND() LIMIT 1`
        )) as VideoRecord[];

        if (videosToProcess.length === 0) {
            logger.info('No more videos 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 video_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 videos.');
    } catch (error) {
        logger.error('An error occurred during video generation:', error);
    } finally {
        process.exit();
    }
}

main();
