#| default_exp thumb

#|export
import io, base64, re
import nbformat
from PIL import Image
from fastcore.foundation import L
from fastcore.test import test_fail
from io import BytesIO

#|export
def _decode_img(base64_str):
    img = Image.open(BytesIO(base64.b64decode(base64_str)))
    if img.mode == "RGBA":
        img = img.convert("RGB")
    return img

def get_img(nb_path, label='thumbnail'):
    "Get image from notebook with a quarto cell directive with `#|label: {label}`"
    out_plots = None
    lbl = re.compile(f'#\|(\s*)label:(\s*){label}')
    nb = nbformat.read(open(nb_path), as_version=4)
    for cell in nb.cells:
        if lbl.search(cell.source):
            out_plots = [x for x in cell.outputs if x.output_type == "display_data" and 'data' in x]
            if out_plots:
                data = out_plots[0]['data']
                if 'image/png' not in data: 
                    raise Exception(f'{nb_path}: thumbnails are only supported for `image/png`, found {data.keys()}')
                return _decode_img(data['image/png'])
            else:
                raise Exception(f'{nb_path}: cell with `#|label: {label}` does not have an output type of `display_data`')
    if out_plots is None:
        raise Exception(f'{nb_path} does not contain a cell with `#|label: {label}`')

img = get_img('test_nbs/geom_col.ipynb', label='two_variable_bar_plot')
img

# test if you provide a label that doesn't exist
test_fail(get_img, contains='test_nbs/geom_col.ipynb does not contain a cell with `#|label: does_not_exist`', 
          args=('test_nbs/geom_col.ipynb', 'does_not_exist'))

# test if you label a cell that doesn't have a plot like a dataframe
test_fail(get_img, contains='test_nbs/geom_col.ipynb: cell with `#|label: no_plot` does not have an output type of `display_data`',
          args=('test_nbs/geom_col.ipynb', 'no_plot'))

# make sure image is a PngImageFile
assert isinstance(img, Image.Image)

#|export
def img2thumb(img:Image.Image, size=(260,260)):
    "Convert image to thumbnail."
    thumb_size = size[0] * 2, size[1] * 2
    img.thumbnail(thumb_size)
    return img

img2thumb(img)

#|export
def nb2thumb(nb_path, label='thumbnail', size=(260,260)) -> Image.Image:
    "Extract thumbnail corresponding to the cell with the comment `#|label: {label}` from a notebook."
    img = get_img(nb_path=nb_path, label=label)
    return img2thumb(img, size=size)

thumb = nb2thumb('test_nbs/geom_col.ipynb', label='two_variable_bar_plot')
thumb

#|hide
t1 = nb2thumb('test_nbs/geom_col.ipynb', label='two_variable_bar_plot')
t2 = nb2thumb('test_nbs/PlotnineAnimation.ipynb', label='spiral')
t3 = nb2thumb('test_nbs/geom_density.ipynb', label='density_curve')
t4 = nb2thumb('test_nbs/geom_map.ipynb', label='map')
t5 = nb2thumb('test_nbs/geom_segment.ipynb', label='ranges')
t6 = nb2thumb('test_nbs/geom_segment.ipynb', label='rank')

images = L([t1, t2, t3, t4, t5, t6])