import time
import json
import sys
import struct
import io

from importlib import import_module
from os.path import join

BASE_DIR = r'd:\obw_benchmark'

lst = [
    # bz2
    'linux-2.6.39.4.tar.bz2',
    'firefox-79.0.linux-i686.tar.bz2',
    'ffmpeg-4.3.1.tar.bz2',
    'gimp-2.10.20.tar.bz2',
    'sde-external-8.56.0-2020-07-05-lin.tar.bz2',

    # lzma
    'linux-5.7.10.tar.xz',
    'linux-2.6.39.4.tar.xz',
    'gcc-9.3.0.tar.xz',
    'Python-3.8.5.tar.xz',
    'firefox-79.0.source.tar.xz',

    # gzip(zlib)
    'linux-5.7.10.tar.gz',
    'linux-2.6.39.4.tar.gz',
    'gcc-9.3.0.tar.gz',
    'Python-3.8.5.tgz',
    'openjdk-14.0.2_linux-x64_bin.tar.gz',
    'postgresql-10.12-1-linux-x64-binaries.tar.gz',
]

def compute_speed(ref, changed):
    speed = ref / changed
    percent = (changed - ref) * 100.0 / ref

    if speed == 1.0:
        return "no change"
    elif speed > 1.0:
        return "%.2fx faster (%+.0f%%)" % (speed, percent)
    else:
        return "%.2fx slower (%+.0f%%)" % (1.0 / speed, percent)

def compare(ref, changed):
    with open(ref) as f:
        o1 = json.load(f)

    with open(changed) as f:
        o2 = json.load(f)  

    for k, v in o2.items():
        if k in o1:
            name1, in_len1, out_len1, time1 = o1[k]
            name2, in_len2, out_len2, time2 = v
            in_len1 = format(in_len1, ',')
            out_len1 = format(out_len1, ',')
            
            r = compute_speed(time1, time2)
            print(f'{name1}')
            print(f'input size: {in_len1}, output size: {out_len1}')
            print(f'best of 5: [{ref}] {time1:.3f} sec -> [{changed}] {time2:.3f} sec, {r}\n')

def _skip_gzip_header(dat):
    FTEXT, FHCRC, FEXTRA, FNAME, FCOMMENT = 1, 2, 4, 8, 16

    offset = 2

    (method, flag,
     _last_mtime) = struct.unpack("<BBIxx", dat[offset:offset+8])
    offset += 8

    if flag & FEXTRA:
        extra_len, = struct.unpack("<H", dat[offset:offset+2])
        offset += 8
    if flag & FNAME:
        while True:
            s = dat[offset:offset+1]
            offset += 1
            if not s or s==b'\000':
                break
    if flag & FCOMMENT:
        while True:
            s = dat[offset:offset+1]
            offset += 1
            if not s or s==b'\000':
                break
    if flag & FHCRC:
        offset += 2
    
    return dat[offset:]

def test_raw(fn):
    # file type
    is_gzip = False   
    t = fn.lower()
    if t.endswith('.xz'):
        module = 'lzma'
    elif t.endswith('.bz2'):
        module = 'bz2'
    elif t.endswith('.gz'):
        module = 'zlib'
        is_gzip = True
    elif t.endswith('.tgz'):
        module = 'zlib'
        is_gzip = True
    else:
        raise Exception('Unknown file type.')

    # import
    m = import_module(module)
    decompress = m.decompress
    print('imported:', module)
    
    # read file
    path = join(BASE_DIR, fn)
    with open(path, 'rb') as f:
        dat = f.read()
  
    if is_gzip:
        dat = _skip_gzip_header(dat)

    in_len = len(dat)
    print(f'read %s, size: %s' % (fn, format(len(dat), ",")))

    # benchmark
    l = []
    for i in range(5):
        if is_gzip:
            t1 = time.perf_counter()
            output = decompress(dat, wbits=-15)
            t2 = time.perf_counter()
        else:
            t1 = time.perf_counter()
            output = decompress(dat)
            t2 = time.perf_counter()

        
        out_len = len(output)
        output = None
        
        l.append(t2-t1)
        print(f'  {t2-t1}')
        
    dat = None
    
    best = min(l)
    print(f'output size: ', format(out_len, ","))
    print(f'best: {best}\n')
    
    return [fn, in_len, out_len, best]

def test_file(fn):
    # file type
    is_gzip = False
    t = fn.lower()
    if t.endswith('.xz'):
        import lzma
        File = lzma.LZMAFile
    elif t.endswith('.bz2'):
        import bz2
        File = bz2.BZ2File
    elif t.endswith('.gz'):
        import gzip
        File = gzip.GzipFile
        is_gzip = True
    elif t.endswith('.tgz'):
        import gzip
        File = gzip.GzipFile
        is_gzip = True
    else:
        raise Exception('Unknown file type.')
   
    # read file
    path = join(BASE_DIR, fn)
    with open(path, 'rb') as f:
        dat = f.read()
    
    in_len = len(dat)
    print(f'read %s, size: %s' % (fn, format(len(dat), ",")))

    # benchmark
    l = []
    for i in range(5):
        b = io.BytesIO(dat)
        if is_gzip:
            f = File(mode='r', fileobj=b)
        else:
            f = File(b, 'r')

        t1 = time.perf_counter()
        output = f.read()
        t2 = time.perf_counter()

        out_len = len(output)
        output = None
        
        l.append(t2-t1)
        print(f'  {t2-t1}')
        
    dat = None
    
    best = min(l)
    print(f'output size: ', format(out_len, ","))
    print(f'best: {best}\n')
    
    return [fn, in_len, out_len, best]

def run_raw(NAME):
    result = {}
    for fn in lst:
        ret = test_raw(fn)
        result[fn] = ret

    with open(NAME+'_raw', 'w') as f:
        json.dump(result, f)

def run_file(NAME):
    result = {}
    for fn in lst:
        ret = test_file(fn)
        result[fn] = ret

    with open(NAME+'_file', 'w') as f:
        json.dump(result, f)

if __name__ == '__main__':
    if len(sys.argv) == 3 and sys.argv[1] == '-raw':
        run_raw(sys.argv[2])
    elif len(sys.argv) == 3 and sys.argv[1] == '-file':
        run_file(sys.argv[2])
    elif len(sys.argv) == 4 and sys.argv[1] == '-compare':
        compare(sys.argv[2], sys.argv[3])
    else:
        s = '''\
Usage:
    benchmark_real.py -raw name
        Test bz2/lzma/zlib decompressor, and dump result to file: name_raw

    benchmark_real.py -file name
        Test calling .read() on BZ2File/LZMAFile/GzipFile object, and dump
        result to file: name_file

    benchmark_real.py -compare name1 name2
        Compare the result stored in files name1 and name2.
'''
        print(s)