# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/02_utils.capture_display.ipynb.

# %% auto 0
__all__ = ['PassthroughCapturer', 'svg_to_png', 'display_to_messages']

# %% ../../nbs/02_utils.capture_display.ipynb 2
from IPython.core.displaypub import DisplayPublisher
from IPython.core.display_functions import _new_id
from IPython.display import display, clear_output, Javascript
from IPython import get_ipython
from typing import List
from PIL import Image
from collections import defaultdict
import warnings
import sys
import re
import cairosvg
import xml.etree.ElementTree as ET
from io import BytesIO
import base64

from .misc import nict

# %% ../../nbs/02_utils.capture_display.ipynb 3
class PassthroughCapturer(DisplayPublisher):
    """A DisplayPublisher that stores AND displays!"""

    raw_outputs: List

    def __init__(self):
        self._publisher = None
        self.raw_outputs = []

    def publish(
        self, data, metadata=None, source=None, *, transient=None, update=False
    ):
        self.raw_outputs.append(nict(data=data, metadata=metadata, transient=transient, update=update))
        self._publisher.publish(data, metadata=metadata, transient=transient, update=update)

    def clear_output(self, wait=False):
        self.raw_outputs = []
        self._publisher.clear_output(wait)

    def start(self):
        ip = get_ipython()
        self._publisher = ip.display_pub
        ip.display_pub = self

        # print("starting passthrough: ", id(self._publisher), id(self))

    # Note: stop must be called from the cell that started the
    # passthrough, or from the associated post-cell callback
    def stop(self):
        ip = get_ipython()
        assert self._publisher
        ip.display_pub = self._publisher

    def get_outputs(self, with_js=False):
        coalesced_outputs = []
        display_id_map = defaultdict(list)

        outputs = self.raw_outputs
        if not with_js:
            outputs = [o for o in outputs if 'application/javascript' not in o['data']]

        # dicts are ordered in Python 3.7+
        for output in outputs:
            display_id = output['transient']['display_id'] if output['transient'] else _new_id()
            if not output['update']:
                display_id_map[display_id] += [output["data"]]
            else:
                display_id_map[display_id][-1] = output["data"]

        for display_id, output_list in display_id_map.items():
            coalesced_outputs += output_list

        return coalesced_outputs

# %% ../../nbs/02_utils.capture_display.ipynb 6
from claudette import text_msg, img_msg

# %% ../../nbs/02_utils.capture_display.ipynb 7
def svg_to_png(svg_content, max_size):
    root = ET.fromstring(svg_content)


    def parse_dimension(value):
        if value.endswith('%'):
            return float(value[:-1]) / 100 * max_size
        return float(re.findall(r'[\d.]+', value)[0])

    width = parse_dimension(root.attrib.get('width', str(max_size)))
    height = parse_dimension(root.attrib.get('height', str(max_size)))

    viewBox = root.attrib.get('viewBox')
    if viewBox:
        _, _, vb_width, vb_height = map(float, viewBox.split())
        if width == max_size:
            width = vb_width
        if height == max_size:
            height = vb_height

    scale = min(max_size / width, max_size / height)
    new_width = int(width * scale)
    new_height = int(height * scale)

    # Render PNG
    png =  cairosvg.svg2png(bytestring=svg_content, output_width=new_width, output_height=new_height)
    return png

# %% ../../nbs/02_utils.capture_display.ipynb 9
def display_to_messages(dsp):
    plain_text = dsp.get("text/plain")
    text = list(filter(lambda x:x[1], [
        ("text-html", dsp.get("text/html")),
        ("latex", dsp.get("text/latex")),
        ("markdown", dsp.get("text/markdown")),
        ("json", dsp.get("application/json")),
        ("js", dsp.get("application/javascript"))
    ]))

    images = list(filter(lambda x:x, [ dsp.get("image/png"), dsp.get("image/jpeg") ]))
    svg = dsp.get("image/svg+xml")
    if svg: images.append(svg_to_png(svg, 256))

    messages = [f"<{t[0]}>{t[1]}</{t[0]}>" for t in text]

    if not len(messages):
        messages = [plain_text]

    claudette_msgs = [text_msg(m) for m in messages]
    if images:
        imagebytes = base64.b64decode(images[0]) if not isinstance(images[0], bytes) else images[0]
        claudette_msgs.append(img_msg(imagebytes)) # Grab only one image

    return claudette_msgs
