import * as fs from 'fs/promises';
import * as path from 'path';
import dotenv from 'dotenv';
import { downloadImagesFromPinterestSearch } from '../lib/pinterest';
import { convertImage, convertImageWithFile, convertImageWithFileForPose, convertImageWithFileHandbag } from '../lib/image-converter';
import { logger } from '../lib/logger';
import { callLmstudio, callLMStudioAPIWithFile } from '../lib/lmstudio';
import { upscale } from '../lib/image-upscaler';
dotenv.config();

const SERVER1_COMFY_BASE_URL = process.env.SERVER1_COMFY_BASE_URL!;
const SERVER1_COMFY_OUTPUT_DIR = process.env.SERVER1_COMFY_OUTPUT_DIR!;
const imageSize: { width: number; height: number } = { width: 1280, height: 720 };


async function upscaleAndFix(
    baseImage: string,
    faceImage: string,
    outputFilename: string,
    outputDir: string,
): Promise<void> {
    try {

        // Copy both images to ComfyUI input directory
        const inputFolderFullPath = SERVER1_COMFY_OUTPUT_DIR.replace('output', 'input');
        await fs.mkdir(inputFolderFullPath, { recursive: true });

        const baseFilePath = path.join(outputDir, baseImage);
        const referenceFilePath = path.join(outputDir, faceImage);

        const baseFileName = path.basename(baseImage);
        const referenceFileName = path.basename(faceImage);

        const inputBasePath = path.join(inputFolderFullPath, baseFileName);
        const inputReferencePath = path.join(inputFolderFullPath, referenceFileName);

        logger.info(`Copying base image to ComfyUI input: ${inputBasePath}`);
        await fs.copyFile(baseFilePath, inputBasePath);

        logger.info(`Copying reference image to ComfyUI input: ${inputReferencePath}`);
        await fs.copyFile(referenceFilePath, inputReferencePath);

        // Convert images with prompt
        logger.info(`Processing images with convertImageWithFile...`);

        const convertedImagePath = await upscale(
            baseFileName,
            SERVER1_COMFY_BASE_URL,
            SERVER1_COMFY_OUTPUT_DIR,
        )

        logger.info(`Converted image: ${convertedImagePath}`);

        // Copy the converted image to final destination
        const finalOutputPath = path.join(outputDir, outputFilename);
        logger.info(`Copying to final destination: ${finalOutputPath}`);
        await fs.copyFile(convertedImagePath, finalOutputPath);

        logger.info(`✓ Successfully generated: ${finalOutputPath}`);

    } catch (error) {
        logger.error(`Error processing two images:`, error);
        throw error;
    }
}

/**
 * Process a single image: download from Pinterest, convert with prompt, and save
 * @param keyword - Pinterest search keyword
 * @param prompt - Image conversion prompt
 * @param filename - Output filename
 * @param outputDir - Directory to save the generated file
 * @param shouldConvert - Whether to convert the image with prompt or just copy it
 */
