import base64
import logging

logger = logging.getLogger(__name__)

import time
from contextlib import contextmanager

@contextmanager
def timer(name):
    start = time.time()
    yield
    logger.info(f'[{name}] done in {time.time() - start:.2f} s')

def encode_image(file: str) -> str:
    """将图片编码为 Base64"""
    with timer(f'encode [{file}] to base64'):
        with open(file, 'rb') as f:
            return base64.encodebytes(f.read()).decode('utf8').replace('\n', '').replace('\r', '')

import subprocess
import os
def runcmd(command):
    ret = subprocess.run(command,shell=True,stdout=subprocess.PIPE,stderr=subprocess.PIPE,encoding="utf-8",timeout=1)
    if ret.returncode == 0:
        # print("success:",ret)
        return ret
    else:
        raise ValueError("error: "+ str(ret))

def check_pngquant_bin():
    try:
        runcmd('pngquant --help')
    except Exception as e:
        print('error:', e)
        print('用于图片压缩的工具 pngquant 不存在，请先在 https://pngquant.org/ 上下载并安装！！！')
        os._exit(1)

def compress_image_by_zipflipng(src: str):
    from zopflipng import png_optimize

    with open(src, 'rb') as fp:
        data = fp.read()
    result, code = png_optimize(data)
    # if code ==0 ,png compression success
    if code == 0:
        # save png
        with open(src, 'wb') as f:
            f.write(result)
            f.close()
    else:
        logger.error('optimize by zopfilong failed, ignore this image')

def png2webp(src: str, quality=30):
    """ PNG 转 WEBP
    ---
    目前是压缩最佳方式
    """
    with timer('convert png to webp'):
        from PIL import Image
        im = Image.open(src)
        dst = src.replace('.png', ".webp")
        im.save(dst, "WEBP", quality=quality)
        return dst

def png2avif(src: str, quality: int=30):
    """目前最新的图片格式，据说比 webp 更轻
    但考虑到兼容性，目前还是暂时使用 webp
    ---
    经过测试，极限状态下比 webp 更强，可以用于缩略图
    """
    from PIL import Image
    import pillow_avif
    with timer('convert png to avif'):
        im = Image.open(src)
        dst = src.replace('.png', ".avif")
        im.save(dst,'AVIF', quality=quality)
    return dst

def compress_image(file: str, zop_enable=False) -> str:
    """压缩图片，返回图片路径
    pngquant: 有损压缩，压缩率高，图片会失真，但速度快
    zopflipng: 无损压缩，压缩率低，图片不失真，但速度慢，适合后台慢慢运行
    
    ---
    经过比对，不通过压缩，直接将 png 转变为 webp, 能同时保持图片清晰度及最大压缩率
    """
    import pngquant
    dst = file.replace('.png', '.compressed.png')
    import time
    with timer('Image compress by pngquant'):
        pngquant.quant_image(file, dst=dst)
    if zop_enable:
        with timer('Image decompress by zopflipng'):
            compress_image_by_zipflipng(dst)
    return dst

def make_thumbnail(image: str, quality=10) -> str:
    """ 制作缩略图 """
    return 'data:image/cvif;base64, ' + encode_image(png2avif(image, quality=quality))

from PIL import Image

def handle_image(image: Image.Image, webp_quality=50, avif_quality=10):
    from io import BytesIO
    import base64
    # to webp
    webp_buf = BytesIO()
    image.save(webp_buf, 'WEBP', quality=webp_quality)
    # to avif (thumbnail)
    import pillow_avif as _
    avif_buf = BytesIO()
    image.save(avif_buf, 'AVIF', quality=avif_quality)
    # to base64
    def to_base64(b: BytesIO):
        return base64.b64encode(b.getvalue()).decode('utf8').replace('\n', '').replace('\r', '')
    result = to_base64(webp_buf)
    suffix = 'webp'
    thumbnail = 'data:image/cvif;base64, ' + to_base64(avif_buf)
    return (
        result,
        suffix,
        thumbnail
    )
    

if __name__ == '__main__':
    LOG_FORMAT = "%(asctime)s - %(levelname)s - %(message)s"
    logging.basicConfig(level=logging.INFO, format=LOG_FORMAT)
    check_pngquant_bin()
    # print(encode_image('assets/demo.png'))
    # dst = compress_image('assets/demo.png', zop_enable=True)
    png2webp('assets/demo.png', quality=1)
    def test_avif_to_base64():
        """ 测试 avif 图片转 base64 后，文件大小"""
        avif = png2avif('assets/demo.png', quality=10)
        encoded = encode_image(avif)
        print(f'{len(encoded)/1024} kB')
    test_avif_to_base64()
