#!/opt/local/bin/python2.7
import bz2
import cPickle
import imghdr
import math
import multiprocessing
import os
import Queue
import random
import sys

import pHash

from image import Image

NUM_WORKERS = 7

# Build index
# TODO: write more schtuff here

class IndexWorker(multiprocessing.Process):

    def __init__(self, file_list, queue):
        multiprocessing.Process.__init__(self)
        self.file_list = file_list
        self.queue = queue

    def run(self):

        list_size = len(self.file_list)
        ctr = 0
        for file_path in self.file_list:
            ctr += 1
            percent_done = int((float(ctr) / list_size) * 100)
            # We can only index jpegs right now
            if not imghdr.what(file_path) == 'jpeg':
                continue
            print '{:d} -- {:3d}% -- {:s}'.format(self.pid, percent_done, 
                                                  file_path)
            image_obj = Image(file_path)

            self.queue.put(image_obj, True)



def main():

    # Make sure at least one command line argument was provided
    if len(sys.argv) < 2:
        print 'Usage: build_index_faster.py <path>'
        return -1;

    # Print out the path we're walking
    path = sys.argv[1]
    print 'Building image index from path: {:s}\n'.format(path)

    master_file_list = []
    image_index = []

    # Walk the root path and create a list of all files in the root path
    for root, dirs, files in os.walk(path):
        for file in files:
            master_file_list.append(os.path.join(root, file))

    random.shuffle(master_file_list)

    print 'Done counting, {:d} files total.'.format(len(master_file_list))

    workers = []
    queue = multiprocessing.Queue(64)

    def chunker(seq, size):
        return (seq[pos:pos + size] for pos in xrange(0, len(seq), size))

    chunk_size = int(math.ceil(len(master_file_list) / NUM_WORKERS))

    # Create worker processes
    for work_chunk in chunker(master_file_list, chunk_size):
        workers.append(IndexWorker(work_chunk, queue))

    # Start workers
    for worker in workers:
        worker.start()

    print 'Waiting for workers to complete'

    # Loop until all workers are finished
    while any([w.is_alive() for w in workers]):
        try:
            image_obj = queue.get(True, 2)
            image_index.append(image_obj)
        except Queue.Empty:
            print 'Queue was empty'
            continue

    print 'Workers are complete: {:d} images found'.format(len(image_index))

    # Write index as a bzip2'd file to save space
    index_file = bz2.BZ2File('images_index.bz2', 'w')
    cPickle.dump(image_index, index_file)
    index_file.close()

    print '\nIndex file images_index.bz2 created. Done.'

if __name__ == '__main__':
    sys.exit(main())