async function processImage(
    keyword: string,
    prompt: string,
    filename: string,
    outputDir: string,
    shouldConvert: boolean = true
): Promise<void> {
    try {
        logger.info(`\n=== Processing: ${filename} ===`);
        logger.info(`Keyword: ${keyword}`);
        logger.info(`Should convert: ${shouldConvert}`);

        // Step 1: Download image from Pinterest
        logger.info(`Step 1: Downloading image from Pinterest with keyword: "${keyword}"...`);
        const downloadedImages = await downloadImagesFromPinterestSearch(keyword, 1);

        if (downloadedImages.length === 0) {
            logger.error(`Failed to download image for keyword: "${keyword}"`);
            return;
        }

        const downloadedImagePath = downloadedImages[0];
        logger.info(`Downloaded image: ${downloadedImagePath}`);

        const finalOutputPath = path.join(outputDir, filename);

        if (shouldConvert) {
            logger.info(`Prompt: ${prompt}`);

            // Step 2: Copy image to ComfyUI input directory
            const inputFolderFullPath = SERVER1_COMFY_OUTPUT_DIR.replace('output', 'input');
            await fs.mkdir(inputFolderFullPath, { recursive: true });

            const imageFileName = path.basename(downloadedImagePath);
            const inputImagePath = path.join(inputFolderFullPath, imageFileName);

            logger.info(`Step 2: Copying image to ComfyUI input folder: ${inputImagePath}`);
            await fs.copyFile(downloadedImagePath, inputImagePath);

            // Step 3: Convert image with prompt
            logger.info(`Step 3: Converting image with prompt...`);
            const convertedImagePath = await convertImage(
                prompt,
                imageFileName,
                SERVER1_COMFY_BASE_URL,
                SERVER1_COMFY_OUTPUT_DIR,
                imageSize
            );
            logger.info(`Converted image: ${convertedImagePath}`);

            // Step 4: Copy the converted image to final destination
            logger.info(`Step 4: Copying to final destination: ${finalOutputPath}`);
            await fs.copyFile(convertedImagePath, finalOutputPath);
        } else {
            // Just copy the downloaded image directly to the output directory with the specified filename
            logger.info(`Step 2: Copying directly to final destination: ${finalOutputPath}`);
            await fs.copyFile(downloadedImagePath, finalOutputPath);
        }

        logger.info(`✓ Successfully generated: ${finalOutputPath}`);

    } catch (error) {
        logger.error(`Error processing image for keyword "${keyword}":`, error);
        throw error;
    }
}

/**
 * Convert an existing image with a prompt
 * @param prompt - Image conversion prompt
 * @param imagePath - Path to the existing image
 * @param outputFilename - Output filename
 * @param outputDir - Directory to save the converted file
 */
async function convertImageWithPrompt(
    prompt: string,
    imagePath: string,
    outputFilename: string,
    outputDir: string
): Promise<void> {
    try {
        logger.info(`\n=== Converting Image: ${outputFilename} ===`);
        logger.info(`Source: ${imagePath}`);
        logger.info(`Prompt: ${prompt}`);

        // Step 1: Copy image to ComfyUI input directory
        const inputFolderFullPath = SERVER1_COMFY_OUTPUT_DIR.replace('output', 'input');
        await fs.mkdir(inputFolderFullPath, { recursive: true });

        const imageFileName = path.basename(imagePath);
        const inputImagePath = path.join(inputFolderFullPath, imageFileName);

        logger.info(`Step 1: Copying image to ComfyUI input folder: ${inputImagePath}`);
        await fs.copyFile(imagePath, inputImagePath);

        // Step 2: Convert image with prompt
        logger.info(`Step 2: Converting image with prompt...`);
        const convertedImagePath = await convertImage(
            prompt,
            imageFileName,
            SERVER1_COMFY_BASE_URL,
            SERVER1_COMFY_OUTPUT_DIR,
            imageSize
        );
        logger.info(`Converted image: ${convertedImagePath}`);

        // Step 3: Copy the converted image to final destination
        const finalOutputPath = path.join(outputDir, outputFilename);
        logger.info(`Step 3: Copying to final destination: ${finalOutputPath}`);
        await fs.copyFile(convertedImagePath, finalOutputPath);

        logger.info(`✓ Successfully converted: ${finalOutputPath}`);

    } catch (error) {
        logger.error(`Error converting image:`, error);
        throw error;
    }
}

/**
 * Process two images together: combine base image with reference image using prompt
 * @param prompt - Processing prompt
 * @param baseFile - Base image filename (in generated folder)
 * @param referenceFile - Reference image filename (in generated folder)
 * @param outputFilename - Output filename
 * @param outputDir - Directory to save the generated file
 */
