// ==UserScript==
// @name         Bing Copilot Image auto-downloader
// @namespace    http://tampermonkey.net/
// @version      0.20
// @license      MIT
// @description  Automatic image downloader for Bing Copilot Image Creator.
// @match        https://copilot.microsoft.com/images/create*?*autosavetimer=*
// @grant        GM_download
// @grant        GM_openInTab
// @grant        GM_getValue
// @grant        GM_setValue
// @grant        GM_deleteValue
// @grant        GM_listValues
// @grant        GM_registerMenuCommand
// @require      http://code.jquery.com/jquery-3.7.1.min.js
// @downloadURL none
// ==/UserScript==
//
// I just pasted this together from things found scattered around the internet.  Starting with: https://github.com/Emperorlou/MidJourneyTools
//
// To enable periodic downloading of newly-created images, go to a 'recent creations' page, and add "&autosavetimer=60" to the URL;
// something like: `https://www.bing.com/images/create/-/1234?autosavetimer=60`.
//
// This implementation is designed to be left unattended - periodically reloading itself.  If you click a link then `autosavetimer=` will
// be removed from the URL and the script will stop working.  Removing `?*autosavetimer=*` from `@match` above _might_ work, but it's not
// well tested at this stage.

var $ = window.jQuery;

(function() {
    'use strict';
    const filenamePrefix = "bing/";
    const downloadables = "img[src$='&pid=ImgGn']";
    const referrerPath = "/images/create/";
    const downloadInterval = 300;
    const sourceContentTag = " #girrc";
    const tagPrefix = "CopilotImageDownloader";
    var pollRate = 60000;
    var activeDownloads = 0;
    var loadErrors = 0;
    var lastReload = 0;
    var lastReauth = Date.now();
    var sourceContent;

    var scanBufferID = null;
    var statusTimeoutID = null;
    var statusBufferID = null;
    var reloadTimerID = null;
    var reauthWindowID = null;

    function jitter(x) {
        return (Math.random() * 0.4 + 0.8) * x;
    }

    $(document).ready(() => {
        var params = new URLSearchParams(window.location.search);
        if (params.get('autosavetimer')) {
            pollRate = params.get('autosavetimer') * 1000;
        }

        scanBufferID = document.createElement("div");
        scanBufferID.setAttribute("id", "scanbuffer");
        scanBufferID.setAttribute("hidden", "");
        document.body.append(scanBufferID);
        statusBufferID = document.createElement("dialog");
        statusBufferID.setAttribute("id", "logmessage");
        statusBufferID.setAttribute("style", "z-index: 100;");
        document.body.append(statusBufferID);
        logger("Automatic image downloader is active.");

        checkMissedDownloads();

        // TODO: Try to figure this out dynamically:
        sourceContent = location.href;

        lastReload = Date.now();
        reloadTimerID = setTimeout(reload, jitter(1000));
        setInterval(function() {
            const timeout = pollRate * 5 + 20 * downloadInterval;
            if (Date.now() - lastReload > timeout) {
                console.log("Reload function seems to have stopped.");
                reload();
            }
        }, pollRate * 1.5);
    });

    GM_registerMenuCommand("Recheck missed downloads", function() {
        const count = checkMissedDownloads();
        if (count > 0) {
            logger("rescheduled " + count + " downloads");
        } else {
            logger("nothing to do");
        }
    });

    GM_registerMenuCommand("Clean up expired records", function() {
        var expired = [];
        for (const tag of GM_listValues()) {
            if (tag.startsWith(tagPrefix + "_info_")) {
                var img = new Image();
                Object.assign(img, GM_getValue(tag));
                if (Date.now() - img.stamp > 31 * 24 * 60 * 60 * 1000) {
                    expired.push(img);
                    GM_deleteValue(img.infoTag);
                    GM_deleteValue(img.busyTag);
                }
            }
        }
        if (expired.length > 0) {
            logger("Found " + expired.length + " old files");
            saveImageLog(expired);
        } else {
            logger("nothing to do");
        }
    });

    GM_registerMenuCommand("Download all records", function() {
        var records = [];
        for (const tag of GM_listValues()) {
            if (tag.startsWith(tagPrefix + "_info_")) {
                var img = new Image();
                Object.assign(img, GM_getValue(tag));
                records.push(img);
            }
        }
        if (records.length > 0) {
            saveImageLog(records);
        } else {
            logger("nothing to do");
        }
    });

    function checkMissedDownloads() {
        var delay = 100;
        var count = 0;
        for (const tag of GM_listValues()) {
            if (tag.startsWith(tagPrefix + "_info_")) {
                var img = new Image();
                Object.assign(img, GM_getValue(tag));
                // TODO: separate this bit out into separate function:
                if (img.scheduleDownload(delay)) {
                    console.log("Rescheduled download of", img.filename);
                    delay += jitter(downloadInterval);
                    count++;
                }
            }
        }
        return count;
    }

    function reload() {
        reloadTimerID = null;
        logger("Rescanning...");
        if (activeDownloads > 0) {
            logger("There are " + activeDownloads + " outstanding.");
        }
        const target = $(scanBufferID);
        var result = target.load(sourceContent + sourceContentTag, function(response, status, xhr) {
            var delay = 100;
            if ( status == "error" ) {
                console.log("problem loading content:", response, status, xhr);
                logger(null);
                if (loadErrors > 0) {
                    logger("previous failures: " + loadErrors);
                }
                logger("problem doing rescan: " + status + ": " + response);
                logger("xhr: " + xhr);
                loadErrors++;
            } else {
                var allImages = $(target.find(downloadables).get().reverse());
                var allRefs = [];
                if (allImages.length < 10) {
                    console.log("Scan buffer doesn't have many images.  Is something wrong?");
                    if (loadErrors == 0) {
                        console.log("all images:", $(target.find("img").get().reverse()));
                    }
                    loadErrors++;
                } else {
                    closeauthwindow();
                    loadErrors = 0;
                }
                for (const i of allImages) {
                    var img = new Image(i);
                    if (img.ref) allRefs.push(img.ref);
                    if (img.scheduleDownload(delay)) {
                        delay += jitter(downloadInterval);
                    }
                }
                if (allRefs.length > 40) {
                    allRefs.sort();
                    // Pick another base URL from which to scan for updates,
                    // in case the initial one eventually expires.
                    // Taking the middle of a sorted list minimises the risk
                    // of accidentally picking up an outlier that doesn't fit
                    // the pattern.
                    var middleref = allRefs[Math.floor(allRefs.length / 2)];
                    middleref = middleref;
                    if (middleref != sourceContent) {
                        sourceContent = middleref;
                        //console.log("new source URL is:", sourceContent);
                    }
                }
                if (activeDownloads == 0) {
                    logger(null);
                }
            }
            if (loadErrors > 3) {
                reauthenticate();
                loadErrors = 1;
            }
            if (reloadTimerID) clearTimeout(reloadTimerID);
            reloadTimerID = setTimeout(reload, jitter(pollRate) + delay);
        });
        if (result.length < 1) {
            console.log("Some kind of load error?");
            reloadTimerID = setTimeout(reload, jitter(pollRate));
        }
        lastReload = Date.now();
    }

    function closeauthwindow() {
        if (reauthWindowID) {
            reauthWindowID.close();
            reauthWindowID = null;
        }
        return true;
    }

    function reauthenticate() {
        if (reauthWindowID && Date.now() - lastReauth < 600000) {
            console.log("too soon to try reauthenticating");
        } else {
            closeauthwindow();
            // TODO: determine this link automatically
            var reauthLink = "https://copilot.microsoft.com/fd/auth/signin" +
                "?action=interactive&provider=windows_live_id" +
                "&cobrandid=03f1ec5e-1843-43e5-a2f6-e60ab27f6b91" +
                "&noaadredir=1&FORM=GENUS1" +
                "&return_url=" + encodeURIComponent("https://copilot.microsoft.com/images/create/");
            // really want the return_url to be something that lets us close the window, but I don't know how to do that.
            console.log("trying to load", reauthLink);
            reauthWindowID = GM_openInTab(reauthLink, { insert: true });
            if (reauthWindowID) {
                lastReauth = Date.now();
            } else {
                console.log("reauth failed");
            }
        }
    }

    function saveImageLog(images) {
        const json = JSON.stringify(images, function(k, v) {
            if (k == 'stamp') return new Date(v).toJSON();
            return v;
        }, 2);
        const blob = encodeURIComponent(json);
        const data = "data:application/json;charset=UTF-8," + blob;
        GM_download({
            url: data,
            name: "image_downloads.txt",
            saveAs: true,
            conflictAction: "uniquify",
            onload: function() {
                console.log("saved images");
            },
            onerror: function(e) {
                console.log("error saving log:", e);
            },
            ontimeout: function(e) {
                console.log("timeout saving log:", e);
            }
        });
    }

    function logger(text) {
        if (statusTimeoutID) {
            statusBufferID.innerHTML = "";
            clearTimeout(statusTimeoutID);
            statusTimeoutID = null;
        }
        if (text) {
            statusBufferID.innerHTML += "<p>" + text + "</p>";
            statusBufferID.show();
        } else {
            statusTimeoutID = setTimeout(function() {
                statusBufferID.innerHTML = "";
                statusBufferID.close();
                statusTimeoutID = null;
            }, 1000);
        }
    }

    class Image {
        constructor(img) {
            // TODO: accept generic result of GM_getValue()
            if (img) {
                this.url = get_download_url(img);
                this.id = get_img_id(this.url);
                this.ref = get_href(img);
                this.alt = img.getAttribute("alt", null);
                this.stamp = Date.now();
                this.done = false;
                if (!this.isSaved) {
                    // TODO: race condition where successful download might be forgotten.
                    GM_setValue(this.infoTag);
                }
            }
        }
        get filename() {
            const src_filename = this.id;
            const pageid = get_page_id(this.ref) || "page";
            const desc = get_page_prompt(this.ref) || this.alt || "image";

            return filenamePrefix + this.id + "_" + pageid + "_" + desc + ".jpg";
        }
        scheduleDownload(delay) {
            if (!this.setBusy()) return false;
            logger("downloading: " + this.filename);
            setTimeout(function() {
                const download = GM_download({
                    url: this.url,
                    name: this.filename,
                    saveAs: false,
                    conflictAction: "uniquify",
                    onload: function() {
                        this.setSaved();
                    }.bind(this),
                    onerror: function(e) {
                        logger("error downloading: " + this.filename, e);
                        this.clearBusy();
                    }.bind(this),
                    ontimeout: function(e) {
                        logger("timeout downloading: " + this.filename, e);
                        this.clearBusy();
                    }.bind(this)
                });
            }.bind(this), delay);
            return true;
        }
        get infoTag() { return tagPrefix + "_info_" + this.id; }
        get busyTag() { return tagPrefix + "_busy_" + this.id; }
        setBusy() {
            if (!this.isReady) return false;
            GM_setValue(this.busyTag, Date.now());
            activeDownloads++;
            return true;
        }
        clearBusy() {
            GM_deleteValue(this.busyTag);
            activeDownloads--;
            if (activeDownloads == 0) {
                logger(null);
            } else if (activeDownloads < 0) {
                logger("Oops, download count underflow!");
                activeDownloads = 0;
            }
        }
        setSaved() {
            this.done = true;
            GM_setValue(this.infoTag, this);
            this.clearBusy();
        }
        get isSaved() {
            const stored = GM_getValue(this.infoTag) || this;
            this.done = stored.done;
            if (this.done) GM_deleteValue(this.busyTag);
            return this.done;
        }
        get isReady() {
            const timeout = 60*1000;
            if (this.isSaved) return false;
            const stamp = GM_getValue(this.busyTag, null);
            if (!stamp) return true;
            if (Date.now() - stamp > timeout) {
                console.log("file has been busy too long (lost event?): " + this.id);
                GM_deleteValue(this.busyTag);
                return true;
            }
            console.log("download already scheduled:", this.id);
            return false;
        }
    }

    // sample: https://tse4.mm.bing.net/th?id=OIG2.AbCdEfGhIjKlMnOp123.&w=100&h=100&c=6&o=5&pid=ImgGn
    function get_img_id(src) {
        var url = new URL(src);
        var id = url.searchParams.get('id') || url.pathname.split('/').pop();
        if (id == null || id.length < 20) {
            console.log("couldn't parse image id from:", src, " got:", id);
        }
        return id;
    }

    // sample: /images/create/kebab-case-prompt/1-0123456789abcedf0123456789abcdef?FORM=GUH2CR
    //         https://copilot.microsoft.com/images/create?q=prompt%20with%20spaces&rt=4&FORM=GENCRE&id=1-0123456789abcedf0123456789abcdef
    function get_page_id(ref) {
        var url = new URL(ref);
        var id = url.searchParams.get('id') || url.searchParams.get('pageId');
        if (id == null) {
            var path = url.pathname.split('/');
            while (path.length && path.shift() != 'create')
                ;
            if (path.length == 2 && path[1].length >= 32) id = path[1];
        }
        if (id == null) {
            console.log("couldn't parse referrer id from:", ref);
        }
        return id;
    }

    // sample: /images/create/kebab-case-prompt/1-0123456789abcedf0123456789abcdef?FORM=GUH2CR
    function get_page_prompt(ref) {
        var url = new URL(ref);
        var q = url.searchParams.get('q');
        if (q == null) {
            var path = url.pathname.split('/');
            while (path.length && path.shift() != 'create')
                ;
            if (path.length == 2 && path[1].length >= 32) q = path[0];
        }
        if (q == null) {
            console.log("couldn't parse referrer prompt from:", ref);
        }
        return q;
    }

    function get_download_url(img) {
      var url = new URL(img.attributes.src.nodeValue);
      url.searchParams.delete("w");
      url.searchParams.delete("h");
      url.searchParams.delete("c");
      url.searchParams.delete("o");
      return url.href;
    }

    function get_href(elem) {
        while (elem) {
            if (elem.hasAttribute('href')) return elem.href;
            elem = elem.parentElement;
        }
        return null;
    }
})();