#!/opt/local/bin/python3

#!/usr/bin/python3

import argparse
import glob
import math
from multiprocessing import Process
from multiprocessing import Lock
import multiprocessing
import os
import re
import subprocess
import sys

LEPTON_PATH = '/Users/dsorber/Documents/lepton/lepton'

# Bash text formatting codes
BOLD = '\033[1m'
ENDC = '\033[0m'
RED  = '\033[31m'
YELLOW = '\033[33m'
PURPLE = '\033[35m'

def file_walker(top_path, regex):
    
    files = []
    
    for dirpath, dirnames, filenames in os.walk(top_path):
        for filename in filenames:
            fullpath = os.path.join(dirpath, filename)
            if re.search(regex, fullpath, re.IGNORECASE):
                #~ print(fullpath)
                files.append(fullpath)
        for dirname in dirnames:
            file_walker(os.path.join(dirpath, dirname), regex)
    return files
    
    
class CompressWrapper(Process):

    def __init__(self, index, image_queue, remove_orig, print_lock):
        Process.__init__(self)
        self.index = index
        self.image_queue = image_queue
        self.remove_orig = remove_orig
        self.print_lock = print_lock
        self.errors = []
        self.running = False

    def run(self):
        self.running = True
        
        num_images = len(self.image_queue)
        ten_percent = math.ceil(num_images * 0.1)
        percent = 0
        count = 0
        
        # Open log file for this instance to use
        output = open('compressor_{:03d}.log'.format(self.index), 'w')
        
        # Iterate over each image and compress it using the lepton binary
        for idx, image in enumerate(self.image_queue):
            output.write('[{:2d}] Image {:d} of {:d}:\n'.format(self.index, idx + 1, num_images))
            output.write('\t{:s}\n'.format(image))
            
            if count == ten_percent:
                percent += 10
                count = 0
                self.print_lock.acquire()
                print('Worker {:d} -- {:d}% complete'.format(self.index, percent))
                self.print_lock.release()
            else:
                count += 1
            
            compressed_image = '{:s}.lep'.format(image)
            cmd = '{:s} "{:s}" "{:s}"'.format(LEPTON_PATH, image, compressed_image)
            #~ rc = subprocess.run(cmd, shell=True)
            #~ if rc.returncode == 0:
            rc = subprocess.call(cmd, shell=True, stdout=output, stderr=output)
            if rc == 0:
                output.write('SUCCESS\n')

                # Only remove original file if specified
                if self.remove_orig:
                    os.remove(image)
            else:
                # Append this image to the list of errors
                output.write('ERROR\n')
                self.errors.append(image)
            output.write('\n\n')
            
        output.close()
        self.print_lock.acquire()
        print('Worker {:d} -- [ DONE ]'.format(self.index))
        self.print_lock.release()
        
    def get_errors(self):
        return self.errors
        
    
def do_decompress(top_path, remove_orig):
    
    # Find all *.lep files (ignoring case)
    images = file_walker(sys.argv[1], '\.lep$')
    num_images = len(images)
    
    print('Found {:d} *.lep files'.format(num_images))
    
    # Iterate over each image and decompress it using the lepton binary
    for idx, image in enumerate(images):
        print('Image {:d} of {:d}:'.format(idx + 1, num_images))
        print('\t{:s}'.format(image))
        
        cmd = '{:s} "{:s}" "{:s}"'.format(LEPTON_PATH, image, image[:-4])
        #~ rc = subprocess.run(cmd, shell=True)
        #~ if rc.returncode == 0:
        rc = subprocess.call(cmd, shell=True)
        if rc == 0:
            print('{:s}SUCCESS{:s}'.format(BOLD, ENDC))
            
            # Only remove original file if specified
            if remove_orig:
                os.remove(image)
        else:
            print('{:s}{:s}ERROR{:s}'.format(BOLD, RED, ENDC))
        print('\n')

def chunks(l, n):
    return (l[i:i+n] for i in xrange(0, len(l), n))

def main():
       
    # Make sure the specified path to lepton is correct
    if not (os.path.isfile(LEPTON_PATH) and os.access(LEPTON_PATH, os.X_OK)):
        print('{:s}{:s}ERROR:{:s} unable to find lepton executable'
              .format(BOLD, RED, ENDC))
        return -1

    # Parse command line arguments
    parser = argparse.ArgumentParser(description='Lepton Wrapper')
    
    parser.add_argument('path',
                        help='path in which to look for images')
    
    parser.add_argument('-d', '--decompress', action='store_true', 
                        default=False, help='decompress *.lep files')
                        
    parser.add_argument('-r', '--remove', action='store_true', default=False,
                        help='remove orignal files after compress/decompress')
    
    args = parser.parse_args()

    
    # Either compress or decompress depending on mode parameter
    if args.decompress:
        do_decompress(args.path, args.remove==True)
    else:
        
        # Find all *.jpg (or *.jpeg) images (ignoring case)
        images = file_walker(sys.argv[1], '(\.jpg$)|(\.jpeg$)')
        num_images = len(images)
        print('Found {:d} *.jpg images\n'.format(num_images))
        
        num_workers = max(1, multiprocessing.cpu_count())
        print('Using {:d} workers'.format(num_workers))

        # Divide input data into chunks
        chunk_size = math.ceil(num_images / num_workers)
        chunks = [images[x:x+chunk_size] 
                  for x in range(0, num_images, chunk_size)]
        
        # Start compressor worker threads
        the_lock = Lock()
        workers = []
        for worker_idx in range(num_workers):
            worker = CompressWrapper(worker_idx, chunks[worker_idx], 
                                     args.remove==True, the_lock)
            worker.start()
            print('\tWorker {:d} - spawned'.format(worker_idx))
            workers.append(worker)
        print()
        
        # Join compressor worker threads and look for any errors
        all_errors = []
        for worker in workers:
            worker.join()
            all_errors.extend(worker.get_errors())
            
        # Check for errors
        if all_errors:
            print('Errors were encountered when attempting to compress the '
                  'following images:')
            for error in all_errors:
                print('\t{:s}'.format(error))
        else:           
            print('COMPLETE')
        
    
if __name__ == '__main__':
    sys.exit(main())