async function processTwoImages(
    prompt: string,
    baseFile: string,
    referenceFile: string,
    outputFilename: string,
    outputDir: string,
    isPose: boolean = false
): Promise<void> {
    try {
        logger.info(`\n=== Processing: ${outputFilename} ===`);
        logger.info(`Base: ${baseFile}, Reference: ${referenceFile}`);
        logger.info(`Prompt: ${prompt}`);

        // Copy both images to ComfyUI input directory
        const inputFolderFullPath = SERVER1_COMFY_OUTPUT_DIR.replace('output', 'input');
        await fs.mkdir(inputFolderFullPath, { recursive: true });

        const baseFilePath = path.join(outputDir, baseFile);
        const referenceFilePath = path.join(outputDir, referenceFile);

        const baseFileName = path.basename(baseFile);
        const referenceFileName = path.basename(referenceFile);

        const inputBasePath = path.join(inputFolderFullPath, baseFileName);
        const inputReferencePath = path.join(inputFolderFullPath, referenceFileName);

        logger.info(`Copying base image to ComfyUI input: ${inputBasePath}`);
        await fs.copyFile(baseFilePath, inputBasePath);

        logger.info(`Copying reference image to ComfyUI input: ${inputReferencePath}`);
        await fs.copyFile(referenceFilePath, inputReferencePath);

        // Convert images with prompt
        logger.info(`Processing images with convertImageWithFile...`);

        const convertedImagePath = isPose ? await convertImageWithFileForPose(
            prompt,
            baseFileName,
            referenceFileName,
            SERVER1_COMFY_BASE_URL,
            SERVER1_COMFY_OUTPUT_DIR,
            imageSize
        ) : await convertImageWithFile(
            prompt,
            baseFileName,
            referenceFileName,
            SERVER1_COMFY_BASE_URL,
            SERVER1_COMFY_OUTPUT_DIR,
            imageSize
        )

        logger.info(`Converted image: ${convertedImagePath}`);

        // Copy the converted image to final destination
        const finalOutputPath = path.join(outputDir, outputFilename);
        logger.info(`Copying to final destination: ${finalOutputPath}`);
        await fs.copyFile(convertedImagePath, finalOutputPath);

        logger.info(`✓ Successfully generated: ${finalOutputPath}`);

    } catch (error) {
        logger.error(`Error processing two images:`, error);
        throw error;
    }
}


/**
 * Process two images together: combine base image with reference image using prompt
 * @param prompt - Processing prompt
 * @param baseFile - Base image filename (in generated folder)
 * @param referenceFile - Reference image filename (in generated folder)
 * @param outputFilename - Output filename
 * @param outputDir - Directory to save the generated file
 */
async function processTwoImagesHandbag(
    prompt: string,
    baseFile: string,
    referenceFile: string,
    outputFilename: string,
    outputDir: string,
): Promise<void> {
    try {
        logger.info(`\n=== Processing: ${outputFilename} ===`);
        logger.info(`Base: ${baseFile}, Reference: ${referenceFile}`);
        logger.info(`Prompt: ${prompt}`);

        // Copy both images to ComfyUI input directory
        const inputFolderFullPath = SERVER1_COMFY_OUTPUT_DIR.replace('output', 'input');
        await fs.mkdir(inputFolderFullPath, { recursive: true });

        const baseFilePath = path.join(outputDir, baseFile);
        const referenceFilePath = path.join(outputDir, referenceFile);

        const baseFileName = path.basename(baseFile);
        const referenceFileName = path.basename(referenceFile);

        const inputBasePath = path.join(inputFolderFullPath, baseFileName);
        const inputReferencePath = path.join(inputFolderFullPath, referenceFileName);

        logger.info(`Copying base image to ComfyUI input: ${inputBasePath}`);
        await fs.copyFile(baseFilePath, inputBasePath);

        logger.info(`Copying reference image to ComfyUI input: ${inputReferencePath}`);
        await fs.copyFile(referenceFilePath, inputReferencePath);

        // Convert images with prompt
        logger.info(`Processing images with convertImageWithFile...`);

        const convertedImagePath = await convertImageWithFileHandbag(
            prompt,
            baseFileName,
            referenceFileName,
            SERVER1_COMFY_BASE_URL,
            SERVER1_COMFY_OUTPUT_DIR,
            imageSize
        )

        logger.info(`Converted image: ${convertedImagePath}`);

        // Copy the converted image to final destination
        const finalOutputPath = path.join(outputDir, outputFilename);
        logger.info(`Copying to final destination: ${finalOutputPath}`);
        await fs.copyFile(convertedImagePath, finalOutputPath);

        logger.info(`✓ Successfully generated: ${finalOutputPath}`);

    } catch (error) {
        logger.error(`Error processing two images:`, error);
        throw error;
    }
}

