#!/usr/bin/env python3
# ----------------------------------------------------------------------------
# Copyright (c) 2018--, Qurro development team.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the file LICENSE.txt, distributed with this software.
# ----------------------------------------------------------------------------

import json
import copy
import os


def extract_json_from_line(line):
    """Extracts the JSON from a line of JS that we know has a { in it.

       It's assumed that this line has already had .strip() called on it (i.e.
       it has no leading or trailing whitespace).
    """

    # The -1 in the slicing operation below removes the trailing semicolon
    # (since this is a line of JS code)
    return line[line.index("{") : -1]


def get_jsons(main_js_loc, as_dict=True, return_nones=False, json_prefix=""):
    """Extracts the rank/sample plot and count JSONs from a main.js
       file generated by Qurro.

       If any of the JSONs aren't available, raises a ValueError unless
       return_nones is True (in which case nonexistent JSONs will be
       represented by a None in the output).

       Parameters
       ----------
       main_js_loc: str
          The location of a main.js file, which should contain assignments
          to the rankPlotJSON and samplePlotJSON (local) variables. We make
          the assumption that some of the lines in this file are written as
          follows:

          [optional whitespace]var rankPlotJSON = {1};[optional whitespace]
          [optional whitespace]var samplePlotJSON = {2};[optional whitespace]
          [optional whitespace]var countJSON = {3};[optional whitespace]

          where {1} is the rank plot JSON, {2} is the sample plot JSON, and {3}
          is the count JSON (this one doesn't define a plot, it just specifies
          the feature counts for each sample).

          This function just extracts {1}, {2}, and {3} and returns the JSONs
          as either dicts or strings.

       as_dict: bool
          If True, this will load the JSONs as dicts by calling json.loads().
          If False, this will just return the strings.

       return_nones: bool
          If False, this will raise a ValueError if a declaration of either
          JSON is not found in the main.js file.
          If True, this will just return None(s) in place of the corresponding
          missing JSON(s).

       json_prefix: str
          An optional prefix that will be appended to any JSON names we try to
          find. The default value of "" means that we'll look for "normal" JSON
          names (e.g. "var rankPlotJSON = {", "var samplePlotJSON = {",
          "var countJSON = {") -- you can set this to another value to find
          differently named JSONs (for example, if the prefix is "SST", then
          this will look for JSONs of the format "var SSTrankPlotJSON = {").


       Returns
       -------
       (rank_plot_json, sample_plot_json, count_json): each plot could be a
                                                       dict, str, or None

          By default this returns (dict, dict, dict). Passing as_dict=False
          will cause (str, str, str) to be returned instead. Passing
          return_nones=True will allow Nones to be in the output (so
          (None, None, None), (dict, None, None), (None, str, None), etc.
          are all possible output types).
    """
    rank_plot_json_str = None
    sample_plot_json_str = None
    count_json_str = None

    rp_def = "var {}rankPlotJSON = {{".format(json_prefix)
    sp_def = "var {}samplePlotJSON = {{".format(json_prefix)
    c_def = "var {}countJSON = {{".format(json_prefix)
    with open(main_js_loc, "r") as mf:
        for line in mf:
            sline = line.strip()

            if sline.startswith(rp_def):
                rank_plot_json_str = extract_json_from_line(sline)

            elif sline.startswith(sp_def):
                sample_plot_json_str = extract_json_from_line(sline)

            elif sline.startswith(c_def):
                count_json_str = extract_json_from_line(sline)
                break

    if (
        rank_plot_json_str is None
        or sample_plot_json_str is None
        or count_json_str is None
    ):
        if not return_nones:
            raise ValueError(
                "At least one JSON not found in {}.".format(main_js_loc)
            )

    if as_dict:

        def str_to_json(s):
            if s is None:
                return None
            else:
                return json.loads(s)

        return (
            str_to_json(rank_plot_json_str),
            str_to_json(sample_plot_json_str),
            str_to_json(count_json_str),
        )
    else:
        return rank_plot_json_str, sample_plot_json_str, count_json_str


