import argparse
import os
import progressbar  # progressbar2
import signal
import time

from git import Repo, GitCommandError  # gitpython
from multiprocessing import Lock, Process, Queue
from Queue import Empty

# Paths
closest_save_path = "closest.dat"
progress_save_path = ".progress"

# Multiprocessing
stopping = False


class Diff:
    def __init__(self, sha, count):
        self.sha = sha
        self.count = int(count)


def save_closest(path, dictionary):
    f = open(path, "w")
    for key in dictionary:
        obj = dictionary[key]
        f.write(key + "\t" + str(obj.sha) + "\t" + str(obj.count) + "\n")
    f.close()


def load_closest(path):
    dictionary = {}
    try:
        f = open(path, "r")
        lines = f.read().splitlines()
        for line in lines:
            if line != "":
                try:
                    tokens = line.split("\t")
                    dictionary[tokens[0]] = Diff(tokens[1], tokens[2])
                except IndexError:
                    pass
                except ValueError:
                    pass
        f.close()
    except IOError:
        pass
    return dictionary


def do_check_path(queue, git, git_lock, path, sha):
    global stopping

    git_lock.acquire()
    try:
        diff = git.diff("--numstat", "-w", "--ignore-blank-lines", path)
        diff_details = diff.split("\t")
        if len(diff_details) == 1 and diff_details[0] == "":
            queue.put(["closest_rm", path])
            queue.put(["stop", path])
            return
        changes = git.diff("-U0", "-w", "--ignore-blank-lines", path).split("\n")[4:]
        changes_closest = git.diff("-U0", "-w", "--ignore-blank-lines", sha, "--", path).split("\n")[4:]
    finally:
        git_lock.release()

    cherrypick = True

    for i, change in enumerate(changes):
        if stopping:
            break
        if change[0:2] == "@@":
            continue
        for j, change_closest in enumerate(changes_closest):
            if stopping:
                break
            if change_closest[0:2] == "@@":
                continue
            if change == change_closest:
                cherrypick = False

    if cherrypick:
        git_lock.acquire()
        try:
            git.checkout(path)
            queue.put(["closest_rm", path])
        finally:
            git_lock.release()

    queue.put(["stop", path])


def check_path(queue, git, git_lock, path, sha):
    try:
        do_check_path(queue, git, git_lock, path, sha)
    except GitCommandError:
        queue.put(["stop", path])


def signal_handler(*_):
    global stopping
    stopping = True


def main(args):
    global stopping, progress_save_path, closest_save_path

    signal.signal(signal.SIGINT, signal_handler)

    queues = list()

    repo = Repo(args.kernel_path)
    git = repo.git
    git_lock = Lock()

    last_branch = ""
    last_paths = list()

    closest_revisions = load_closest(closest_save_path)

    try:
        f = open(progress_save_path, "r")
        lines = f.read().splitlines()
        try:
            last_paths = lines[1:]
        except IndexError:
            pass
        f.close()
    except IOError:
        pass

    last_save = 0

    paths = closest_revisions.keys()
    n = len(paths)
    offset = 0
    for path in last_paths:
        try:
            del paths[paths.index(path)]
            offset += 1
        except ValueError:
            pass
    if n > 0 and offset < n:
        widgets = [
            progressbar.Percentage(),
            " (",
            progressbar.Counter(),
            " of " + str(n) + ") ",
            progressbar.Bar(),
            " ",
            progressbar.Timer(),
            " " * 50
        ]
        bar = progressbar.ProgressBar(max_value=n, widgets=widgets, term_width=150)
        j = offset + 1
        for path in paths:
            bar.update(j)
            if stopping:
                break
            if len(queues) < args.max_threads:
                bar.update(j)
                queue = Queue()
                Process(target=check_path, args=(queue, git, git_lock, path, closest_revisions[path].sha)).start()
                queues.append(queue)
            while len(queues) == args.max_threads or (paths.index(path) == len(paths) - 1 and len(queues) > 0):
                bar.update(j)
                for queue in queues[:]:
                    bar.update(j)
                    try:
                        data = queue.get(False)
                        data_action = data[0]
                        data_path = data[1]
                        if data_action == "stop":
                            if not stopping:
                                label = " -- " + data_path
                                widgets[7] = label + " " * (50 - len(label))
                                j += 1
                                last_paths.append(data_path)
                            queues.remove(queue)
                        if data_action == "closest_rm":
                            if data_path in closest_revisions:
                                del closest_revisions[data_path]
                        now = time.time()
                        if now - last_save > 10:
                            f = open(progress_save_path, "w")
                            for last_path in last_paths:
                                f.write("\n" + last_path)
                            f.close()
                            save_closest(closest_save_path, closest_revisions)
                            last_save = now
                    except Empty:
                        pass
        if not stopping:
            bar.finish()
        if not stopping:
            last_paths = list()

    if stopping and last_branch != "":
        f = open(progress_save_path, "w")
        f.write(last_branch)
        for path in last_paths:
            f.write("\n" + path)
        f.close()

    if not stopping:
        os.remove(progress_save_path)

    save_closest(closest_save_path, closest_revisions)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("-j",
                        type=int,
                        dest="max_threads",
                        default=16,
                        help="Maximum number of threads (default: 16)")
    parser.add_argument("-k", "--kernel_path",
                        type=str,
                        dest="kernel_path",
                        required=True,
                        help="Location of the Kernel folder")

    main(parser.parse_args())