import dotenv from 'dotenv';
import path from 'path';
import fs from 'fs/promises';
import { spawn } from 'child_process';

import { logger } from '../lib/logger';
import { generateImage } from '../lib/image-generator';
import { generateVideo } from '../lib/video-generator';

dotenv.config();

type Size = { width: number; height: number };

interface InfinitySceneConfig {
    initialImage: string;
    videoPrompt: string;
}

interface Server {
    baseUrl?: string;
    outputDir?: string;
    inputDir?: string;
    name: string;
}

const DEFAULT_SIZE: Size = { width: 720, height: 1280 };
const GENERATED_DIR = path.resolve('generated');

/**
 * Load ComfyUI servers from env:
 *  - SERVER1_COMFY_BASE_URL, SERVER1_COMFY_OUTPUT_DIR
 *  - SERVER2_COMFY_BASE_URL, SERVER2_COMFY_OUTPUT_DIR
 * (inputDir is inferred by replacing "output" with "input")
 */
function loadServers(): Server[] {
    const servers: Server[] = [
        {
            name: 'SERVER1',
            baseUrl: process.env.SERVER1_COMFY_BASE_URL,
            outputDir: process.env.SERVER1_COMFY_OUTPUT_DIR,
        },
    ]
        .filter((s) => !!s.baseUrl && !!s.outputDir)
        .map((s) => ({
            ...s,
            inputDir: s.outputDir!.replace(/output/i, 'input'),
        }));

    if (servers.length === 0) {
        logger.warn('No servers configured. Please set SERVER{N}_COMFY_BASE_URL and SERVER{N}_COMFY_OUTPUT_DIR in .env');
    } else {
        for (const s of servers) {
            logger.info(`Configured ${s.name}: baseUrl=${s.baseUrl}, outputDir=${s.outputDir}, inputDir=${s.inputDir}`);
        }
    }

    return servers;
}

async function ensureDirs() {
    await fs.mkdir(GENERATED_DIR, { recursive: true });
}

async function copyImageToAllServerInputs(servers: Server[], localGeneratedImagePath: string): Promise<string> {
    const fileName = path.basename(localGeneratedImagePath);
    for (const s of servers) {
        if (!s.inputDir) continue;
        const dest = path.join(s.inputDir, fileName);
        try {
            await fs.mkdir(s.inputDir, { recursive: true });
            await fs.copyFile(localGeneratedImagePath, dest);
            logger.debug(`Copied ${fileName} to ${s.name} input: ${dest}`);
        } catch (err) {
            logger.warn(`Failed to copy ${fileName} to ${s.name} input: ${err}`);
        }
    }
    return fileName; // name used by Comfy workflow inputs
}

function pickServer(servers: Server[], idx: number): Server {
    if (servers.length === 0) {
        throw new Error('No servers configured.');
    }
    return servers[idx % servers.length];
}

async function fileExists(p: string): Promise<boolean> {
    try {
        await fs.access(p);
        return true;
    } catch {
        return false;
    }
}

/** ---------- ffmpeg / ffprobe helpers ---------- */
function runFfmpeg(args: string[], { cwd }: { cwd?: string } = {}): Promise<void> {
    return new Promise((resolve, reject) => {
        const proc = spawn('ffmpeg', args, { cwd, stdio: ['ignore', 'pipe', 'pipe'] });

        let stderr = '';
        proc.stdout.on('data', (d) => logger.debug(d.toString()));
        proc.stderr.on('data', (d) => {
            const msg = d.toString();
            stderr += msg;
            logger.debug(msg); // ffmpeg prints progress to stderr
        });
        proc.on('error', (err: any) => reject(err));
        proc.on('close', (code) => {
            if (code === 0) resolve();
            else reject(new Error(`ffmpeg exited with code ${code}. ${stderr.slice(0, 1000)}`));
        });
    });
}

function runFfprobe(args: string[], { cwd }: { cwd?: string } = {}): Promise<{ stdout: string; stderr: string }> {
    return new Promise((resolve, reject) => {
        const proc = spawn('ffprobe', args, { cwd, stdio: ['ignore', 'pipe', 'pipe'] });
        let stdout = '';
        let stderr = '';
        proc.stdout.on('data', (d) => (stdout += d.toString()));
        proc.stderr.on('data', (d) => (stderr += d.toString()));
        proc.on('error', (err) => reject(err));
        proc.on('close', (code) => {
            if (code === 0) resolve({ stdout, stderr });
            else reject(new Error(`ffprobe exited with code ${code}. ${stderr.slice(0, 1000)}`));
        });
    });
}

/**
 * Try to get precise total frame count using ffprobe.
 * 1) nb_read_frames (requires -count_frames, may be slow but exact)
 * 2) nb_frames (container-level, sometimes missing)
 * Returns undefined if both unavailable.
 */
