import dotenv from 'dotenv';
import path from 'path';
import fs from 'fs/promises';

import { logger } from '../../lib/logger';
import { callOpenAI } from '../../lib/openai';
import { generateImage as generateFaceImage } from '../../lib/image-generator-face';
import { generateVideo } from '../../lib/video-generator';

dotenv.config();

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

interface MusicSpotCharacter {
    bodyType: string;
    hairStyle: string;
}

interface MusicSpotCut {
    cutId: number;
    pose: string;
    action: string;
    camera?: string[]; // list of 5 camera variants per cut
}

interface MusicSpotScene {
    sceneId: number;
    time: string;
    location: string;
    outfit: string;
    cuts: MusicSpotCut[];
}

interface MusicSpotConfig {
    character: MusicSpotCharacter;
    scenes: MusicSpotScene[];
}

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

const DEFAULT_SIZE: Size = { width: 720, height: 1280 };
const FACE_SRC = path.resolve('src/musicspot_generator/face.png');
const GENERATED_DIR = path.resolve('generated');

function loadServers(): Server[] {
    const servers: Server[] = [
        {
            name: 'SERVER1',
            baseUrl: process.env.SERVER1_COMFY_BASE_URL,
            outputDir: process.env.SERVER1_COMFY_OUTPUT_DIR,
        }
        /*{
            name: 'SERVER2',
            baseUrl: process.env.SERVER2_COMFY_BASE_URL,
            outputDir: process.env.SERVER2_COMFY_OUTPUT_DIR,
        }*/,
    ]
        .filter(s => !!s.baseUrl && !!s.outputDir)
        .map(s => ({
            ...s,
            // Convert output dir to input dir by convention
            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 copyFaceToServers(servers: Server[]): Promise<{ fileName: string; absPerServer: Record<string, string> }> {
    const faceFileName = 'face.png';
    const absPerServer: Record<string, string> = {};

    // Validate face source
    try {
        await fs.access(FACE_SRC);
    } catch {
        throw new Error(`Face image not found at ${FACE_SRC}`);
    }

    for (const srv of servers) {
        if (!srv.inputDir) continue;
        const dest = path.join(srv.inputDir, faceFileName);
        try {
            await fs.mkdir(srv.inputDir, { recursive: true });
            await fs.copyFile(FACE_SRC, dest);
            absPerServer[srv.name] = path.resolve(dest);
            logger.info(`Copied face image to ${srv.name} input: ${dest}`);
        } catch (err) {
            logger.warn(`Failed to copy face image to ${srv.name}: ${err}`);
        }
    }

    return { fileName: faceFileName, absPerServer };
}

function buildImagePromptRequest(
    character: MusicSpotCharacter,
    scene: MusicSpotScene,
    cut: MusicSpotCut,
    cameraIntent: string
): string {
    // Ask OpenAI to return JSON: { "imagePrompt": "..." }
    return `
Return exactly one JSON object, nothing else: { "imagePrompt": "..." }.

Write "imagePrompt" in around 110–140 words to generate a still portrait image (720x1280 vertical).
Keep a consistent character identity using the provided face image (identity preservation), but do not mention any camera brand/model.

Describe clearly and concretely:
- Character: ${character.bodyType}; hair: ${character.hairStyle}
- Time: ${scene.time}
- Location: ${scene.location}
- Outfit: ${scene.outfit}
- Pose: ${cut.pose}
- Action/Expression: ${cut.action}
- Camera framing/composition intention: ${cameraIntent}
- Lighting/mood/style: cohesive and realistic, natural skin tones, soft depth of field.

Avoid: brand names, copyrighted characters, extreme or explicit content, text overlays, watermarks, multiple people. Focus on a single subject medium-full portrait, tasteful and aesthetic. Use simple sentences.

Only respond with JSON.
  `.trim();
}

function buildVideoPromptRequest(
    character: MusicSpotCharacter,
    scene: MusicSpotScene,
    cut: MusicSpotCut,
    cameraIntent: string
): string {
    // Ask OpenAI to return JSON: { "videoPrompt": "..." }
    // Strong constraints to avoid "cut/zoom" etc. Keep a single continuous 8s shot.
    return `
Return exactly one JSON object and nothing else: { "videoPrompt": "..." }.

Write "videoPrompt" in 100–140 words. Present tense. Concrete, simple sentences.

HARD RULES:
- One continuous 8-second shot (oner). No edits.
- Fixed location and general vantage; maintain spatial continuity.
- No zooms, no rack zoom, no smash/push-in, no cuts, no transitions, no "meanwhile".
- Camera motion: at most a slight pan/tilt or subtle dolly within 1 meter.
- Keep framing consistent (vertical 720x1280). Avoid technical brand names or lens jargon.

Incorporate the following camera intention: "${cameraIntent}".
If it conflicts with HARD RULES (e.g., zoom, push-in, extreme moves), reinterpret it into a subtle, compliant motion (e.g., gentle glide, slight pan/tilt) while preserving the creative intent.

Describe:
1) Main action: ${cut.action}
2) Pose/composition: ${cut.pose}
3) Scene/time/location/outfit: ${scene.time}; ${scene.location}; outfit: ${scene.outfit}
4) Lighting/mood/style coherent with the character: ${character.bodyType}; hair: ${character.hairStyle}

Prohibited (case-insensitive): cut, cuts, cutting, quick cut, insert, close-up, extreme close-up, zoom, zooming, push-in, pull-out, whip, switch angle, change angle, montage, cross-cut, smash cut, transition, meanwhile, later.

Only respond with JSON.
  `.trim();
}

async function getImagePromptFromOpenAI(req: string): Promise<string> {
    const res = await callOpenAI(req);
    const prompt = res?.imagePrompt || res?.image_prompt || res?.prompt;
    if (!prompt || typeof prompt !== 'string') {
        throw new Error('OpenAI failed to return imagePrompt JSON.');
    }
    return prompt.trim();
}

async function getVideoPromptFromOpenAI(req: string): Promise<string> {
    const res = await callOpenAI(req);
    const prompt = res?.videoPrompt || res?.video_prompt || res?.prompt;
    if (!prompt || typeof prompt !== 'string') {
        throw new Error('OpenAI failed to return videoPrompt JSON.');
    }
    return prompt.trim();
}

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

async function copyImageToAllServerInputs(servers: Server[], localGeneratedImagePath: string): Promise<string> {
    const fileName = path.basename(localGeneratedImagePath);
    const copies: string[] = [];
    for (const s of servers) {
        if (!s.inputDir) continue;
        const dest = path.join(s.inputDir, fileName);
        try {
            await fs.copyFile(localGeneratedImagePath, dest);
            copies.push(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; // return the name used for Comfy workflows
}

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

        // Load scenes.json
        const configRaw = await fs.readFile(path.resolve('src/musicspot_generator/scenes.json'), 'utf-8');
        const cfg: MusicSpotConfig = JSON.parse(configRaw);

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

        // Ensure face.png in each server's input
        const { fileName: faceFileName } = await copyFaceToServers(servers);

        // Two-phase per scene:
        // Phase 1: Generate ALL images for the scene (to keep image model warm)
        // Phase 2: Generate ALL videos for the images of that scene (to keep video model warm)

        let imageTaskIndex = 0;
        let videoTaskIndex = 0;

        for (const scene of cfg.scenes) {
            logger.info(`=== Scene ${scene.sceneId}: Phase 1 - Image generation start ===`);
            type SceneImageItem = {
                sceneId: number;
                cutId: number;
                cameraIntent: string;
                imgFileName: string;
                imgPath: string;
            };
            const sceneImageItems: SceneImageItem[] = [];

            // Phase 1: images
            for (const cut of scene.cuts) {
                const cameraVariants = Array.isArray(cut.camera) && cut.camera.length > 0
                    ? cut.camera
                    : ['eye-level medium shot', 'slight left 30°', 'slight right 30°', 'slight high angle', 'slight low angle'];

                for (let camIdx = 0; camIdx < cameraVariants.length; camIdx++) {
                    const cameraIntent = cameraVariants[camIdx];
                    const variantIndex = camIdx + 1;

                    // 1) Generate image prompt for this camera
                    logger.info(`Scene ${scene.sceneId} - Cut ${cut.cutId} - Cam${variantIndex}: generating image prompt...`);
                    const imgPromptReq = buildImagePromptRequest(cfg.character, scene, cut, cameraIntent);
                    let imagePrompt: string;
                    try {
                        imagePrompt = await getImagePromptFromOpenAI(imgPromptReq);
                    } catch (err) {
                        logger.error(`OpenAI image prompt failed for scene ${scene.sceneId} cut ${cut.cutId} cam ${variantIndex}: ${err}`);
                        continue;
                    }

                    // 2) Generate one image using face conditioning for this specific camera
                    const serverForImage = pickServer(servers, imageTaskIndex++);
                    const imgFileName = `musicspot_s${scene.sceneId}_c${cut.cutId}_v${variantIndex}.png`;
                    logger.info(`Generating image (${imgFileName}) on ${serverForImage.name}...`);

                    try {
                        // Use only the face file name for the workflow image input (Comfy expects it in its input dir)
                        const finalImagePath = await generateFaceImage(
                            imagePrompt,
                            faceFileName,
                            imgFileName,
                            serverForImage.baseUrl!,
                            serverForImage.outputDir!,
                            DEFAULT_SIZE
                        );
                        logger.info(`Image generated: ${finalImagePath}`);

                        sceneImageItems.push({
                            sceneId: scene.sceneId,
                            cutId: cut.cutId,
                            cameraIntent,
                            imgFileName,
                            imgPath: finalImagePath,
                        });
                    } catch (err) {
                        logger.error(`Image generation failed (${imgFileName}) on ${serverForImage.name}: ${err}`);
                        continue;
                    }
                }
            }

            logger.info(`=== Scene ${scene.sceneId}: Phase 1 complete. Generated ${sceneImageItems.length} image(s). ===`);
            logger.info(`=== Scene ${scene.sceneId}: Phase 2 - Video generation start ===`);

            // Phase 2: videos for all images generated in this scene
            for (const item of sceneImageItems) {
                // Find original cut info for prompts
                const cut = scene.cuts.find(c => c.cutId === item.cutId);
                if (!cut) {
                    logger.warn(`Cut ${item.cutId} not found for scene ${scene.sceneId}, skipping video.`);
                    continue;
                }

                // 3) Generate video prompt for this camera
                logger.info(`Generating video prompt for ${item.imgFileName} (scene ${scene.sceneId}, cut ${item.cutId})...`);
                const vidPromptReq = buildVideoPromptRequest(cfg.character, scene, cut, item.cameraIntent);

                let videoPrompt: string;
                try {
                    videoPrompt = await getVideoPromptFromOpenAI(vidPromptReq);
                } catch (err) {
                    logger.error(`OpenAI video prompt failed for ${item.imgFileName}: ${err}`);
                    continue;
                }

                // 4) Copy the base image to every server's input folder
                const imageFileNameForComfy = await copyImageToAllServerInputs(servers, item.imgPath);

                // 5) Generate video on a chosen server (round-robin)
                const serverForVideo = pickServer(servers, videoTaskIndex++);
                const videoFileName = item.imgFileName.replace(/\.png$/i, '.mp4');
                logger.info(`Generating video (${videoFileName}) on ${serverForVideo.name} using ${imageFileNameForComfy}...`);

                try {
                    const videoPath = await generateVideo(
                        videoPrompt,
                        imageFileNameForComfy,
                        videoFileName,
                        serverForVideo.baseUrl!,
                        serverForVideo.outputDir!,
                        DEFAULT_SIZE
                    );
                    logger.info(`Video generated: ${videoPath}`);
                } catch (err) {
                    logger.error(`Video generation failed (${videoFileName}) on ${serverForVideo.name}: ${err}`);
                }
            }

            logger.info(`=== Scene ${scene.sceneId}: Phase 2 complete. ===`);
        }

        logger.info('Music spot generation completed.');
    } catch (err) {
        logger.error('Fatal error in music spot generator:', err);
    } finally {
        // Optional: do not exit to allow long-running processes; but align with other scripts:
        // process.exit();
    }
}

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