"""Functions to gather context for your LLM from the current notebook cells, files, documents, urls."""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/01_context.ipynb.

# %% auto 0
__all__ = ['get_notebook_data', 'get_mime_text', 'cell2out', 'cell2xml', 'nb2xml', 'get_notebook_context']

# %% ../nbs/01_context.ipynb 2
import asyncio
from ipykernel.comm import Comm

import nbformat
from fastcore.utils import *
from fastcore.xml import to_xml, Src, Source,Out,Outs,Cell

# %% ../nbs/01_context.ipynb 11
def _notebook_data():
    future = asyncio.Future()
    
    def on_msg(msg):
        if not future.done():
            future.set_result(msg['content']['data'])
    
    comm = Comm(target_name='wordslab_notebook_comm', show_warning=False)
    comm.on_msg(on_msg)
    comm.send({'request': 'get_notebook_data'})

    return future

async def get_notebook_data(timeout=1):
    future = _notebook_data()
    try:
        return await asyncio.wait_for(future, timeout=timeout)
    except asyncio.TimeoutError:
        try:
            future = _notebook_data()
            return await asyncio.wait_for(future, timeout=timeout)
        except asyncio.TimeoutError:   
            raise TimeoutError("Failed to receive notebook context from Jupyterlab frontend: install wordslab-notebooks-lib extension, or increase the timeout parameter in seconds, or try to refresh the web page.")

# %% ../nbs/01_context.ipynb 20
def get_mime_text(data):
    "Get text from MIME bundle, preferring markdown over plain"
    if 'text/markdown' in data: return ''.join(list(data['text/markdown']))
    if 'text/plain' in data: return ''.join(list(data['text/plain']))

# %% ../nbs/01_context.ipynb 21
def cell2out(o):
    "Convert single notebook output to XML format"
    if hasattr(o, 'data'): 
        txt = get_mime_text(o.data)
        if txt: return Out(txt, mime='markdown' if 'text/markdown' in o.data else 'plain')
    if hasattr(o, 'text'):
        txt = o.text if isinstance(o.text, str) else ''.join(o.text)
        return Out(txt, type='stream', name=o.get('name', 'stdout'))
    if hasattr(o, 'ename'): return Out(f"{o.ename}: {o.evalue}", type='error')

# %% ../nbs/01_context.ipynb 22
def cell2xml(cell):
    "Convert notebook cell to concise XML format"
    cts = Source(''.join(cell.source)) if hasattr(cell, 'source') and cell.source else None
    out_items = L(getattr(cell,'outputs',[])).map(cell2out).filter()
    outs = []
    if out_items: outs = Outs(*out_items)
    parts = [p for p in [cts, outs] if p]
    return Cell(*parts, type=cell.cell_type)

# %% ../nbs/01_context.ipynb 23
def nb2xml(nb, until_cell_id):
    cells_xml = []
    for c in nb.cells:
        if c.id == until_cell_id: break
        if c.cell_type in ('code','markdown'):
            cells_xml.append(to_xml(cell2xml(c), do_escape=False))
    return '\n'.join(cells_xml)     

# %% ../nbs/01_context.ipynb 25
async def get_notebook_context(timeout=0.5):
    data = await get_notebook_data(timeout=timeout)
    notebook_content = data["notebook"]
    nb = nbformat.from_dict(notebook_content)
    cell_id = data["cell_id"]
    return nb2xml(nb, cell_id)