async function getVideoFrameCount(inputVideoPath: string): Promise<number | undefined> {
    // Attempt 1: nb_read_frames
    try {
        const { stdout } = await runFfprobe([
            '-v', 'error',
            '-count_frames',
            '-select_streams', 'v:0',
            '-show_entries', 'stream=nb_read_frames',
            '-of', 'default=nokey=1:noprint_wrappers=1',
            inputVideoPath,
        ]);
        const n = parseInt(stdout.trim(), 10);
        if (Number.isFinite(n) && n > 0) return n;
    } catch (e) {
        logger.debug(`ffprobe nb_read_frames failed: ${e instanceof Error ? e.message : e}`);
    }
    // Attempt 2: nb_frames
    try {
        const { stdout } = await runFfprobe([
            '-v', 'error',
            '-select_streams', 'v:0',
            '-show_entries', 'stream=nb_frames',
            '-of', 'default=nokey=1:noprint_wrappers=1',
            inputVideoPath,
        ]);
        const n = parseInt(stdout.trim(), 10);
        if (Number.isFinite(n) && n > 0) return n;
    } catch (e) {
        logger.debug(`ffprobe nb_frames failed: ${e instanceof Error ? e.message : e}`);
    }
    return undefined;
}

/**
 * Fallback when frame count is unavailable:
 * - Get duration (sec)
 * - Seek extremely near the end with -sseof (negative seek from EOF)
 * - Reverse that tiny tail and grab the first decoded frame
 */
async function extractLastFrameFallback(inputVideoPath: string, outputImagePath: string): Promise<void> {
    await fs.mkdir(path.dirname(outputImagePath), { recursive: true });
    // Read duration to decide a safe small tail (0.2s or 1% of duration)
    let tail = 0.2;
    try {
        const { stdout } = await runFfprobe([
            '-v', 'error',
            '-select_streams', 'v:0',
            '-show_entries', 'format=duration',
            '-of', 'default=nokey=1:noprint_wrappers=1',
            inputVideoPath,
        ]);
        const dur = parseFloat(stdout.trim());
        if (Number.isFinite(dur) && dur > 0) {
            tail = Math.max(0.05, Math.min(0.5, dur * 0.01)); // 1% (min 0.05s, max 0.5s)
        }
    } catch { /* ignore */ }

    const args = [
        '-y',
        '-sseof', `-${tail}`,
        '-i', inputVideoPath,
        '-vf', 'reverse',
        '-vframes', '1',
        '-q:v', '2',
        outputImagePath,
    ];
    await runFfmpeg(args);
}

/**
 * Extract the last frame of a video into a PNG, accurately.
 * Preferred: precise frame index with select=eq(n\,last)
 * Fallback: tiny tail + reverse (handles weird containers/codecs)
 */
async function extractLastFrameAccurate(inputVideoPath: string, outputImagePath: string): Promise<void> {
    const total = await getVideoFrameCount(inputVideoPath);
    if (total && total > 0) {
        const last = Math.max(0, total - 1);
        await fs.mkdir(path.dirname(outputImagePath), { recursive: true });
        const args = [
            '-y',
            '-i', inputVideoPath,
            '-vf', `select=eq(n\\,${last})`,
            '-vframes', '1',
            '-q:v', '2',
            outputImagePath,
        ];
        try {
            await runFfmpeg(args);
            return;
        } catch (e) {
            logger.warn(`select by frame index failed, falling back: ${e instanceof Error ? e.message : e}`);
        }
    }
    await extractLastFrameFallback(inputVideoPath, outputImagePath);
}

/**
 * Concatenate two videos into a single MP4.
 * First try concat demuxer with stream copy (fast, no re-encode).
 * If it fails, fall back to re-encoding with concat filter.
 */
async function concatVideosFFmpeg(input1: string, input2: string, outputPath: string): Promise<void> {
    await fs.mkdir(path.dirname(outputPath), { recursive: true });

    // Try concat demuxer with copy (requires same codecs/params)
    const listTxtPath = path.join(path.dirname(outputPath), `concat_${Date.now()}.txt`);
    const listContent = `file '${path.resolve(input1).replace(/'/g, "'\\''")}'\nfile '${path.resolve(input2).replace(/'/g, "'\\''")}'\n`;
    await fs.writeFile(listTxtPath, listContent, 'utf-8');

    try {
        await runFfmpeg(['-y', '-f', 'concat', '-safe', '0', '-i', listTxtPath, '-c', 'copy', outputPath]);
    } catch (e) {
        logger.warn(`Concat with -c copy failed, falling back to re-encode: ${e instanceof Error ? e.message : e}`);
        // Fallback: re-encode with concat filter, video only
        await runFfmpeg([
            '-y',
            '-i', input1,
            '-i', input2,
            '-filter_complex', 'concat=n=2:v=1:a=0',
            '-c:v', 'libx264',
            '-crf', '18',
            '-preset', 'veryfast',
            '-pix_fmt', 'yuv420p',
            outputPath,
        ]);
    } finally {
        try { await fs.unlink(listTxtPath); } catch { /* ignore */ }
    }
}