def plot_jsons_equal(json1, json2):
    """Determines if two Vega-Lite JSON objects (dicts) are effectively equal.

       This attempts to "standardize" the name assigned by Altair to the
       datasets (which is somehow based on a hash of the input dataset).
       I've noticed that this can vary even when the data is otherwise
       similar, if the two JSONs only differ by this name then this will
       consider them equal.

       This also attempts to "standardize" the names of selectors generated by
       Altair, in the case that both JSONs have a "selection" property with
       exactly one property within the "selection" property. (This is because
       selector names include a number that is incremented every time Altair
       creates a new selection; see issue #164 on Qurro's GitHub page for
       context.) NOTE that this only bothers trying to standardize selector
       stuff if both JSONs have .data.name properties. In practice every Altair
       JSON we check for Qurro should have both of these, so this shouldn't
       matter. (The only reason the check is done that way is because I'm tired
       and don't want to restructure this function for something it'll never be
       used for in practice.)

       DON'T USE THIS FUNCTION WITH NON-VEGA-LITE JSON OBJECTS -- instead, just
       use something like the == operator.
    """
    # Try to standardize dataset names
    try:
        dn1 = json1["data"]["name"]
        dn2 = json2["data"]["name"]
    except (KeyError, TypeError):
        # At least one of the JSONs didn't have a set data name. Or it was
        # None. However, the inputs could still ostensibly be equal --
        # e.g. they're both {} or both None or something
        return json1 == json2
    # We make the (conservative) assumption that Altair isn't going to generate
    # a dataset with the name "qurro_test_dataname".
    test_dn = "qurro_test_dataname"
    # Avoid modifying json1 and json2 by creating deep copies of them.
    # Use of copy.deepcopy() based on
    # https://www.peterbe.com/plog/be-careful-with-using-dict-to-create-a-copy
    json1_c = copy.deepcopy(json1)
    json2_c = copy.deepcopy(json2)
    json1_c["data"]["name"] = test_dn
    json2_c["data"]["name"] = test_dn
    # https://stackoverflow.com/a/16475444
    json1_c["datasets"][test_dn] = json1_c["datasets"].pop(dn1)
    json2_c["datasets"][test_dn] = json2_c["datasets"].pop(dn2)

    # Try to standardize selector names
    if "selection" in json1_c and "selection" in json2_c:
        if len(json1_c["selection"]) == len(json2_c["selection"]) == 1:
            sn1 = list(json1_c["selection"].keys())[0]
            sn2 = list(json2_c["selection"].keys())[0]
            test_sn = "qurro_test_selectorname"
            # Uses the same principle as above with replacing the datasets.
            json1_c["selection"][test_sn] = json1_c["selection"].pop(sn1)
            json2_c["selection"][test_sn] = json2_c["selection"].pop(sn2)

    return json1_c == json2_c


def try_to_replace_line_json(line, json_type, new_json, json_prefix=""):
    """Attempts to replace a JSON declaration if it's on the line.

       Parameters
       ----------
       line: str
          A line from a JavaScript code file. It's assumed that, if it declares
          a JSON, this declaration will only take up one line (i.e. it
          will be of the form
          "[whitespace?]var [JSON prefix?][JSON name] = {JSON contents};").

          If a replacement is made, everything on and after the { in this line
          will be replaced with the contents of the new JSON, followed by
          ";\n".

       json_type: str
          One of "rank", "sample", or "count". Other values will result in a
          ValueError being thrown.

       new_json: dict
          A JSON to try replacing the current variable declaration (if present)
          on the input line with.

       json_prefix: str (default value: "")
          An optional prefix that will be appended to any JSON names we try to
          replace. If this is anything but "", this *won't replace normal JSON
          lines* (e.g. "var rankPlotJSON = {") -- instead, this will only
          replace lines with the given prefix (e.g. if the prefix is "SST",
          then only JSON lines of the format "var SSTrankPlotJSON = {" will be
          replaced.

       Returns
       -------
       (line, replacement_made): str, bool
          If no replacement was made, replacement_made will be False and line
          will just equal to the input line.

          If a replacement was made, replacement_made will be True and line
          will be equal to the new line with the JSON replaced.
    """

    prefixToReplace = ""
    if json_type == "rank":
        prefixToReplace = "var {}rankPlotJSON = {{"
    elif json_type == "sample":
        prefixToReplace = "var {}samplePlotJSON = {{"
    elif json_type == "count":
        prefixToReplace = "var {}countJSON = {{"
    else:
        raise ValueError(
            "Invalid json_type argument. Must be 'rank', "
            "'sample', or 'count'."
        )

    prefixToReplace = prefixToReplace.format(json_prefix)

    if line.lstrip().startswith(prefixToReplace):
        return (
            (
                line[: line.index("{")]
                + json.dumps(new_json, sort_keys=True)
                + ";\n"
            ),
            True,
        )
    return line, False


