# Copyright (c) Facebook, Inc. and its affiliates.
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import uuid
import html
import typing as t
import json
from pathlib import Path

import IPython.display
from ipykernel.comm import Comm
from . import experiment as exp
from .render import escapejs, make_experiment_standalone_page


class GetSelectedFailure(Exception):
    pass


class NotebookJSBundleInjector:
    injected = False

    @classmethod
    def ensure_injected(cls) -> None:
        if cls.injected:
            return
        bundle = Path(__file__).parent / "static" / "built" / "hiplot.bundle.js"
        IPython.display.display(IPython.display.Javascript(bundle.read_text("utf-8")))
        cls.injected = True


def jupyter_make_full_width(content: str) -> str:
    w_id = f"wrap_html_{uuid.uuid4().hex[:6]}"
    return f"""
<div id="{w_id}">{content}</div>
<script type="text/javascript">
(function() {{
    const elem = document.getElementById({json.dumps(w_id)});
    elem.style.width = "100vw";
    const removeElems = elem.parentElement.parentElement.getElementsByClassName("prompt");
    for (var i = 0; i < removeElems.length; ++i) {{
        removeElems[i].remove();
    }}
    elem.parentElement.style.overflowX = "visible";

    const scale_to_100pct_screen = function() {{
        elem.style.marginLeft = '0px';
        elem.style.marginLeft = (- elem.getBoundingClientRect().x) + 'px';
    }};
    window.addEventListener('resize', function() {{
        scale_to_100pct_screen();
    }});
    scale_to_100pct_screen();
}})();
</script>
"""


class IPythonExperimentDisplayed(exp.ExperimentDisplayed):
    def __init__(self, xp: exp.Experiment, comm_name: str) -> None:
        self._exp = xp
        self._num_recv = 0
        self._selected_ids: t.List[str] = []
        self._last_data_per_type: t.Dict[str, t.Any] = {}

        def target_func(comm: Comm, open_msg: t.Dict[str, t.Any]) -> None:  # pylint: disable=unused-argument
            # comm is the kernel Comm instance
            # msg is the comm_open message

            # Register handler for later messages
            @comm.on_msg  # type: ignore
            def _recv(msg: t.Dict[str, t.Any]) -> None:
                self._num_recv += 1
                msg_data = msg["content"]["data"]
                print(msg_data)
                self._last_data_per_type[msg_data["type"]] = msg_data["data"]

        try:
            ip: Any = get_ipython()  # type: ignore  # pylint: disable=undefined-variable
            ip.kernel.comm_manager.register_target(comm_name, target_func)
        except NameError:  # NameError: name 'get_ipython' is not defined
            # We are not in an ipython environment - for example in testing
            pass

    no_data_received_error = GetSelectedFailure(
        """No data received from the front-end. Please make sure that:
    1. You don't call "get_selected" on the same cell
    2. The interface has loaded
    3. You are in a Jupyter notebook (Jupyter lab is *not* supported)"""
    )

    def get_selected(self) -> t.List[exp.Datapoint]:
        last_selection = self._last_data_per_type.get("selection")
        if last_selection is None:
            raise self.no_data_received_error
        selected_set = set(last_selection["selected"])
        datapoints = [i for i in self._exp.datapoints if i.uid in selected_set]
        assert len(datapoints) == len(selected_set)
        return datapoints

    def get_brush_extents(self) -> t.Dict[str, t.Dict[str, t.Any]]:
        last_msg = self._last_data_per_type.get("brush_extents")
        if last_msg is None:
            raise self.no_data_received_error
        return last_msg  # type: ignore


def display_exp(
        xp: exp.Experiment,
        force_full_width: bool = False,
        store_state_url: t.Optional[str] = None
) -> IPythonExperimentDisplayed:
    comm_id = f"comm_{uuid.uuid4().hex[:6]}"
    displayed_xp = IPythonExperimentDisplayed(xp, comm_id)
    options: t.Dict[str, t.Any] = {
        'experiment': xp._asdict()
    }
    if store_state_url is not None:
        options.update({"persistent_state_url_prefix": store_state_url})
    else:
        options.update({"persistent_state": None})
    index_html = make_experiment_standalone_page(options=options)
    # Remove line that references the script bundle - prevents an HTTP error in the notebook
    index_html = index_html.replace('src="static/built/hiplot.bundle.js"', '')
    index_html = index_html.replace(
        "/*ON_LOAD_SCRIPT_INJECT*/",
        f"""/*ON_LOAD_SCRIPT_INJECT*/
const comm_id = {escapejs(comm_id)};
try {{
    console.log("Setting up communication channel with Jupyter: ", comm_id);
    var comm = Jupyter.notebook.kernel.comm_manager.new_comm(comm_id, {{'type': 'hello'}});
    Object.assign(options, {{"comm": comm}});
}}
catch(err) {{
    console.warn('Unable to create Javascript <-> Python communication channel' +
        ' (are you in a Jupyter notebook? Jupyter labs is *not* supported!)');
}}
        """)

    if force_full_width:
        index_html = jupyter_make_full_width(index_html)
    NotebookJSBundleInjector.ensure_injected()
    IPython.display.display(IPython.display.HTML(index_html))
    return displayed_xp
