#!/usr/bin/env node

const fs = require('fs').promises;
const process = require('process');
const puppeteer = require('puppeteer');
const image = require('./image.js');
const uuid = require('uuid').v4;

// Size of the rendering of the web page
const size = { width: 1280, height: 720 };

// Logging helper
function write(stream, chunk, encoding='utf8') {
    return new Promise((resolve, reject) => {
        const errListener = (err) => {
            stream.removeListener('error', errListener);
            reject(err);
        };
        stream.addListener('error', errListener);
        const callback = () => {
            stream.removeListener('error', errListener);
            resolve(undefined);
        };
        stream.write(chunk, encoding, callback);
    });
}

async function print(data) {
    await write(process.stdout, data);
}

async function eprint(data) {
    await write(process.stderr, data);
}

async function println(data) {
    await write(process.stdout, (data ? data : "") + '\n');
}

async function eprintln(data) {
    await write(process.stderr, (data ? data : "") + '\n');
}

async function info(data) {
    await eprintln("\x1b[34;1minfo\x1b[0m\x1b[1m:\x1b[0m " + data);
}

async function warning(data) {
    await eprintln("\x1b[33;1mwarning\x1b[0m\x1b[1m:\x1b[0m " + data);
}

async function error(data) {
    await eprintln("\x1b[31;1merror\x1b[0m\x1b[1m:\x1b[0m " + data);
}

function help() {
    const name = "\x1b[32mlocator\x1b[0m";
    const version = "0.1.0";
    const description = "Helper tool to analyse HTML content produced from marp slides";
    const command = "locator";

    const usage = "\x1b[33mUSAGE:\x1b[0m";
    const args = "\x1b[33mARGUMENTS:\x1b[0m";

    const helpShort = "\x1b[32m-h\x1b[0m";
    const helpLong = "\x1b[32m--help\x1b[0m";

    const inputShort = "\x1b[32m-i\x1b[0m";
    const inputLong = "\x1b[32m--input\x1b[0m";

    const outputShort = "\x1b[32m-o\x1b[0m";
    const outputLong = "\x1b[32m--output\x1b[0m";

    const thresholdShort = "\x1b[32m-t\x1b[0m";
    const thresholdLong = "\x1b[32m--threshold\x1b[0m";

    const forceShort = "\x1b[32m-f\x1b[0m";
    const forceLong = "\x1b[32m--force\x1b[0m";

    const shrinkShort = "\x1b[32m-s\x1b[0m";
    const shrinkLong = "\x1b[32m--shrink\x1b[0m";

    const flattenLong = "\x1b[32m--flatten\x1b[0m";

    const allMasksShort = "\x1b[32m-a\x1b[0m";
    const allMasksLong = "\x1b[32m--all-masks\x1b[0m";

    println(`${name} ${version}
${description}

${usage}
    ${command} -i <HTML-FILE>

${args}
    ${helpShort}, ${helpLong}                Displays this help and quit
    ${inputShort}, ${inputLong} <INPUT>       Path to the HTML input file
    ${outputShort}, ${outputLong} <OUTPUT>     Save mask images and annotations as json in this directory
    ${thresholdShort}, ${thresholdLong} <VALUE>   Threshold for RGB mask computation (between 0 and 1)
    ${shrinkShort}, ${shrinkLong}              Shrink horizontally leaves' bounding boxes
    ${flattenLong}                 Flattens the tree into a list before serializing in JSON
    ${forceShort}, ${forceLong}               Delete the output directory before generating masks it again
    ${allMasksShort}, ${allMasksLong}           Compute and save all masks`
    );
}

