import * as fs from 'fs/promises';
import * as path from 'path';
import axios from 'axios';
import dotenv from 'dotenv';

dotenv.config();

interface VideoSize {
    width: number;
    height: number;
}

async function generateStyledVideo(
    imagePrompt: string,
    videoPrompt: string,
    imageName: string,
    newFileName: string,
    comfyBaseUrl: string,
    comfyOutputDir: string,
    size: VideoSize = { width: 320, height: 640 }
): Promise<string> {
    const COMFY_BASE_URL = comfyBaseUrl.replace(/\/$/, '');
    const COMFY_OUTPUT_DIR = comfyOutputDir;
    const workflow = JSON.parse(await fs.readFile('src/comfyworkflows/prototyping_style_flux_wan22.json', 'utf-8'));

    // Set prompts
    workflow['95']['inputs']['text'] = imagePrompt;
    workflow['105']['inputs']['text'] = videoPrompt;

    // Set image name
    workflow['114']['inputs']['image'] = imageName;

    // Set sizes
    workflow['104']['inputs']['width'] = size.width;
    workflow['104']['inputs']['height'] = size.height;
    workflow['97']['inputs']['width'] = size.width;
    workflow['97']['inputs']['height'] = size.height;

    const response = await axios.post(`${COMFY_BASE_URL}/prompt`, { prompt: workflow });
    const promptId = response.data.prompt_id;

    let history;
    do {
        await new Promise(resolve => setTimeout(resolve, 1000));
        const historyResponse = await axios.get(`${COMFY_BASE_URL}/history/${promptId}`);
        history = historyResponse.data[promptId];
    } while (!history || Object.keys(history.outputs).length === 0);

    const files = await fs.readdir(COMFY_OUTPUT_DIR!);

    // Find the latest MP4 file
    const generatedVideos = files.filter(file => file.endsWith('.mp4'));
    const videoStats = await Promise.all(
        generatedVideos.map(async (file) => {
            const stat = await fs.stat(path.join(COMFY_OUTPUT_DIR!, file));
            return { file, mtime: stat.mtime };
        })
    );
    videoStats.sort((a, b) => b.mtime.getTime() - a.mtime.getTime());
    const latestVideo = videoStats[0].file;

    // Find the latest PNG file
    const generatedImages = files.filter(file => file.endsWith('.png'));
    const imageStats = await Promise.all(
        generatedImages.map(async (file) => {
            const stat = await fs.stat(path.join(COMFY_OUTPUT_DIR!, file));
            return { file, mtime: stat.mtime };
        })
    );
    imageStats.sort((a, b) => b.mtime.getTime() - a.mtime.getTime());
    const latestImage = imageStats[0].file;

    // Create new file paths
    const newVideoPath = path.resolve('./generated', newFileName);
    const newImagePath = path.resolve('./generated', newFileName.replace('.mp4', '.png'));

    await fs.mkdir('./generated', { recursive: true });

    // Copy both files
    const sourceVideoPath = path.join(COMFY_OUTPUT_DIR!, latestVideo);
    await fs.copyFile(sourceVideoPath, newVideoPath);

    const sourceImagePath = path.join(COMFY_OUTPUT_DIR!, latestImage);
    await fs.copyFile(sourceImagePath, newImagePath);

    // Optionally, unlink the source files
    // await fs.unlink(sourceVideoPath);
    // await fs.unlink(sourceImagePath);

    return newVideoPath;
}

export { generateStyledVideo };