import dotenv from 'dotenv';
import * as fs from 'fs/promises';
import * as path from 'path';
import { logger } from '../lib/logger';
import { getPinUrlFromPinterest, downloadImageFromPin } from '../lib/pinterest';
import { generateImage } from '../lib/image-generator-face';
import { callLMStudioAPIWithFile } from '../lib/lmstudio';

dotenv.config();

const MODE: "keywords" | "pinIds" = "keywords";

const KEYWORDS = [
    'a girl in scary forest',
    'a girl in graveyard',
    ''];

const pinIDs = [
    "22377329393970367",
    "18999629674210230",
    "3166662232983784",
    "291537775902572408",
    "2744449769232655",
    "9429480465939847",
    "34058540926328062",
    "1071153092617107265",
    "6825836928646465",
    "1407443629997072",
    "333407178685095962",
    "15833036184288417",
    "6825836928284784",
    "2181499815469509",
    "199706564723106062",
    "1759287348280571",
    "56083957854040032",
    "3025924743999802",
    "2955556001576084",
    "1407443627212889",
    "836965911982723974",
    "97460779431981493",
    "282600945363725869",
    "/59532026387104913",
    "70437490453979",
    "152489137384620437",
    "50947039528553588",
    "73042825197955754",
    "624593042089280419",
    "351912466315529",
    "624030092104188250",
    "21673641951379251",
    "27021666506512503",
    "3377768467678091",
    "985231163409578",
    "17240411068654164"

]
const TARGET_COUNT = Number(process.env.IMAGE_COUNT || 20);
const PROMPT =
    `
    change camera angle to fullbody shot from image1, 
    change background to light gray with faing gradient, 
    change clothes to shite sports bra and shite cotton short pants
    `;
const LmStudioPrompt = `
describe the image in 50 words including, scene, lighting, describe character(s), add some beautiful accent like light, fog, starts, lamps, whatever suits with scene.

then return in this format
{"prompt":""}
`;

type ServerCfg = { baseUrl: string; outputDir: string; inputDir: string };

function getServerConfig(): ServerCfg {
    const candidates = [
        //{ baseUrl: process.env.SERVER1_COMFY_BASE_URL, outputDir: process.env.SERVER1_COMFY_OUTPUT_DIR },
        { baseUrl: process.env.SERVER2_COMFY_BASE_URL, outputDir: process.env.SERVER2_COMFY_OUTPUT_DIR },
    ].filter((s): s is { baseUrl: string; outputDir: string } => !!s.baseUrl && !!s.outputDir);

    if (candidates.length === 0) {
        throw new Error(
            'No ComfyUI server configured. Please set SERVER1_COMFY_BASE_URL/OUTPUT_DIR or SERVER2_COMFY_BASE_URL/OUTPUT_DIR in .env'
        );
    }
    const chosen = candidates[0];
    const inputDir = chosen.outputDir.replace('output', 'input');
    return { ...chosen, inputDir };
}

function sleep(ms: number) {
    return new Promise((res) => setTimeout(res, ms));
}

async function ensureDir(p: string) {
    try {
        await fs.mkdir(p, { recursive: true });
    } catch {
        // ignore
    }
}

