// ==UserScript==
// @name         ImageFX Downloader
// @namespace    http://tampermonkey.net/
// @version      0.3.2
// @description  Downloads ImageFX generations with prompt, seed and generation metadata
// @author       You
// @match        https://labs.google/fx*tools/image-fx*
// @icon         https://www.google.com/s2/favicons?sz=64&domain=labs.google
// @grant        none
// ==/UserScript==

(function() {
    'use strict';

    const CONFIG = {
        enableAutoDownload: true, // Автоматическое скачивание всех генераций
        enableManualDownload: true, // Ручное скачивание качает генерации с метаданными
        enableWideInputField: true, // Широкое поле для промпта
        imageFormat: 'jpg', // Варианты: 'jpg', 'png', 'webp'. Сохранение меты реализовано только для JPG.
        imageQuality: 0.95 // Качество от 0.0 до 1.0 for jpeg/webp, ignored for png
    };

    const originalFetch = window.fetch;
    const piexifScriptUrl = 'https://cdn.jsdelivr.net/npm/piexifjs';
    let piexifLoaded = false;

    const loadPiexif = () => {
        if (piexifLoaded) return Promise.resolve(window.piexif);

        return new Promise(resolve => {
            const script = document.createElement('script');
            script.src = piexifScriptUrl;
            script.onload = () => {
                piexifLoaded = true;
                resolve(window.piexif);
            };
            document.head.appendChild(script);
        });
    };

    const generateFilename = (seed) => {
        const now = new Date();
        const timestamp = now.toISOString().replace(/[:T]/g, '-').split('.')[0];
        return `ImageFX - ${timestamp} - ${seed}.${CONFIG.imageFormat}`;
    };

    const createImageWithMetadata = async (imageData, metadata) => {
        const dataUrl = imageData.startsWith('data:') ? imageData : `data:image/png;base64,${imageData}`;
        const response = await fetch(dataUrl);
        const blob = await response.blob();
        const img = await createImageBitmap(blob);

        const canvas = document.createElement('canvas');
        canvas.width = img.width;
        canvas.height = img.height;
        canvas.getContext('2d').drawImage(img, 0, 0);

        // Only add EXIF metadata for JPEG format
        if (CONFIG.imageFormat === 'jpg') {
            const piexif = await loadPiexif();
            const jpegDataUrl = canvas.toDataURL('image/jpeg', CONFIG.imageQuality);
            const exifObj = piexif.load(jpegDataUrl);

            const utf8Metadata = unescape(encodeURIComponent(JSON.stringify(metadata)));
            exifObj['0th'][piexif.ImageIFD.ImageDescription] = utf8Metadata;

            const exifBytes = piexif.dump(exifObj);
            const newDataUrl = piexif.insert(exifBytes, jpegDataUrl);
            return { dataUrl: newDataUrl, fileName: generateFilename(metadata.seed) };
        } else {
            const mimeType = `image/${CONFIG.imageFormat}`;
            const dataUrl = canvas.toDataURL(mimeType, CONFIG.imageFormat !== 'png' ? CONFIG.imageQuality : undefined);
            return { dataUrl, fileName: generateFilename(metadata.seed) };
        }
    };

    const downloadImage = ({ dataUrl, fileName }) => {
        const byteString = atob(dataUrl.split(',')[1]);
        const buffer = new Uint8Array(byteString.length);

        for (let i = 0; i < byteString.length; i++) {
            buffer[i] = byteString.charCodeAt(i);
        }

        const mimeType = `image/${CONFIG.imageFormat}`;
        const blob = new Blob([buffer], { type: mimeType });
        const link = document.createElement('a');
        link.href = URL.createObjectURL(blob);
        link.download = fileName;
        document.body.appendChild(link);
        link.click();
        link.remove();
        URL.revokeObjectURL(link.href);
    };

    if (CONFIG.enableAutoDownload) {
        const processGeneratedImages = async (data) => {
            if (!data?.imagePanels?.[0]?.generatedImages) return;

            const { generatedImages, prompt } = data.imagePanels[0];
            const batchSize = generatedImages.length;

            for (let index = 0; index < batchSize; index++) {
                const { encodedImage, seed } = generatedImages[index];
                if (!encodedImage) continue;

                const metadata = { prompt, seed, index, batchSize };
                const imageData = await createImageWithMetadata(encodedImage, metadata);
                downloadImage(imageData);
            }
        };

        window.fetch = async function(input, init) {
            const url = typeof input === 'string' ? input : input.url;
            if (!url.includes('v1:runImageFx')) {
                return originalFetch.apply(this, arguments);
            }

            const response = await originalFetch.apply(this, arguments);

            try {
                const data = await response.clone().json();
                if (response.status === 200) {
                    processGeneratedImages(data);
                }
            } catch (err) {}

            return response;
        };
    }

    if (CONFIG.enableManualDownload) {
        const findImageIndex = (images, targetSrc) => {
            if (!targetSrc) return -1;
            const targetPrefix = targetSrc.substring(0, 100);

            return images.findIndex(img =>
                img.src && img.src.substring(0, 100) === targetPrefix
            );
        };

        const findImageContainer = (imgElement) => {
            let container = imgElement;
            let hasSwiper = false;

            while (container && container !== document.body) {
                if (container.classList?.contains('swiper')) {
                    hasSwiper = true;
                    break;
                }
                container = container.parentElement;
            }

            const xpath = hasSwiper ?
                '/html/body/div/div/div/div/div[1]/div[1]/div/div/div[2]' :
                '/html/body/div/div/div/div/div[1]/div[1]/div';

            return {
                container: document.evaluate(xpath, document, null, XPathResult.FIRST_ORDERED_NODE_TYPE, null).singleNodeValue,
                hasSwiper
            };
        };

        const getImageMetadata = (imgElement) => {
            if (!imgElement) return {};

            const { container, hasSwiper } = findImageContainer(imgElement);

            let allImages;
            if (!container) {
                let parentContainer = imgElement.parentElement;
                for (let i = 0; i < 2; i++) {
                    parentContainer = parentContainer?.parentElement;
                }
                allImages = parentContainer ? Array.from(parentContainer.querySelectorAll('img')) : [imgElement];
            } else {
                allImages = hasSwiper ?
                    Array.from(container.querySelectorAll('img')) :
                    Array.from(container.children)
                    .filter(el => el.querySelector('img'))
                    .map(el => el.querySelector('img'));
            }

            const index = hasSwiper ?
                findImageIndex(allImages, imgElement.src) :
                allImages.indexOf(imgElement);

            const promptElement = document.querySelector('span[data-slate-string="true"]');
            const prompt = promptElement?.innerText.trim() || 'Unknown Prompt';

            const seedInput = document.querySelector('#imagefx-seed-input');
            const seed = seedInput?.matches('input[type="number"]') ?
                seedInput.value :
                -1;

            return { prompt, seed, index, batchSize: allImages.length };
        };

        const findDisplayedImage = (buttonElement) => {
            let current = buttonElement;
            let depth = 0;

            while (current && depth < 6) {
                const img = current.querySelector('img');
                if (img) return img;
                current = current.parentElement;
                depth++;
            }
            return null;
        };

        const handleManualDownload = async (e) => {
            e.preventDefault();
            e.stopPropagation();

            const imgElement = findDisplayedImage(e.currentTarget);
            if (!imgElement) {
                console.warn("Image element not found.");
                return;
            }

            const metadata = getImageMetadata(imgElement);
            const response = await fetch(imgElement.src);
            const blob = await response.blob();
            const img = await createImageBitmap(blob);

            const canvas = document.createElement('canvas');
            canvas.width = img.width;
            canvas.height = img.height;
            canvas.getContext('2d').drawImage(img, 0, 0);

            let imageData;
            if (CONFIG.imageFormat === 'jpg') {
                imageData = await createImageWithMetadata(
                    canvas.toDataURL('image/jpeg', CONFIG.imageQuality),
                    metadata
                );
            } else {
                const mimeType = `image/${CONFIG.imageFormat}`;
                const dataUrl = canvas.toDataURL(mimeType, CONFIG.imageFormat !== 'png' ? CONFIG.imageQuality : undefined);
                imageData = { dataUrl, fileName: generateFilename(metadata.seed) };
            }

            downloadImage(imageData);
        };

        const overrideDownloadButtons = () => {
            document.querySelectorAll("button").forEach(button => {
                if (button.textContent.includes("download") && !button.dataset.overridden) {
                    button.dataset.overridden = "true";
                    button.addEventListener("click", handleManualDownload, true);
                }
            });
        };

        const observer = new MutationObserver(overrideDownloadButtons);
        observer.observe(document.body, { childList: true, subtree: true });

        overrideDownloadButtons();
    }

    if (CONFIG.enableWideInputField) {
        const applyWideStyling = () => {
            const selectors = {
                inputZone: 'div.sc-3943c7d6-2.gJJoYK',
                windowBlock: '.sc-51eeec5c-1'
            };

            Object.values(selectors).forEach(selector => {
                const element = document.querySelector(selector);
                if (element) element.style.maxWidth = '100%';
            });
        };

        applyWideStyling();
    }
})();