import { convertImage } from '../lib/image-converter';
import * as fs from 'fs-extra';
import * as path from 'path';
import dotenv from 'dotenv';

dotenv.config();

const inputDir = path.join(__dirname, '../../input');
const outputDir = path.join(__dirname, '../../generated/clearned');

const comfyUrl = process.env.SERVER1_COMFY_BASE_URL;
const comfyOutputDir = process.env.SERVER1_COMFY_OUTPUT_DIR;

if (!comfyUrl || !comfyOutputDir) {
    console.error("ComfyUI URL or Output Directory is not set in environment variables.");
    process.exit(1);
}

const comfyInputDir = comfyOutputDir.replace("output", "input");

async function processImages() {
    await fs.ensureDir(outputDir);

    const files = await fs.readdir(inputDir);
    let index = 1;

    for (const file of files) {
        const sourceFilePath = path.join(inputDir, file);
        const stats = await fs.stat(sourceFilePath);

        if (stats.isFile()) {
            console.log(`Processing ${file}...`);

            const comfyInputPath = path.join(comfyInputDir, file);

            try {
                // 1. Copy file to ComfyUI input directory
                await fs.copy(sourceFilePath, comfyInputPath);
                console.log(`Copied ${file} to ComfyUI input.`);

                const prompt = "请从图1中提取主要主体，把背景设置为浅灰色，并让主体正面朝向，制作成产品照片。";

                // 2. Call convertImage with correct parameters
                const generatedFilePath = await convertImage(prompt, file, comfyUrl!, comfyOutputDir!);

                if (generatedFilePath && await fs.pathExists(generatedFilePath)) {
                    const outputFilename = `clearned_${index}.png`;
                    const finalOutputPath = path.join(outputDir, outputFilename);

                    // 3. Move the generated file to the final destination
                    await fs.move(generatedFilePath, finalOutputPath, { overwrite: true });
                    console.log(`Saved cleaned image to ${finalOutputPath}`);
                    index++;

                    // 4. Delete the original file from the script's input directory
                    await fs.unlink(sourceFilePath);
                    console.log(`Deleted original file: ${file}`);
                }

                // 5. Clean up the file from ComfyUI input directory
                await fs.unlink(comfyInputPath);
                console.log(`Cleaned up ${file} from ComfyUI input.`);

            } catch (error) {
                console.error(`Failed to process ${file}:`, error);
                // If something fails, make sure to clean up the copied file if it exists
                if (await fs.pathExists(comfyInputPath)) {
                    await fs.unlink(comfyInputPath);
                }
            }
        }
    }
}

processImages().catch(console.error);