// ==UserScript==
// @name           Huggingface Image Downloader
// @description    Add buttons to quickly download images
// @author         Isaiah Odhner
// @namespace      https://isaiahodhner.io
// @version        1.0
// @license        MIT
// @match          https://stabilityai-stable-diffusion.hf.space/*
// @icon           https://www.google.com/s2/favicons?sz=64&domain=huggingface.co
// @grant          none
// @downloadURL none
// ==/UserScript==

let foundGrid = false;

setInterval(() => {
	const input = document.querySelector('#prompt-text-input input, [name=prompt], [placeholder*="prompt"]');

	const dlButtons = [];
	for (const img of document.querySelectorAll(".grid img")) {
		const existingA = img.parentElement.querySelector("a");
		if (existingA) {
			if (existingA._imgSrc !== img.src) {
				existingA.remove();
				const index = dlButtons.indexOf(existingA);
				if (index > -1) {
					dlButtons.splice(index);
				}
			} else {
				continue; // don't add a duplicate <a> or change the supposed prompt it was generated with
			}
		}

		const a = document.createElement("a");
		a.style.position = "absolute";
		a.style.opacity = "0";
		a.style.top = "0";
		a.style.left = "0";
		a.style.background = "black";
		a.style.color = "white";
		a.style.borderRadius = "5px";
		a.style.padding = "5px";
		a.style.margin = "5px";
		a.style.fontSize = "50px";
		a.style.lineHeight = "50px";
		a.textContent = "Download";
		a._imgSrc = img.src;

		let prompt = input.value;
		// Sanitize for file name, replacing symbols rather than removing them
		prompt = prompt.replace(/\//g, "⧸");
		prompt = prompt.replace(/\\/g, "⧹");
		prompt = prompt.replace(/</g, "ᐸ");
		prompt = prompt.replace(/>/g, "ᐳ");
		prompt = prompt.replace(/:/g, "꞉");
		prompt = prompt.replace(/\|/g, "∣");
		prompt = prompt.replace(/\?/g, "？");
		prompt = prompt.replace(/\*/g, "∗");
		prompt = prompt.replace(/(^|[-—\s(\["])'/g, "$1\u2018");  // opening singles
		prompt = prompt.replace(/'/g, "\u2019");                  // closing singles & apostrophes
		prompt = prompt.replace(/(^|[-—/\[(‘\s])"/g, "$1\u201c"); // opening doubles
		prompt = prompt.replace(/"/g, "\u201d");                  // closing doubles
		prompt = prompt.replace(/--/g, "\u2014");                 // em-dashes
		prompt = prompt.replace(/\.\.\./g, "…");                  // ellipses
		prompt = prompt.replace(/~/g, "\u301C");                  // Chrome at least doesn't like tildes
		prompt = prompt.trim();
		a.download = `SD - ${prompt}.jpeg`;

		a.href = img.src;
		img.parentElement.append(a);
		dlButtons.push(a);

		// Can't be delegated
		a.addEventListener("click", (event) => {
			// Prevent also zooming into the image when clicking Download
			event.stopImmediatePropagation();
		});
	}
	const grid = document.querySelector(".grid");
	if (grid && !foundGrid) {
		foundGrid = true;
		grid.addEventListener("mouseover", (event) => {
			for (const a of dlButtons) {
				a.style.opacity = "0";
			}
			const cell = event.target.closest(".gallery-item");
			if (cell) {
				cell.querySelector("a[download]").style.opacity = "1";
			}
		});
		grid.addEventListener("mouseout", (event) => {
			if (event.target instanceof HTMLImageElement) {
				const cell = event.target.closest(".gallery-item");
				const newCell = event.relatedTarget.closest(".gallery-item");
				if (cell === newCell) return;
				cell.querySelector("a[download]").style.opacity = "0";
			}
		});
		document.addEventListener("mouseleave", (event) => {
			for (const a of document.querySelectorAll(".gallery-item a[download]")) {
				a.style.opacity = "0";
			}
		});
	}
}, 300);