import * as fs from 'fs';
import { promises as fsPromises } from 'fs';
import * as path from 'path';
import * as crypto from 'crypto';
import dotenv from 'dotenv';
import { callLMStudioAPIWithFile } from '../lib/lmstudio';
import { generateImage } from '../lib/image-generator';
import { convertImageWithFile } from '../lib/image-converter';

dotenv.config();

const inputDir = path.resolve('./input');
const outputDir = path.resolve('./generated');

async function main() {
    const comfyBaseUrl = process.env.SERVER1_COMFY_BASE_URL;
    const comfyOutputDir = process.env.SERVER1_COMFY_OUTPUT_DIR;

    if (!comfyBaseUrl || !comfyOutputDir) {
        throw new Error('SERVER1_COMFY_BASE_URL and SERVER1_COMFY_OUTPUT_DIR must be set in .env file');
    }
    const comfyInputDir = comfyOutputDir.replace("output", "input");

    // Ensure output directory exists
    if (!fs.existsSync(outputDir)) {
        fs.mkdirSync(outputDir, { recursive: true });
    }

    const files = fs.readdirSync(inputDir);

    for (const file of files) {
        const imagePath = path.join(inputDir, file);

        // Calculate file hash
        const fileBuffer = fs.readFileSync(imagePath);
        const hashSum = crypto.createHash('sha256');
        hashSum.update(fileBuffer);
        const hexHash = hashSum.digest('hex');

        const genFileName = `${hexHash}_gen.png`;
        const convertedFileName = `${hexHash}_converted.png`;

        const genFilePath = path.join(outputDir, genFileName);
        const convertedFilePath = path.join(outputDir, convertedFileName);

        if (fs.existsSync(genFilePath) && fs.existsSync(convertedFilePath)) {
            console.log(`Skipping ${file} as it has already been processed.`);
            continue;
        }

        console.log(`Processing ${imagePath}...`);

        // 1. Ask LLM to make a prompt for this photo
        const llmPrompt = `
        Extract the following information from the attached image: face detail, hairstyle, eye color, hair color.
        Output instruction:
        { result: "face detail, hairstyle, eye color, hair color" }
        `;
        const extractedInfoJSON = await callLMStudioAPIWithFile(imagePath, llmPrompt);

        if (!extractedInfoJSON) {
            console.error(`Could not get prompt from LLM for ${imagePath}`);
            continue;
        }

        const extractedInfo = extractedInfoJSON.result;

        console.log(`LLM response: ${extractedInfo}`);

        // 2. Generate photo using the prompt
        const prefix = "Highly detailed face portrait of 20 years old girl, light gray background with faint gradient, look at camera in front";
        const finalPrompt = prefix + extractedInfo;

        console.log(`Generating image with prompt: ${finalPrompt}`);

        const generatedImagePath = await generateImage(
            finalPrompt,
            genFileName,
            comfyBaseUrl,
            comfyOutputDir,
            'flux'
        );

        if (generatedImagePath) {
            console.log(`Generated image saved to: ${generatedImagePath}`);

            // 3. Convert the image
            const generatedImageFileName = path.basename(generatedImagePath);
            const originalImageFileName = file;

            const destGeneratedImagePath = path.join(comfyInputDir, generatedImageFileName);
            const destOriginalImagePath = path.join(comfyInputDir, originalImageFileName);

            // Copy files to comfy input directory
            await fsPromises.copyFile(generatedImagePath, destGeneratedImagePath);
            await fsPromises.copyFile(imagePath, destOriginalImagePath);
            console.log(`Copied ${generatedImageFileName} and ${originalImageFileName} to ${comfyInputDir}`);

            const conversionPrompt = "Replace the hairs of image1 with hairs of image2, change face of image1 with face of image2.";
            const convertedResultPath = await convertImageWithFile(
                conversionPrompt,
                generatedImageFileName,
                originalImageFileName,
                comfyBaseUrl,
                comfyOutputDir
            );

            if (convertedResultPath) {
                // Rename the output file to the final converted name
                await fsPromises.rename(convertedResultPath, convertedFilePath);
                console.log(`Converted image saved to: ${convertedFilePath}`);
            } else {
                console.error(`Failed to convert image for ${originalImageFileName}`);
            }

        } else {
            console.error(`Failed to generate image for prompt: ${finalPrompt}`);
        }
    }
}

main().catch(console.error);