async function main() {

    let outputDir = null;
    let filename = null;
    let threshold = undefined;
    let forceMode = false;
    let shrinkBoxes = true;
    let flatten = false;
    let allMasks = false;

    let argIndex = 2;

    while (argIndex < process.argv.length) {

        switch (process.argv[argIndex]) {
            case "-i":
            case "--input":
                filename = process.argv[argIndex + 1];
                argIndex += 2;
                break;

            case "-o":
            case "--output":
                outputDir = process.argv[argIndex + 1];
                argIndex += 2;
                break;

            case "-f":
            case "--force":
                forceMode = true;
                argIndex++;
                break;

            case "-t":
            case "--threshold":
                threshold = parseFloat(process.argv[argIndex + 1]);
                argIndex += 2;
                break;

            case "-s":
            case "--shrink":
                shrinkBoxes = true;
                argIndex++;
                break;

            case "--flatten":
                flatten = true;
                argIndex++;
                break;

            case "-h":
            case "--help":
                help();
                process.exit(0);

            case "-a":
            case "--all-masks":
                allMasks = true;
                argIndex++;
                break;

            default:
                error("unknown option " + process.argv[argIndex]);
                help();
                process.exit(1);
                break;
        }

    }

    if (filename === null) {
        error("program needs a filename argument");
        help();
        process.exit(1);
    }

    if (allMasks && outputDir === null) {
        error("in order to compute all masks, you need to specify an output directory");
        process.exit(1);
    }

    if (threshold !== undefined && isNaN(threshold)) {
        error(tmp + " is not a valid threshold value");
        process.exit(1);
    }

    if (outputDir !== null) {

        if (forceMode === false) {

            try {
                await fs.mkdir(outputDir);
            } catch (e) {
                error("couldn't create directory " + outputDir + ": " + e);
                process.exit(1);
            }

        } else {

            try {
                // If we can just create the directory, we don't need to do anything more
                await fs.mkdir(outputDir);

            } catch (e) {

                // If it fails, we must try to delete it to recreate it because we're in force mode

                try {
                    // Try to access .locator file
                    // If it exists, it is likely that we generated the directory, and therefore, it is safe to delete
                    await fs.access(outputDir + "/.locator", fs.constants.F_OK);
                } catch (e) {
                    // If the file doesn't exist, we don't really know what we would be deleting so we should avoid it
                    error(outputDir + " doesn't seem to have been generated by locator, not deleting and quitting");
                    process.exit(1);
                }

                await fs.rm(outputDir, { recursive: true, force: true });
                await fs.mkdir(outputDir);
            }


        }

        let lock = await fs.open(outputDir + "/.locator", 'a');
        await lock.close();
    }

    // Path to the HTML file to analyse (given as relative path from current directory)
    // We need the full path so that puppeteer is able to access it
    const path = filename.startsWith('/') ? filename : process.cwd() + '/' + filename;

    // Check that the file exists
    try {
        await fs.access(path, fs.constants.F_OK);
    } catch (e) {
        error('no such file: ' + path);
        process.exit(1);
    }

    // Initialize browser
    const browser = await puppeteer.launch();
    const page = await browser.newPage();
    await page.setViewport(size);
    await page.goto('file://' + path);

    let output = [];

    for (let pageIndex = 1;; pageIndex++) {

        let pageIndexStr = "" + pageIndex;

        let pageId = "";
        for (let c of pageIndexStr) {
            pageId += "\\3" + c;
        }

        let root = await page.$('#' + pageId);
        if (root === null) {
            break;
        }

        info("analysing slide " + pageIndex);

        // Hide slides controls
        await page.evaluate(() => {
            for (let elt of document.getElementsByClassName('bespoke-marp-osc')) {
                elt.style.visibility = "hidden";
            }
        });

        if (outputDir !== null) {
            // Take a first screenshot
            await page.screenshot({path: outputDir + '/' + 'screenshot1.png'});
        }

        // Edit the page to shrink elements in order to get better bounding boxes
        if (shrinkBoxes) {
            info("shrinking bounding boxes");
            await addSpan(root);
            info("boundingboxes shrunk");
            eprintln();
        }

        if (outputDir !== null) {
            // Take another screenshot and check the modification we made didn't change the layout of the page
            await page.screenshot({path: outputDir + '/' + 'screenshot2.png'});

            // Compare both screenshots
            let file1 = await fs.readFile(outputDir + '/' + 'screenshot1.png');
            let file2 = await fs.readFile(outputDir + '/' + 'screenshot2.png');
            let filesAreSame = file1.map((x, i) => x === file2[i]).reduce((a, b) => a && b, true);

            if (!filesAreSame) {
                // Check psnr
                let psnr = await image.psnr(outputDir + '/' + 'screenshot1.png', outputDir + '/' + 'screenshot2.png');

                // Crash if they're different
                if (psnr > 0) {
                    warning(filename + " produced slight diff: psnr = " + psnr);
                } else {
                    await error("page edit changed the layout: psnr = " + psnr);
                    process.exit(1);
                }
            }
        }


        if (outputDir !== null) {
            // Produce a screenshot without text

            // Adds a style that makes text invisible
            await page.evaluate(() => {
                let style = document.createElement('style');
                style.innerHTML = "* { color: rgba(0, 0, 0, 0) !important;";
                style.id = "no-text-style";
                document.head.appendChild(style);
            });

            // Perform the screenshot / mask computation
            await page.screenshot({path: outputDir + '/' + 'no-text-' + pageIndex + '.png'});
            await image.segmentationMask(
                outputDir + '/screenshot2.png',
                outputDir + '/no-text-' + pageIndex + '.png',
                outputDir + '/no-text-mask-' + pageIndex + '.png',
                threshold,
            );

            // Restore original element
            await page.evaluate(() => {
                let style = document.getElementById('no-text-style');
                document.head.removeChild(style);
            });
        }


        // Analyse the root and output the result
        info("performing analysis");
        let analyse = await analyseElement(root, page, outputDir, allMasks, threshold);
        info("analysis done");

        // Append the current slide to the output
        if (flatten) {

            analyse = flattenTree(analyse);
            for (let elt of anaylse) {
                output.push(elt);
            }

        } else {

            output.push(analyse);

        }

        // Go to the next page
        await page.evaluate(el => {
            let buttons = document.getElementsByTagName('button');
            for (let button of buttons) {
                if (button.title === "Next slide") {
                    button.click();
                }
            }
        });

        eprintln();
    }

    info("saving output");
    let json = JSON.stringify(output, undefined, 4);
    if (outputDir === null) {
        console.log(json);
    } else {
        await fs.writeFile(outputDir + '/annotations.json', json);
    }

    await browser.close();

}

