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

# %% auto 0
__all__ = ['nict', 'GenericEncoder', 'safe_int', 'strip_traceback_junk', 'parse_cell_line', 'svg_to_png',
           'display_data_to_blocks', 'cell_to_blocks', 'make_cell_from_captured', 'merge_messages']

# %% ../nbs/06_utils.ipynb 1
import pprint

# %% ../nbs/06_utils.ipynb 2
class nict(dict):
    def __init__(self, d=dict(), /, **kwargs):
        if d is None: d = {}
        else:
            assert isinstance(d, dict), f'expected dict, got {type(d)}'
            d = d.copy()

        d.update(kwargs)
        for k,v  in d.items():
            if isinstance(v, dict):
                d[k] = nict(**v)
            elif isinstance(v, list):
                d[k] = [nict(**item) if isinstance(item, dict) else item for item in v]
            elif isinstance(v, tuple):
                d[k] = tuple(nict(**item) if isinstance(item, dict) else item for item in v)
            else:
                d[k] = v

        super().__init__(**d)

    def __getattr__(self, key):
        return self[key] if key in self else super().__getattribute__(key)

    def __setattr__(self, key, value):
        self[key] = value

    def __delattr__(self, key):
        del self[key]

    def __dir__(self):
        return list(self.keys()) + dir(super())

    def _to_dict(self):
        return {k: v._to_dict() if isinstance(v, nict) else v for k, v in self.items()}

# %% ../nbs/06_utils.ipynb 13
import json
from datetime import datetime

# %% ../nbs/06_utils.ipynb 14
class GenericEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, datetime):
            return obj.isoformat()
        try:
            return obj.__dict__
        except AttributeError:
            try:
                return str(obj)
            except:
                return f"<{type(obj).__name__}>"

# %% ../nbs/06_utils.ipynb 15
def safe_int(s, default=0):
    try:
        return int(s)
    except ValueError:
        return default

# %% ../nbs/06_utils.ipynb 16
import re

# %% ../nbs/06_utils.ipynb 17
def strip_traceback_junk(text):
    """Strips colors and other junk from a traceback"""

    # This line just easts up tokens. We don't need it.
    text = text.replace("---------------------------------------------------------------------------", "")
    ansi_escape = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])')
    return ansi_escape.sub('', text)

# %% ../nbs/06_utils.ipynb 19
from claudette import text_msg, img_msg
import base64

# %% ../nbs/06_utils.ipynb 20
def parse_cell_line(line):
    num_cells, is_friendlly, is_reply, is_config = 0, False, False, False
    if line.lstrip('# \t').rstrip().startswith("%%fr"):
        is_friendlly = True
        args = line.split()
        for arg in args[1:]:
            if "+" in arg: num_cells = safe_int(arg[1:])
            if arg == "reply": is_reply = True
            if arg == "config" : is_config = True

    return num_cells, is_friendlly, is_reply, is_config

# %% ../nbs/06_utils.ipynb 21
import cairosvg
import xml.etree.ElementTree as ET

# %% ../nbs/06_utils.ipynb 22
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/06_utils.ipynb 24
def display_data_to_blocks(dsp):
    """
    Converts a display_data (a dict of mime types to data) to a list of claudette messages (text, image, with xml tags)
    """

    # The dsp for an svg might look like this:
    # {
    #     "text/plain": "<SVG object>",
    #     "image/svg+xml": "<svg>....</svg>",
    #     "image/png": "<base64 image>", # or binary string
    #     "image/jpeg": "<base64 image>", # or binary string
    # }
    # Or for markdown:
    # {
    #     "text/plain": "Markdown text",
    #     "text/markdown": "## Markdown text",
    # }

    plain_text = dsp.get("text/plain")

    # Map mime types to their xml tags
    text_types = 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"))
    ]))
    messages = [f"<{t[0]}>{t[1]}</{t[0]}>" for t in text_types]

    images = list(filter(lambda x:x, [ dsp.get("image/png"), dsp.get("image/jpeg") ]))
    # Svg is usually rendered by ipython as an image, but let's handle unrendered svg as well
    svg = dsp.get("image/svg+xml")
    if svg: images.append(svg_to_png(svg, 256))

    # Different text typed outputs come with plain text.
    # We only include the plain text if no other text type is present.
    if not len(messages):
        messages = [plain_text]

    claudette_blocks = [text_msg(m) for m in messages]
    if images:
        image = images[0] # Grab only one image, they are ordered as png, jpeg, rendered svg
        # If the image is a base64 string, convert it to binary
        if isinstance(image, str):
            image = base64.b64decode(image)
        claudette_blocks.append(img_msg(image))

    return claudette_blocks


