import puppeteer from 'puppeteer';
import fs from 'fs';
import path from 'path';
import { logger } from './logger';

const downloadPath = path.resolve(__dirname, '../../download');

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

async function downloadImage(url: string, filepath: string) {
    const response = await fetch(url);
    const buffer = await response.arrayBuffer();
    fs.writeFileSync(filepath, Buffer.from(buffer));
}

export async function downloadPinterestImages(keyword: string, numberOfPages: number): Promise<string[]> {
    const browser = await puppeteer.launch({ headless: true });
    const page = await browser.newPage();
    await page.setUserAgent('Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/125.0.0.0 Safari/537.36');
    await page.setViewport({ width: 1920, height: 1080 });

    const encodedKeyword = encodeURIComponent(keyword);
    await page.goto(`https://www.pinterest.com/search/pins/?q=${encodedKeyword}`, { waitUntil: 'networkidle2' });

    logger.debug('Searching for:', keyword);

    try {
        await page.waitForSelector('img[src*="i.pinimg.com"]', { timeout: 10000 });
    } catch (error) {
        logger.error('Could not find images on the page. Pinterest might have changed its layout or is blocking the scraper.');
        await browser.close();
        return [];
    }

    let imageCount = 0;
    const downloadedUrls = new Set<string>();
    const downloadedImagePaths: string[] = [];

    for (let i = 0; i < numberOfPages; i++) {
        logger.debug(`Scraping page ${i + 1}...`);
        try {
            const imageUrls = await page.evaluate(() => {
                const images = Array.from(document.querySelectorAll('img[src*="i.pinimg.com"]'));
                return images.map(img => {
                    const image = img as HTMLImageElement;
                    const srcset = image.srcset;
                    if (srcset) {
                        const sources = srcset.split(',').map(s => s.trim().split(' ')[0]);
                        return sources[sources.length - 1];
                    }
                    return image.src;
                }).filter(Boolean);
            });

            for (const url of imageUrls) {
                if (!downloadedUrls.has(url)) {
                    downloadedUrls.add(url);
                    const extension = path.extname(new URL(url).pathname) || '.jpg';
                    const filename = `${keyword.replace(/ /g, '_')}_${imageCount++}${extension}`;
                    const filepath = path.join(downloadPath, filename);
                    logger.debug(`Downloading ${url} to ${filepath}`);
                    await downloadImage(url, filepath);
                    downloadedImagePaths.push(filepath);
                }
            }

            const previousHeight = await page.evaluate('document.body.scrollHeight');
            await page.evaluate('window.scrollTo(0, document.body.scrollHeight)');

            try {
                await page.waitForFunction(`document.body.scrollHeight > ${previousHeight}`, { timeout: 10000 });
            } catch (e) {
                logger.debug('No more content to load.');
                break;
            }

            await new Promise(resolve => setTimeout(resolve, 1000)); // Small delay for images to render
        } catch (error) {
            logger.error(`An error occurred while scraping page ${i + 1}:`, error);
            logger.debug('Stopping the scraping process due to an error.');
            break;
        }
    }

    await browser.close();
    logger.debug('Done.');
    return downloadedImagePaths;
}

export async function downloadImagesFromPinterestPin(pinUrl: string, scrollCount: number = 5): Promise<string[]> {
    const browser = await puppeteer.launch({ headless: false });
    const page = await browser.newPage();
    await page.setUserAgent('Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/125.0.0.0 Safari/537.36');
    await page.setViewport({ width: 1920, height: 1080 });

    const downloadedImagePaths: string[] = [];
    const downloadedUrls = new Set<string>();

    try {
        await page.goto(pinUrl, { waitUntil: 'networkidle2' });
        logger.debug(`Navigated to pin page: ${pinUrl}`);

        const pageImageUrls = new Set<string>();

        for (let i = 0; i < scrollCount; i++) {
            logger.debug(`Scrolling page ${i + 1}/${scrollCount} for ${pinUrl}...`);
            const imageUrlsOnPage = await page.evaluate(() => {
                const images = Array.from(document.querySelectorAll('img[src*="i.pinimg.com"]'));
                const urls = images.map(img => {
                    const image = img as HTMLImageElement;
                    if (image.srcset) {
                        const sources = image.srcset.split(',').map(s => s.trim().split(' '));
                        const fourXSource = sources.find(s => s[1] === '4x');
                        if (fourXSource) {
                            return fourXSource[0];
                        }
                    }
                    return null; // Return null if no 4x source is found
                });
                return urls.filter((url): url is string => url !== null); // Filter out nulls and assert type
            });

            for (const url of imageUrlsOnPage) {
                pageImageUrls.add(url);
            }

            const previousHeight = await page.evaluate('document.body.scrollHeight');
            await page.evaluate('window.scrollTo(0, document.body.scrollHeight)');

            try {
                await page.waitForFunction(`document.body.scrollHeight > ${previousHeight}`, { timeout: 10000 });
            } catch (e) {
                logger.debug('No more content to load on this page.');
                break;
            }
            await new Promise(resolve => setTimeout(resolve, 1000)); // Wait for images to load
        }

        let imageCount = 0;
        for (const imageUrl of pageImageUrls) {
            if (!downloadedUrls.has(imageUrl)) {
                downloadedUrls.add(imageUrl);
                const pinIdMatch = pinUrl.match(/\/pin\/(\d+)\/?/);
                const basePinId = pinIdMatch ? pinIdMatch[1] : `image_${Date.now()}`;

                const extension = path.extname(new URL(imageUrl).pathname) || '.jpg';
                const filename = `${basePinId}_related_${imageCount++}${extension}`;
                const filepath = path.join(downloadPath, filename);

                logger.debug(`Downloading ${imageUrl} to ${filepath}`);
                await downloadImage(imageUrl, filepath);
                downloadedImagePaths.push(filepath);
            }
        }

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

    await browser.close();
    logger.debug('Done processing pin.');
    return downloadedImagePaths;
}