// Traverses the text nodes of the element and put every text into a single span.
async function addSpan(element) {
    let elts = await element.$$('*');

    for (let elt of elts) {
        let value = await elt.evaluate(el => el.textContent, element);
        if (value !== "") {
            let html = await elt.evaluate(el => {
                if (el.innerHTML.indexOf('<li>') === -1 && el.innerHTML.indexOf('<p>') === -1) {
                    let tmp = el.innerHTML;
                    el.innerHTML = '<span class="sized-span">' + el.innerHTML + '</span>';
                    return [tmp, el.innerHTML];
                } else {
                    return null;
                }
            });

            if (html !== null) {
                await info("replaced \x1b[34m" + html[0].replace(/(\r\n|\n|\r)/gm, "") + "\x1b[0m by \x1b[34m" + html[1].replace(/(\r\n|\n|\r)/gm, "") + '\x1b[0m');
            }
        }
    }

}

// Recursive function to analyse an HTML element.
// The output is written in hierarchy.
async function analyseElement(element, page, outputDir = null, allMasks = false, threshold = undefined) {
    // Get some information on the element
    let tagAttr = await element.getProperty('tagName');
    let tagName = await tagAttr.jsonValue();

    let classAttr = await element.getProperty('className');
    let className = await classAttr.jsonValue();

    let textContent = await element.evaluate(el => el.textContent, element);

    let box = await element.boundingBox();

    // Register it into the return value
    let analyse = {};
    analyse.tag = tagName;
    analyse.class = className;

    analyse.box = box;
    if (box !== null) {
        box.x /= size.width;
        box.width /= size.width;
        box.y /= size.height;
        box.height /= size.height;
    }

    if (outputDir !== null) {
        analyse.uuid = uuid();
    }

    analyse.children = [];

    if (outputDir !== null && allMasks) {

        info("computing screenshots \x1b[34m" + analyse.uuid + "\x1b[0m");

        // Perform a screenshot where the element is hidden
        let previousVisibility = await element.evaluate(el => {
            let previousVisibility = el.style.visibility;
            el.style.visibility = "hidden";
            return previousVisibility;
        });

        await page.screenshot({path: outputDir + '/' + analyse.uuid + '.png'});

        // Compute mask
        await image.segmentationMask(
            outputDir + '/screenshot2.png',
            outputDir + '/' + analyse.uuid + '.png',
            outputDir + '/' + analyse.uuid + '-mask.png',
            threshold,
        );

        await element.evaluate((el, previousVisibility) => {
            el.style.visibility = previousVisibility;
        }, [ previousVisibility ]);
    }

    // Extract the text content if it is a span (we made those spans by ourselves in the addSpan function)
    if (tagName === 'SPAN' && textContent !== "") {
        analyse.text = textContent;
    }

    if (tagName !== "svg") {

        // Select the children of this HTML element.
        let children = await element.$$('> *');

        for (let child of children) {
            // Recursively analyse the children
            analyse.children.push(await analyseElement(child, page, outputDir, allMasks, threshold));
        }

    }

    return analyse;
}

// Flattens the tree into a list.
function flattenTree(input, acc = []) {
    let children = input.children;
    let child = children[0];
    delete input["children"];

    switch (children.length) {
        case 0:
            acc.push(input);
            break;

        case 1:
            if (child.tag === "SPAN" && child.class.indexOf("sized-span") !== -1) {
                delete child["children"];
                acc.push(child);
                break;
            }
            // There is purposefully no break here, if the condition above is false, we want to do the default treatment

        default:
            acc.push(input);
            for (let child of children) {
                flattenTree(child, acc);
            }
            break;
    }

    return acc;
}


main();