def cell_to_blocks(cell):
    "Converts a cell and its outputs to a list of claudette messages (text, image, with xml tags)"

    assert "source" in cell
    if cell.cell_type != "code":
        assert not "outputs" in cell or not len(cell.outputs), "Got a non-code cell with outputs??"
        return False, False, text_msg(cell.source), []

    is_reply = False
    is_friendlly = False
    output_blks = []

    lines = cell.source.strip().splitlines()
    if lines:
        _, is_friendlly, is_reply, _ = parse_cell_line(lines[0])

        # Remove the cell magic
        if is_friendlly:
            lines = lines[1:]
        code_text = "\n".join(lines)

        outputs = cell.get("outputs", [])
        for o in outputs:
            if o.output_type == "stream":
                output_blks.append(text_msg(f"<{o.name}>{o.text}</{o.name}>"))
            if o.output_type in ("display_data", "execute_result"):
                # We produce a lot of empty JS outputs when running the chat in nbclassic mode,
                # and they don't get cleaned up in time. Ignore them.
                if "application/javascript" not in o.data or o.data["application/javascript"]:
                    output_blks += display_data_to_blocks(o.data)
            if o.output_type == "error":
                output_blks.append(text_msg("<error>" + strip_traceback_junk("\n".join(o.traceback)) + "</error>"))

        return is_friendlly, is_reply, code_text, output_blks
    return False, False, None, []

# %% ../nbs/06_utils.ipynb 25
from IPython import get_ipython

# %% ../nbs/06_utils.ipynb 26
def make_cell_from_captured(source:str, res, stdout:str, stderr:str, displays, tracebacks):
    ip = get_ipython()
    cell = nict()
    cell.metadata  = {}
    cell.cell_type = "code"
    cell.source    = source
    cell.outputs   = []
    if stdout: cell.outputs.append(nict(output_type="stream", name="stdout", text=stdout))
    if stderr: cell.outputs.append(nict(output_type="stream", name="stderr", text=stderr))
    for d in displays:
        cell.outputs.append(nict(output_type="display_data", data=d))
    if res.result:
        formatted, _ = ip.display_formatter.format(res.result)
        cell.outputs.append(nict(output_type="execute_result", data=formatted))
    for t in tracebacks:
        cell.outputs.append(nict(t, output_type="error"))

    return cell

# %% ../nbs/06_utils.ipynb 27
from typing import List
import claudette

# %% ../nbs/06_utils.ipynb 28
def merge_messages(msgs: List[nict]):
    """
    Takes a list of messages and merges them together, coalescing first by role,
    and then also coalescing text blocks together.
    """

    # first, merge by role
    role_messages = []
    prev_role = None

    for m in msgs:
        if prev_role == m.role: role_messages[-1].content += m.content
        else:                   role_messages.append(m)
        prev_role = m.role

    # Then, for assistant, also coalesce messages together
    for m in role_messages:
        merged_content = []
        for c in m.content:
            if c.type == "text" and merged_content and merged_content[-1].type == "text":
                merged_content[-1].text += ("\n" + c.text)
            else: merged_content.append(c)
        m.content = merged_content
    return role_messages

