import * as fs from 'fs';
import * as path from 'path';
import { callLMStudioAPIWithFile, callLmstudio } from '../lib/lmstudio';
import { embedJsonToPng } from '../lib/util';
import { downloadImagesFromPinterestSearch } from '../lib/pinterest';
import { logger } from '../lib/logger';
import sharp from 'sharp';

const INPUT_DIR = path.join(process.cwd(), 'input');
const OUTPUT_DIR = path.join(process.cwd(), 'generated', 'prompts');

if (!fs.existsSync(OUTPUT_DIR)) {
    fs.mkdirSync(OUTPUT_DIR, { recursive: true });
}

async function generatePromptsForImage(imagePath: string, index: number) {
    const outputFilePath = path.join(OUTPUT_DIR, `cleaned_prompt_${index}.png`);
    logger.info(`Processing image: ${path.basename(imagePath)} -> ${path.basename(outputFilePath)}`);

    try {
        // Step 1: Detect main object and generate colors from the input image
        const colorGenerationPrompt = `
You are a creative assistant. Analyze the provided image.
Identify the main subject product ( not a product name).
Then, list exactly five colors related to this subject:
- Two colors that are common for this object.
- Two colors that are uncommon but plausible.
- One color that is completely crazy or surreal for this object.

Output strictly in this JSON format:
{
  "result": {
    "main_object": "the identified noun",
    "colors": [
      "color1",
      "color2",
      "color3",
      "color4",
      "color5"
    ]
  }
}
`;
        const colorResponse = await callLMStudioAPIWithFile(imagePath, colorGenerationPrompt);
        const { main_object, colors } = colorResponse.result;

        if (!main_object || !Array.isArray(colors) || colors.length !== 5) {
            logger.error(`Failed to get a valid main object and color list for ${imagePath}.`);
            return;
        }

        logger.info(`Main object: "${main_object}", Colors: ${colors.join(', ')}`);

        const prompts: { imagePrompt: string, videoPrompt: string }[] = [];
        const themes = ["special", "unique", "beautiful", "crazy", "funny"];

        // Step 2: Iterate through each color
        for (const color of colors) {
            const randomTheme = themes[Math.floor(Math.random() * themes.length)];
            const pinterestQuery = `${main_object} product photo ${color} background ${randomTheme}`;
            logger.info(`Searching Pinterest for: "${pinterestQuery}"`);

            // Step 3: Get an image from Pinterest
            const downloadedImages = await downloadImagesFromPinterestSearch(pinterestQuery, 1);
            if (downloadedImages.length === 0) {
                logger.warn(`Could not find an image on Pinterest for query: "${pinterestQuery}"`);
                continue;
            }
            const pinterestImagePath = downloadedImages[0];
            logger.info(`Downloaded Pinterest image: ${pinterestImagePath}`);

            // Step 4: Generate a detailed prompt from the Pinterest image
            const imagePromptRequest = `
You are an expert in generating descriptive prompts for image generation models.
Analyze the provided image and describe it in a single, detailed paragraph.
Focus on style, mood, lighting, color palette, sub-objects, and composition.
Do not mention the main object itself. The prompt should be about the scene.

Output strictly in this JSON format:
{
  "result": "your generated prompt here"
}
`;
            const imagePromptResponse = await callLMStudioAPIWithFile(pinterestImagePath, imagePromptRequest);
            const imagePrompt = imagePromptResponse.result;

            if (imagePrompt) {
                logger.info(`Generated image prompt for color ${color}: "${imagePrompt}"`);

                // Step 5: Generate a matching video prompt
                const videoPromptRequest = `
You are a creative director for a short, stylish video ad.
Based on the provided image and the following scene description, generate an attractive video prompt.

Main Subject: ${main_object}
Scene Description: ${imagePrompt}

The video prompt should:
- Be in English and approximately 50 words.
- Describe one clear action involving the main subject.
- Include one specific camera movement (e.g., slow zoom in, orbiting shot, push-in, pull-out).
- Be dynamic and visually appealing.

Output strictly in this JSON format:
{
  "result": "your generated video prompt here"
}
`;
                const videoPromptResponse = await callLMStudioAPIWithFile(pinterestImagePath, videoPromptRequest);
                const videoPrompt = videoPromptResponse.result;

                if (videoPrompt) {
                    logger.info(`Generated video prompt for color ${color}: "${videoPrompt}"`);
                    prompts.push({ imagePrompt, videoPrompt });
                } else {
                    logger.warn(`Failed to generate a video prompt for ${pinterestImagePath}`);
                }
            } else {
                logger.warn(`Failed to generate an image prompt for ${pinterestImagePath}`);
            }
        }

        if (prompts.length === 0) {
            logger.error(`No prompt pairs were generated for ${imagePath}. Aborting.`);
            return;
        }

        // Step 6: Embed all prompts into the original image and save to the new location
        const metadata = {
            prompts: prompts
        };

        // Convert original image to a valid PNG at the output path before embedding
        await sharp(imagePath)
            .toFormat('png')
            .toFile(outputFilePath);

        await embedJsonToPng(outputFilePath, metadata);

        logger.info(`Successfully generated prompts and saved metadata to ${outputFilePath}`);

    } catch (error) {
        logger.error(`An error occurred while processing ${imagePath}:`, error);
    }
}

async function main() {
    try {
        const files = fs.readdirSync(INPUT_DIR);
        const imageFiles = files.filter(file => /\.(png|jpg|jpeg)$/i.test(file));

        if (imageFiles.length === 0) {
            console.log('No images found in the input directory.');
            return;
        }

        for (let i = 0; i < imageFiles.length; i++) {
            const imageFile = imageFiles[i];
            const imagePath = path.join(INPUT_DIR, imageFile);
            await generatePromptsForImage(imagePath, i);
        }

        console.log('All images processed.');
    } catch (error) {
        console.error('An error occurred in the main process:', error);
    }
}

main();