async function collectImages(keyword: string, total: number): Promise<string[]> {
    const results: string[] = [];
    // ensure local download dir exists (pinterest.ts also ensures it, but harmless here)
    await ensureDir(path.join(process.cwd(), 'download'));

    if (MODE == "keywords") {
        while (results.length < total) {
            try {
                const pinUrl = await getPinUrlFromPinterest(keyword);
                if (!pinUrl) {
                    logger.warn('No pin URL found, retrying...');
                    await sleep(1500);
                    return [];
                }
                const remaining = total - results.length;
                // attempt to grab up to 5 per pin to reduce churn
                const batchTarget = Math.min(5, remaining);
                const imgs = await downloadImageFromPin(pinUrl, batchTarget);
                if (imgs && imgs.length > 0) {
                    results.push(...imgs);
                    logger.info(`Downloaded ${imgs.length} image(s) from ${pinUrl}.Progress: ${results.length}/${total}`);
                } else {
                    logger.warn(`Pin yielded no downloadable images: ${pinUrl}`);
                }
                await sleep(1000 + Math.random() * 1000);
            } catch (err) {
                logger.error('Error while collecting images:', err);
                await sleep(2000);
            }
        }
    } else if (MODE == "pinIds") {
        while (results.length < total) {
            const shuffledPinIds = pinIDs.slice().sort(() => 0.5 - Math.random());
            const pinUrl = `https://www.pinterest.com/pin/${shuffledPinIds[0]}`;

            try {
                const remaining = total - results.length;
                // attempt to grab up to 5 per pin to reduce churn
                const batchTarget = Math.min(5, remaining);
                const imgs = await downloadImageFromPin(pinUrl, batchTarget);
                if (imgs && imgs.length > 0) {
                    results.push(...imgs);
                    logger.info(`Downloaded ${imgs.length} image(s) from ${pinUrl}.Progress: ${results.length}/${total}`);
                } else {
                    logger.warn(`Pin yielded no downloadable images: ${pinUrl}`);
                }
                await sleep(1000 + Math.random() * 1000);
            } catch (err) {
                logger.error('Error while collecting images:', err);
                await sleep(2000);
            }
        }
    }

    return results.slice(0, total);
}

let imageIndex = 0;
async function processImages(imagePaths: string[], server: ServerCfg) {
    await ensureDir(server.inputDir);

    for (const localImagePath of imagePaths) {
        const response = await callLMStudioAPIWithFile(localImagePath, LmStudioPrompt);
        const prompt = response.prompt;
        const baseName = path.basename(localImagePath);
        const serverInputPath = path.join(server.inputDir, baseName);

        try {
            // copy source image into ComfyUI input dir so workflow node can find it by filename
            await fs.copyFile(localImagePath, serverInputPath);
            logger.info(`Copied ${localImagePath} -> ${serverInputPath}`);

            // Run conversion (sequential to avoid output race conditions)
            const generatedPath = await generateImage(
                `ultra realistic photo, ${prompt}`,
                baseName,
                `monster_${imageIndex}.png`,
                server.baseUrl,
                server.outputDir,
                { width: 1280, height: 720 } // portrait
            );
            logger.info(`Generated image: ${generatedPath}`);
            imageIndex++;
        } catch (err) {
            logger.error(`Failed to convert ${localImagePath}:`, err);
        } finally {
            // cleanup both server input copy and local downloaded file
            try {
                await fs.unlink(serverInputPath);
            } catch { }
            try {
                await fs.unlink(localImagePath);
            } catch { }
            await sleep(500);
        }
    }
}

async function main() {
    const server = getServerConfig();
    const files = await fs.readdir(path.join(process.cwd(), 'generated'));
    const monsterFiles = files.filter((f) => f.startsWith('monster_'));
    if (monsterFiles.length > 0) {
        const latestFile = monsterFiles.sort().pop();
        if (latestFile) {
            const latestIndex = parseInt(latestFile.replace('monster_', '').replace('.png', ''));
            imageIndex = latestIndex + 1;
        }
    }


    // Infinite loop as requested
    while (true) {

        for (const KEYWORD of KEYWORDS) {
            logger.info(`Starting Pinterest image conversion loop for keyword: "${KEYWORD}" (target ${TARGET_COUNT})`);
            try {
                const images = await collectImages(KEYWORD, TARGET_COUNT);
                logger.info(`Collected ${images.length} image(s). Starting conversion...`);

                if (images.length == 0)
                    continue;

                await processImages(images, server);
                logger.info('Iteration completed. Sleeping before next cycle...');
            } catch (err) {
                logger.error('Iteration failed:', err);
            }

            // brief pause between iterations
            await sleep(10000);
        }

    }
}

process.on('unhandledRejection', (reason) => {
    logger.error('UnhandledRejection:', reason);
});

process.on('uncaughtException', (err) => {
    logger.error('UncaughtException:', err);
});

main().catch((e) => {
    logger.error('Fatal error:', e);
});