def replace_js_json_definitions(
    input_file_loc,
    rank_plot_json,
    sample_plot_json,
    count_json,
    output_file_loc=None,
    json_prefix="",
    verbose=False,
):
    """Writes a version of the input JS file with JSON(s) changed.

       If output_file_loc is None, the input_file_loc will be overwritten.

       You can provide an arbitrary string (well, at least one without any
       whitespace) to ensure that only JSON declarations that start with that
       prefix will be overwritten. For example, you can pass
       json_prefix="SST" to ensure that "var SSTrankPlotJSON = {" lines
       will have their JSONs replaced, but "var rankPlotJSON = {" lines won't
       be replaced.

       If none of the JSONs are different, this won't write anything to the
       output file (or the input file, if output_file_loc is None). This
       function will return 1 in this case, and will return 0 otherwise.

       Note that the JS variable names defined here (rankPlotJSON,
       samplePlotJSON, and countJSON), as well as these variables being defined
       on separate lines of the file, are relied on in the python tests when
       extracting the JSON files from generated main.js files. If you change
       the way these variables are written to in the JS, it may cause the
       python tests to fail.

       The "verbose" flag just determines whether or not to print something
       when trying to go forward with a replacement.
    """

    curr_rank_plot_json, curr_sample_plot_json, curr_count_json = get_jsons(
        input_file_loc, return_nones=True, json_prefix=json_prefix
    )

    # These "diff" boolean variables indicate which of the JSONs are candidates
    # to be replaced.
    diff_rp = not plot_jsons_equal(curr_rank_plot_json, rank_plot_json)
    diff_sp = not plot_jsons_equal(curr_sample_plot_json, sample_plot_json)
    # Since the count JSON isn't a Vega-Lite JSON, we need to just compare it
    # normally using the != operator.
    diff_c = curr_count_json != count_json

    # If straight-up we know that all of the JSONs are equal, then we won't
    # write anything out. Just return 1 immediately.
    if not diff_rp and not diff_sp and not diff_c:
        return 1
    # Similarly, if the input file doesn't contain *any* JSON declarations (not
    # even something like var rankPlotJSON = {};), then we won't be able to
    # replace any of these declarations. So we just return 1 here as well.
    elif (
        curr_rank_plot_json is None
        and curr_sample_plot_json is None
        and curr_count_json is None
    ):
        return 1

    if verbose:
        print(
            "Attempting JSON replacement with prefix {} for {}".format(
                repr(json_prefix), os.path.basename(input_file_loc)
            )
        )

    output_file_contents = ""
    at_least_one_json_changed = False
    with open(input_file_loc, "r") as input_file_obj:
        # read in basic main.js contents. Replace everything after the {
        # in a JSON definition with the actual JSON.
        for line in input_file_obj:
            output_line = line
            changed_yet = False
            if diff_rp:
                output_line, changed_yet = try_to_replace_line_json(
                    line, "rank", rank_plot_json, json_prefix
                )
            if diff_sp and not changed_yet:
                output_line, changed_yet = try_to_replace_line_json(
                    line, "sample", sample_plot_json, json_prefix
                )
            if diff_c and not changed_yet:
                output_line, changed_yet = try_to_replace_line_json(
                    line, "count", count_json, json_prefix
                )
            if changed_yet:
                at_least_one_json_changed = True
            output_file_contents += output_line

    if at_least_one_json_changed:
        if output_file_loc is None:
            output_file_loc = input_file_loc

        with open(output_file_loc, "w") as output_file_obj:
            output_file_obj.write(output_file_contents)
        return 0
    # Let the caller know that nothing was written to the output file location.
    # If this was called on a JS test then this perfectly normal, but if this
    # was called during the execution of Qurro then this is a problem.
    return 1


def check_json_dataset_names(json_dict, *restricted_names):
    """Checks that certain dataset names aren't present in a Vega-Lite JSON.

       This should never really fail, since all of the special dataset names we
       use internally start with "qurro_" (and Altair generates hashes of the
       data somehow to use as dataset names). However, this function is here
       just in the exceedingly rare case that this conflict does actually come
       up (if so, we can figure out a way around it).
    """

    if "datasets" not in json_dict:
        # Should never happen normally but might as well check for this
        raise ValueError("Input JSON doesn't have any datasets defined")

    intersection = set(json_dict["datasets"].keys()) & set(restricted_names)
    if len(intersection) > 0:
        raise ValueError(
            "Found the following disallowed dataset name(s) in a JSON: "
            "{}".format(intersection)
        )
