#!/usr/bin/python
"""
  Time how long to read firmware module - with or without compression
  gene c/2023-07-02
"""
import os
import sys
import time
import lzma
import zstandard
#import pdb

def open_file(path):
    """
    Open file in binary read mode return file handle
    """
    # pylint: disable=R1732
    try:
        fobj = open(path, "rb")

    except OSError as err:
        print(f'Error opening {path} : {err}')
        fobj = None

    return fobj

class File_Read():
    """
    Simple file reader
    """
    # pylint: disable=R0903
    def __init__(self, fpath):
        self.fpath = fpath
        self.base = None
        self.fext = None
        self.data = None
        self.compress = None

        path_exists = os.path.exists(fpath)
        if path_exists and os.path.isfile(fpath):
            # if extension unknown we assume uncomopressed
            (self.base, self.fext) = os.path.splitext(fpath)
            known_ext = ['.zst', '.xz']
            if self.fext and self.fext in known_ext:
                self.path_ok = True
            else:
                self.base = self.fpath
                self.fext = ''
        else:
            print (f'Bad Module file: {fpath}')

    def read(self):
        """
        Read a module
        """

        fobj = open_file(self.fpath)
        if fobj:
            raw_data = fobj.read()
            fobj.close()
        else:
            return

        match self.fext:
            case '.zst':
                dctx = zstandard.ZstdDecompressor()
                self.data = dctx.decompress(raw_data)
                self.compress = 'zst'
            case '.xz':
                self.data = lzma.decompress(raw_data)
                self.compress = 'xz'
            case _:
                self.data = raw_data
                self.compress = 'no'

def scan_dir(tdir):
    """
    wrapper around scan handling exceptions
    """
    scan = None
    if os.path.exists(tdir) and os.path.isdir(tdir) :
        try:
            scan = os.scandir(tdir)
        except OSError as _error:
            scan = None
    return scan



def files_in_dir(path):
    """
    Return list of file paths
    ignore directories or links
    """
    scan = scan_dir(path)
    flist = []
    if scan:
        for item in scan:
            if item.is_file():
                flist.append(item.path)

    return flist

def file_size(path):
    """ get file size """
    try:
        fstat = os.stat(path)
        size = fstat.st_size
    except OSError as _error:
        size = 0
    return size

def get_file_list(argv):
    """
    Return file list from argv
    """

    flist = []
    if len(argv) > 1:
        for item in argv[1:]:
            if os.path.isfile(item):
                flist.append(item)
            elif os.path.isdir(item):
                this_list = files_in_dir(item)
                if this_list:
                    flist += this_list
    return flist

def size_format(size):
    """ 
    format size in bytes into human readable 
    """
    kbyte = 1024
    mbyte = 1024*1024
    if not size:
        size_int = 0
    else:
        size_int = int(size)

    if size_int < kbyte:
        size_str = f'{size_int:8d} B'
    elif size_int < mbyte:
        size_int /= kbyte
        size_str = f'{size_int:8.2f} KB'
    else:
        size_int /= mbyte
        size_str = f'{size_int:8.2f} MB'

    return size_str

def main():
    """
    Time to read file
    Usage:  ./fread-test.py [file1 [file2] ...]
    """
    verb = False
    #pdb.set_trace()
    file_list = get_file_list(sys.argv)

    if not file_list:
        print('Missing file(s)')
        return

    #
    # get size outsize of timing loop
    #
    file_info = {}
    for file in file_list:
        size = file_size(file)
        file_info[file] = size

    total = {}
    for (file,size) in file_info.items():
        start_time = time.time()

        fread = File_Read(file)
        fread.read()

        end_time = time.time()
        delta = end_time - start_time
        compress = fread.compress
        if verb:
            print(f'{file:>20s} : {compress:5s} : {delta:8.3f} secs')

        if total and compress in total:
            total[compress]['time'] += delta
            total[compress]['size'] += size
        else:
            total[compress] = {}
            total[compress]['time'] = delta
            total[compress]['size'] = size

    print('Summary')
    for (comp, info) in total.items():
        dtime = info['time']
        size = size_format(info['size'])

        print(f'{comp:5s} :  {dtime:8.3f} secs {size}')


if __name__ == '__main__':
    main()