#!/usr/bin/env python
# This file is part of Xpra.
# Copyright (C) 2014-2019 Antoine Martin <antoine@xpra.org>
# Xpra is released under the terms of the GNU GPL v2, or, at your option, any
# later version. See the file COPYING for details.

import sys
import os.path
import time
import binascii

from PIL import Image           #@UnresolvedImport
from io import BytesIO


def do_test_encode(rgb_data, w, h, encodings=("png", "png/P", "png/L", "jpeg", "webp"),
                   N=5, Q=(0, 50, 100), S=(0, 1, 50, 90, 100), has_alpha=False):
    from xpra.codecs.pillow.encoder import encode
    from xpra.codecs.image_wrapper import ImageWrapper
    image = ImageWrapper(0, 0, w, h, rgb_data, "BGRA", 32, w*4, planes=ImageWrapper.PACKED, thread_safe=True)
    #buf = "\0" * (w*h*4)
    #buf = get_source_data(w*h*4)
    for encoding in encodings:
        Q_options = Q
        if encoding in ("png", "png/P", "png/L"):
            Q_options = [-1]
        S_options = S
        if encoding=="webp":
            S_options = [-1]
        if encoding=="jpeg":
            S_options = [0, -1]
        for q in Q_options:
            for s in S_options:
                #print("test_encode() quality=%s, speed=%s" % (q, s))
                coding = encoding
                start = time.time()
                for _ in range(N):
                    data = encode(coding, image, q, s, True)
                #def compress(pixels, width, height, quality=50, speed=50):
                end = time.time()
                mps = w*h*N/(end-start)/1024/1024
                cdata = data[1]
                ratio = 100.0 * len(cdata) / len(rgb_data)
                print("%s : %s MPixels/s  %s%% compression ratio : %sx%s to: %s KBytes (quality=%s, speed=%s) %s times in %sms average" % \
                      (coding.ljust(5), ("%.1f" % mps).rjust(5), ("%.1f" % ratio).rjust(5), str(w).rjust(4), str(h).ljust(4), str(len(cdata)/1024).rjust(6), str(q).rjust(3), str(s).rjust(3),
                       N, str(int((end-start)*1000.0/N)).rjust(5)))

def test_files(filenames, extensions=(".png", ".jpg"), recurse=True):
    for x in filenames:
        if os.path.isdir(x):
            if recurse:
                test_files(os.listdir(x), extensions, recurse)
            continue
        try:
            img = Image.open(x)
            img.load()
            has_alpha = img.mode=="RGBA"
            img = img.convert("RGBA")
            rgb_data = img.tobytes("raw", img.mode)
            w, h = img.size
            assert len(rgb_data)==w*h*len(img.mode)
            print("testing with file: %s (%sx%s)" % (x, w, h))
            do_test_encode(rgb_data, w, h, has_alpha=has_alpha)
        except Exception as e:
            print("error on %s: %s" % (x, e))

def test_encode():
    img_data = binascii.unhexlify("[encoded data, 9760 characters]")
    buf = BytesIO(img_data)
    img = Image.open(buf)
    img = img.convert("RGBA")
    rgb_data = img.tobytes("raw", img.mode)
    w, h = img.size
    do_test_encode(rgb_data, w, h)

def test_file_args():
    if len(sys.argv)>0:
        filenames = sys.argv[1:]
        print("testing with folders/files: %s" % ", ".join(filenames))
        test_files(filenames)


def main():
    test_encode()
    test_file_args()


if __name__ == "__main__":
    main()