async function main() {
    try {
        await ensureDirs();

        // Load scene config
        const sceneRaw = await fs.readFile(path.resolve('src/infinityvideo_generator/scene.json'), 'utf-8');
        const scene: InfinitySceneConfig = JSON.parse(sceneRaw);

        const servers = loadServers();
        if (servers.length === 0) {
            return;
        }

        // Optional limiter (env MAX_LOOPS). If not set, loop infinitely.
        const MAX_LOOPS = process.env.MAX_LOOPS ? Math.max(0, parseInt(process.env.MAX_LOOPS, 10)) : undefined;

        // Session identifiers
        const sessionId = Date.now().toString(36);
        let rrIndex = 0; // round-robin index across servers
        let iteration = 0;

        // Step 2: Determine initial image (use existing file if provided)
        const providedInitialImagePath = path.resolve('src/infinityvideo_generator/initial_image.png');
        let initialImagePath: string;
        let initialImageName: string;

        if (await fileExists(providedInitialImagePath)) {
            initialImagePath = providedInitialImagePath;
            initialImageName = path.basename(initialImagePath);
            logger.info(`Using existing initial image at ${initialImagePath}`);
        } else {
            initialImageName = `infinity_${sessionId}_i${iteration}.png`;
            const serverForInitialImage = pickServer(servers, rrIndex++);
            logger.info(`Generating initial image (${initialImageName}) on ${serverForInitialImage.name} (flux)...`);

            initialImagePath = await generateImage(
                scene.initialImage,
                initialImageName,
                serverForInitialImage.baseUrl!,
                serverForInitialImage.outputDir!,
                'flux',
                DEFAULT_SIZE,
            );
            logger.info(`Initial image generated: ${initialImagePath}`);
        }

        // Step 3: Copy image to input folders and generate first video
        const imageNameForComfy = await copyImageToAllServerInputs(servers, initialImagePath);

        const firstVideoName = initialImageName.replace(/\.png$/i, `_v${iteration}.mp4`);
        const serverForFirstVideo = pickServer(servers, rrIndex++);
        logger.info(`Generating first video (${firstVideoName}) on ${serverForFirstVideo.name} using ${imageNameForComfy}...`);

        // Use "light" workflow and short length similar to musicspot videos
        let currentVideoPath = await generateVideo(
            scene.videoPrompt,
            imageNameForComfy,
            firstVideoName,
            serverForFirstVideo.baseUrl!,
            serverForFirstVideo.outputDir!,
            DEFAULT_SIZE,
            true,
            true
        );
        logger.info(`First video generated: ${currentVideoPath}`);

        // Loop:
        // 4) Extract last frame (accurate)
        // 5) Generate a new video from that image
        // 6) Concat current video + new video => becomes the new "current" video
        while (true) {
            if (MAX_LOOPS !== undefined && iteration >= MAX_LOOPS) {
                logger.info(`Reached MAX_LOOPS=${MAX_LOOPS}. Stopping.`);
                break;
            }

            iteration += 1;
            const lastFrameName = `infinity_${sessionId}_lastframe_${iteration}.png`;
            const lastFramePath = path.join(GENERATED_DIR, lastFrameName);

            logger.info(`Extracting last frame (accurate) from ${currentVideoPath} -> ${lastFramePath}`);
            await extractLastFrameAccurate(currentVideoPath, lastFramePath);
            const lastFrameExists = await fileExists(lastFramePath);
            if (!lastFrameExists) {
                throw new Error(`Failed to extract last frame to ${lastFramePath}`);
            }

            // Copy to server inputs
            const frameNameForComfy = await copyImageToAllServerInputs(servers, lastFramePath);

            // Generate new video from last frame
            const newVideoName = `infinity_${sessionId}_video_${iteration}.mp4`;
            const serverForVideo = pickServer(servers, rrIndex++);
            logger.info(`Generating new video (${newVideoName}) on ${serverForVideo.name} using ${frameNameForComfy}...`);

            const newVideoPath = await generateVideo(
                scene.videoPrompt,
                frameNameForComfy,
                newVideoName,
                serverForVideo.baseUrl!,
                serverForVideo.outputDir!,
                DEFAULT_SIZE,
                true,
                true
            );
            logger.info(`New segment generated: ${newVideoPath}`);

            // Concat current + new => new current
            const concatenatedName = `infinity_${sessionId}_concat_${iteration}.mp4`;
            const concatenatedPath = path.join(GENERATED_DIR, concatenatedName);

            logger.info(`Concatenating videos: [${currentVideoPath}] + [${newVideoPath}] -> ${concatenatedPath}`);
            await concatVideosFFmpeg(currentVideoPath, newVideoPath, concatenatedPath);
            logger.info(`Concatenated video: ${concatenatedPath}`);

            // Set as the new current video for next loop
            currentVideoPath = concatenatedPath;
        }

        logger.info('Infinity video generation finished.');
    } catch (err) {
        logger.error('Fatal error in infinity video generator:', err);
    }
}

main().catch((err) => {
    logger.error('Unhandled error:', err);
});
