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

import { logger } from '../lib/logger';
import { callOpenAI } from '../lib/openai';
import { callLMStudio } from '../lib/lmstudio';
// import { generateImage as generateFaceImage } from '../lib/image-generator-face'; // Removed
import { generateImage } from '../lib/image-generator'; // Added

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 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: 1280, height: 720 };
const FOLDER = process.argv[2] || process.env.MUSICSPOT_FOLDER || 'infinitydance';
const FOLDER_SAFE = FOLDER.replace(/[/\\?%*:|"<>]/g, '_');
const FACE_SRC = path.resolve(`src/musicspot_generator/${FOLDER}/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,
            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<string | undefined> {
    const faceFileName = 'face.png';

    // Validate face source
    try {
        await fs.access(FACE_SRC);
    } catch {
        // If face.png doesn't exist, we don't need to copy anything.
        // The caller should handle this case.
        return undefined;
    }

    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);
            logger.info(`Copied face image to ${srv.name} input: ${dest}`);
        } catch (err) {
            logger.warn(`Failed to copy face image to ${srv.name}: ${err}`);
        }
    }

    return faceFileName; // Comfy expects file name present in input dir
}

function buildImagePromptRequest(
    character: MusicSpotCharacter,
    scene: MusicSpotScene,
    cut: MusicSpotCut,
    cameraIntent: string,
    hasFaceImage: boolean // Added parameter
): string {
    let promptInstructions = `
Write "imagePrompt" in around 110–140 words to generate a still portrait image (720x1280 vertical).
`;

    if (hasFaceImage) {
        promptInstructions += `Keep a consistent character identity using the provided face image (identity preservation), but do not mention any camera brand/model.
`;
    } else {
        promptInstructions += `Do not use any face image for identity preservation.
`;
    }

    return `
Return exactly one JSON object, nothing else: { "imagePrompt": "Cinematic realistic photo, (camera framing),(character),(pose),(time),(location),(outfit),(action),(lighting)" }.

${promptInstructions}
Describe clearly and concretely:
- Character: ${character.bodyType}; hair: ${character.hairStyle}
- Camera framing/composition intention: ${cameraIntent}
- Time: ${scene.time}
- Location: ${scene.location}
- Outfit: ${scene.outfit}
- Pose: ${cut.pose}
- Action/Expression: ${cut.action}
- Lighting: please be creative and make beautiful lighting, I like something like luminous, colorful

Important:
- adjust the outfit based on the camera framing, describe only what is visible.
For example, if the framing is a close-up of the face, do not mention the outfit at all.

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();
}

const IMAGE_PROMPT_PROVIDER = (process.env.IMAGE_PROMPT_PROVIDER || 'lmstudio').toLowerCase();

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

async function getImagePrompt(req: string): Promise<string> {
    if (IMAGE_PROMPT_PROVIDER === 'openai') {
        return getImagePromptFromOpenAI(req);
    }
    // default to LM Studio
    return getImagePromptFromLMStudio(req);
}

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

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

        // Check if face.png exists
        let faceImageExists = false;
        try {
            await fs.access(FACE_SRC);
            faceImageExists = true;
            logger.info(`Face image found at ${FACE_SRC}`);
        } catch {
            logger.warn(`Face image not found at ${FACE_SRC}. Images will be generated without face conditioning.`);
        }

        // Load scenes.json
        const configRaw = await fs.readFile(path.resolve(`src/musicspot_generator/${FOLDER}/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, only if it exists
        let faceFileName: string | undefined = undefined;
        if (faceImageExists) {
            faceFileName = await copyFaceToServers(servers);
        } else {
            logger.info('Skipping copyFaceToServers as face image does not exist.');
        }

        // Generate images only (no video here). Intended to be run first.
        let imageTaskIndex = 0;

        for (const scene of cfg.scenes) {
            logger.info(`=== Scene ${scene.sceneId}: Image generation start ===`);
            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;
                    const imgFileName = `${FOLDER_SAFE}_musicspot_s${scene.sceneId}_c${cut.cutId}_v${variantIndex}.png`;
                    const outputPath = path.join(GENERATED_DIR, imgFileName);

                    // Skip generation if target file already exists
                    try {
                        await fs.access(outputPath);
                        logger.info(`Skipping generation, file already exists: ${outputPath}`);
                        continue;
                    } catch {
                        // File does not exist; proceed with generation
                    }

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

                    // 2) Generate one image using the new generateImage function
                    const serverForImage = pickServer(servers, imageTaskIndex++);
                    logger.info(`Generating image (${imgFileName}) on ${serverForImage.name}...`);

                    try {
                        // Use the generic generateImage function from image-generator.ts
                        // The faceFileName is no longer passed as an argument, but its existence
                        // influenced the imagePrompt.
                        const finalImagePath = await generateImage(
                            imagePrompt, // The prompt now contains face instructions if faceImageExists is true
                            imgFileName,
                            serverForImage.baseUrl!,
                            serverForImage.outputDir!,
                            'flux', // Use default imageModel ('qwen')
                            DEFAULT_SIZE
                        );
                        logger.info(`Image generated: ${finalImagePath}`);
                    } catch (err) {
                        logger.error(`Image generation failed (${imgFileName}) on ${serverForImage.name}: ${err}`);
                        continue;
                    }
                }
            }
            logger.info(`=== Scene ${scene.sceneId}: Image generation complete ===`);
        }

        logger.info('Image generation for all scenes completed.');
    } catch (err) {
        logger.error('Fatal error in music spot image generator:', err);
    }
}

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