import os
from copy import deepcopy

from nbformat import NO_CONVERT, read

from voila.execute import executenb

BASE_DIR = os.path.dirname(__file__)
WIDGET_MIME_TYPE_VIEW = "application/vnd.jupyter.widget-view+json"
WIDGET_MIME_TYPE_STATE = "application/vnd.jupyter.widget-state+json"


# based on nbconvert.preprocessors.tests.test_execute.TestExecute
# we cannot import it because pytest would then also execute those tests
def normalize_output(output):
    """
    Normalizes outputs for comparison.
    """
    output = dict(output)
    if "metadata" in output:
        del output["metadata"]
    # tracebacks can be different per installation/python version
    if "traceback" in output:
        del output["traceback"]
    if "application/vnd.jupyter.widget-view+json" in output.get("data", {}):
        output["data"]["application/vnd.jupyter.widget-view+json"][
            "model_id"
        ] = "<MODEL_ID>"


def normalize_outputs(outputs):
    for output in outputs:
        normalize_output(output)


def test_execute_output():
    path = os.path.join(BASE_DIR, "notebooks/output.ipynb")
    nb = read(path, NO_CONVERT)
    nb_voila = deepcopy(nb)
    executenb(nb_voila, timeout=30)

    widget_states = nb.metadata.widgets[WIDGET_MIME_TYPE_STATE]["state"]
    widget_states_voila = nb_voila.metadata.widgets[WIDGET_MIME_TYPE_STATE]["state"]

    for cell_voila, cell in zip(nb_voila.cells, nb.cells):
        for output_voila, output in zip(cell_voila.outputs, cell.outputs):
            if "data" in output and WIDGET_MIME_TYPE_VIEW in output["data"]:
                widget_id = output["data"][WIDGET_MIME_TYPE_VIEW]["model_id"]
                widget_id_voila = output_voila["data"][WIDGET_MIME_TYPE_VIEW][
                    "model_id"
                ]
                widget_state = widget_states[widget_id]
                widget_state_voila = widget_states_voila[widget_id_voila]
                # if the widget is an output widget, it has the outputs, which we also check
                assert normalize_outputs(
                    widget_state.state.get("outputs", [])
                ) == normalize_outputs(widget_state_voila.state.get("outputs", []))
            normalize_output(output_voila)
            normalize_output(output)
            assert output_voila == output