/**
 * Process a complete iteration: download base images and apply sequential transformations
 */
async function processIteration(iteration: number): Promise<void> {
    try {
        const timestamp = Date.now();
        logger.info(`\n${'='.repeat(80)}`);
        logger.info(`ITERATION ${iteration} - Starting with timestamp: ${timestamp}`);
        logger.info(`${'='.repeat(80)}`);

        // Create output directory for this iteration
        const outputDir = path.join(process.cwd(), 'generated', `vton_${timestamp}`);
        await fs.mkdir(outputDir, { recursive: true });
        logger.info(`Output directory created: ${outputDir}`);

        // === PHASE 1: Download base images ===
        logger.info(`\n--- PHASE 1: Downloading base images ---`);

        await processImage(
            'cute girl face high resolution',
            '',
            `model_${timestamp}.png`,
            outputDir,
            false
        );

        await processImage(
            'woman elegant outfit fullbody single',
            '',
            `outfit_${timestamp}.png`,
            outputDir,
            false
        );

        await processImage(
            'photo elegant indoor room',
            '',
            `room_${timestamp}.png`,
            outputDir,
            false
        );

        await processImage(
            'handbag single product photography',
            '请提取照片中的包，并将其正面朝向地放置在亮灰色背景上。',
            `handbag_${timestamp}.png`,
            outputDir,
            true
        );

        await processImage(
            'woman portrait standing',
            '',
            `pose_${timestamp}.png`,
            outputDir,
            false
        );

        // === PHASE 2: Sequential transformations ===
        logger.info(`\n--- PHASE 2: Sequential transformations ---`);

        // Step 1: Generate outfit prompt using LMStudio API
        logger.info('Step 1: Generating outfit prompt with LMStudio API...');
        const outfitImagePath = path.join(outputDir, `outfit_${timestamp}.png`);
        const outfitPromptResponse = await callLMStudioAPIWithFile(
            outfitImagePath,
            'Describe this outfit in detail about  30 words. Focus on color and cloth type. Return the result in this format: {"result":""}'
        );
        const outfitPrompt = outfitPromptResponse.result || outfitPromptResponse;
        logger.info(`Generated outfit prompt: ${outfitPrompt}`);

        // Step 2: Generate location prompt using LMStudio API
        logger.info('Step 2: Generating location prompt with LMStudio API...');
        const roomImagePath = path.join(outputDir, `room_${timestamp}.png`);
        const locationPromptResponse = await callLMStudioAPIWithFile(
            roomImagePath,
            'Describe this location/room in detail about  30 words. Return the result in this format: {"result":""}'
        );
        const locationPrompt = locationPromptResponse.result || locationPromptResponse;
        logger.info(`Generated location prompt: ${locationPrompt}`);

        // Step 3: Generate Chinese prompt using LMStudio API
        logger.info('Step 3: Generating Chinese prompt for model transformation...');
        const chinesePromptRequest = `Generate a Chinese prompt for image transformation that describes:
- Prefix: genereate a portarit photo of a woman in image1
- Use outfit to: ${outfitPrompt}
- Use location to: ${locationPrompt}

Return the result in this format: {"result":""}`;

        const chinesePromptResponse = await callLmstudio(chinesePromptRequest);
        const chinesePrompt = chinesePromptResponse.result || chinesePromptResponse;
        logger.info(`Generated Chinese prompt: ${chinesePrompt}`);

        // Process model with outfit and location using the Chinese prompt
        logger.info('Step 4: Processing model with outfit and location...');
        const modelImagePath = path.join(outputDir, `model_${timestamp}.png`);

        // Copy model image to ComfyUI input directory
        const inputFolderFullPath = SERVER1_COMFY_OUTPUT_DIR.replace('output', 'input');
        await fs.mkdir(inputFolderFullPath, { recursive: true });
        const modelFileName = path.basename(modelImagePath);
        const inputModelPath = path.join(inputFolderFullPath, modelFileName);
        await fs.copyFile(modelImagePath, inputModelPath);

        // Convert image with Chinese prompt and pose 
        await processTwoImages(
            `请将图1中模特的姿势更改为图2的姿势。, ${chinesePrompt}`,
            modelFileName,
            `pose_${timestamp}.png`,
            `model_outfit_location_pose_${timestamp}.png`,
            outputDir,
            true
        );

        // Step 5: Add handbag to model
        await processTwoImagesHandbag(
            '请将图1中的女性修改成手持图2的包。',
            `model_outfit_location_pose_${timestamp}.png`,
            `handbag_${timestamp}.png`,
            `model_outfit_location_handbag1_${timestamp}.png`,
            outputDir
        );


        await processTwoImagesHandbag(
            '请让图1的女性看起来像是在手里拿着图2的包。',
            `model_outfit_location_pose_${timestamp}.png`,
            `handbag_${timestamp}.png`,
            `model_outfit_location_handbag2_${timestamp}.png`,
            outputDir
        );

        await processTwoImagesHandbag(
            '请将图1中的女性修改成双手拿着图2的包。',
            `model_outfit_location_pose_${timestamp}.png`,
            `handbag_${timestamp}.png`,
            `model_outfit_location_handbag3_${timestamp}.png`,
            outputDir
        );

        await upscaleAndFix(
            `model_outfit_location_handbag1_${timestamp}.png`,
            `model_${timestamp}.png`,
            `model_outfit_location_handbag1_upscaled_${timestamp}.png`,
            outputDir
        );

        await upscaleAndFix(
            `model_outfit_location_handbag2_${timestamp}.png`,
            `model_${timestamp}.png`,
            `model_outfit_location_handbag2_upscaled_${timestamp}.png`,
            outputDir
        );

        await upscaleAndFix(
            `model_outfit_location_handbag3_${timestamp}.png`,
            `model_${timestamp}.png`,
            `model_outfit_location_handbag3_upscaled_${timestamp}.png`,
            outputDir
        );


        logger.info(`\n${'='.repeat(80)}`);
        logger.info(`ITERATION ${iteration} COMPLETED!`);
        logger.info(`Generated files are saved in: ${outputDir}`);
        logger.info(`${'='.repeat(80)}\n`);

    } catch (error) {
        logger.error(`Error in iteration ${iteration}:`, error);
        throw error;
    }
}

/**
 * Main execution function with infinite iteration
 */
async function main() {
    let iteration = 1;

    try {
        logger.info('Starting infinite processing loop...');
        logger.info('Press Ctrl+C to stop the process\n');

        while (true) {
            await processIteration(iteration);
            iteration++;

            // Small delay between iterations
            logger.info('Waiting 5 seconds before next iteration...\n');
            await new Promise(resolve => setTimeout(resolve, 5000));
        }
    } catch (error) {
        logger.error('Error in main execution:', error);
        process.exit(1);
    }
}

// Execute main function if this file is run directly
if (require.main === module) {
    main();
}

export { processImage, convertImageWithPrompt, processTwoImages, processIteration, main };
