import puppeteer from 'puppeteer';
import axios from 'axios';
import fs from 'fs';
import dotenv from 'dotenv';

dotenv.config();
import path from 'path';
import { logger } from './logger';

const DOWNLOAD_DIR = path.join(process.cwd(), 'download');

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

async function downloadImage(url: string, keyword: string): Promise<string | null> {
    try {
        const response = await axios.get(url, { responseType: 'arraybuffer' });
        const fileName = `${keyword.replace(/\s/g, '_')}_${new Date().getTime()}${path.extname(url)}`;
        const filePath = path.join(DOWNLOAD_DIR, fileName);
        fs.writeFileSync(filePath, response.data);
        logger.debug(`Image downloaded successfully: ${filePath}`);
        return filePath;
    } catch (error) {
        logger.error(`Failed to download image from ${url}:`, error);
        return null;
    }
}

export async function downloadCivitaiImages(keyword: string, numberOfImages: number = 10): Promise<string[]> {
    const browser = await puppeteer.launch({ headless: process.env.OPEN_BROWSER !== 'true' });
    const page = await browser.newPage();
    const downloadedImagePaths: string[] = [];

    try {
        const searchUrl = `https://civitai.com/search/images?query=${encodeURIComponent(keyword)}`;
        await page.goto(searchUrl, { waitUntil: 'networkidle2' });
        logger.debug(`Navigated to ${searchUrl}`);

        // Wait for the privacy pop-up and click the accept button
        try {
            const acceptButtonSelector = 'button.mantine-Button-filled';
            await page.waitForSelector(acceptButtonSelector, { timeout: 10000 });
            await page.click(acceptButtonSelector);
            logger.debug('Accepted privacy settings.');
        } catch (error) {
            logger.warn('Privacy pop-up not found or could not be dismissed, continuing anyway.');
        }

        // Click the filter button to open the filter options
        const filterButtonSelector = 'button[aria-label="Filters"]';
        await page.waitForSelector(filterButtonSelector);
        await page.click(filterButtonSelector);
        logger.debug('Clicked the filter button.');

        // Wait for the model filter to be available and type in the model name
        const modelFilterSelector = 'input[placeholder="Search models"]';
        await page.waitForSelector(modelFilterSelector);
        await page.type(modelFilterSelector, 'Flux.1 D');
        logger.debug('Typed "Flux.1 D" into the model filter.');

        // Wait for the search results to update and click the correct model
        const modelResultSelector = 'div[aria-label="Flux.1 D"]';
        await page.waitForSelector(modelResultSelector);
        await page.click(modelResultSelector);
        logger.debug('Selected the "Flux.1 D" model.');

        let imageCount = 0;
        while (imageCount < numberOfImages) {
            const imageUrls = await page.$$eval('img', imgs => imgs.map(img => img.src));
            for (const url of imageUrls) {
                if (imageCount >= numberOfImages) break;
                if (url.startsWith('https://image.civitai.com')) {
                    const imagePath = await downloadImage(url, keyword);
                    if (imagePath) {
                        downloadedImagePaths.push(imagePath);
                        imageCount++;
                    }
                }
            }

            if (imageCount < numberOfImages) {
                await page.evaluate('window.scrollTo(0, document.body.scrollHeight)');
                await new Promise(resolve => setTimeout(resolve, 2000)); // Wait for new images to load
            }
        }
    } catch (error) {
        logger.error('An error occurred while downloading images from Civitai:', error);
    } finally {
        await browser.close();
    }

    return downloadedImagePaths;
}
