from PIL import Image
import base64
import io
import requests
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas


def load_img(imgname):
    if isinstance(imgname, io.BytesIO):
        img = Image.open(imgname)
    elif 'data:image/png;base64,' in imgname:
        img = img_from_base64_str(imgname)
    elif imgname[:4] == 'http':
        rs = requests.get(imgname)
        img = Image.open(io.BytesIO(rs.content))
    else:
        img = Image.open(imgname)
    return img

def get_bytes(img):
    buf = io.BytesIO()
    img.save(buf, format='png')
    buf.seek(0)
    return buf

def fig2img(fig):
    if 0: #v1
        buf = io.BytesIO()
        fig.savefig(buf, format='png')
        buf.seek(0)
        img = Image.open(buf)
    else:
        output = io.BytesIO()
        FigureCanvas(fig).print_png(output)
        output.seek(0)
        img = Image.open(output)
    return img

def fig2b64(fig):
    return img_to_base64_str(fig2img(fig))

#https://stackoverflow.com/questions/48229318/how-to-convert-image-pil-into-base64-without-saving/48229407
#by @Taha Mahjoubi
def img_to_base64_str(img):
    buffered = io.BytesIO()
    img.save(buffered, format="PNG")
    buffered.seek(0)
    img_byte = buffered.getvalue()
    img_str = "data:image/png;base64," + base64.b64encode(img_byte).decode() 
    return img_str

def img_from_base64_str(msg):
    msg = msg.replace("data:image/png;base64,", "")
    msg = base64.b64decode(msg)
    buf = io.BytesIO(msg)
    img = Image.open(buf)
    return img