#!/usr/bin/python
import os
import sys
from retrace import *

def get_process_tree(pid, ps_output):
    result = [pid]

    parser = re.compile("^([0-9]+)[ \t]+(%d).*$" % pid)

    for line in ps_output:
        match = parser.match(line)
        if match:
            pid = int(match.group(1))
            result.extend(get_process_tree(pid, ps_output))

    return result

def kill_process_and_childs(process_id, ps_output=None):
    result = True

    if not ps_output:
        ps_output = run_ps()

    for pid in get_process_tree(process_id, ps_output):
        try:
            os.kill(pid, 9)
        except OSError, ex:
            result = False

    return result

if __name__ == "__main__":
    logfile = os.path.join(CONFIG["LogDir"], "cleanup.log")

    with open(logfile, "a") as log:
        log.write(time.strftime("[%Y-%m-%d %H:%M:%S] Running cleanup\n"))

        # kill tasks running > 1 hour
        ps_output = run_ps()
        running_tasks = get_running_tasks(ps_output)
        for pid, taskid, runtime in running_tasks:
            # do not kill tasks started from task manager
            if CONFIG["AllowTaskManager"]:
                task = RetraceTask(taskid)

                if task.get_managed():
                    continue

            # ToDo: 5 = mm:ss, >5 = hh:mm:ss
            if len(runtime) > 5:
                log.write("Killing task %d running for %s\n" % (taskid, runtime))
                kill_process_and_childs(pid, ps_output)

        # kill orphaned tasks
        running_tasks = get_running_tasks()
        running_ids = []
        for pid, taskid, runtime in running_tasks:
            running_ids.append(taskid)

        for taskid in get_active_tasks():
            if not taskid in running_ids:
                log.write("Cleaning up orphaned task %d\n" % taskid)
                try:
                    task = RetraceTask(taskid)
                except:
                    log.write("Unable to create RetraceTask object for task %d\n" % taskid)
                    continue

                task.clean()
                task.set_log("Killed by garbage collector\n", True)

        if CONFIG["ArchiveTaskAfter"] > 0:
            # archive old tasks
            try:
                files = os.listdir(CONFIG["SaveDir"])
            except OSError, ex:
                files = []
                log.write("Error listing task directory: %s\n" % ex)

            for filename in files:
                try:
                    task = RetraceTask(filename)
                except:
                    continue

                if task.get_age() >= CONFIG["ArchiveTaskAfter"]:
                    log.write("Archiving task %s\n" % filename)
                    if not os.path.isdir(CONFIG["DropDir"]):
                        os.makedirs(CONFIG["DropDir"])

                    targetfile = os.path.join(CONFIG["DropDir"],
                                              "%s-%s.tar.gz" % (filename, time.strftime("%Y%m%d%H%M%S")))
                    with open(os.devnull, "w") as null:
                        child = Popen(["tar", "czf", targetfile, task.get_savedir()],
                                      stdout=PIPE, stderr=STDOUT)
                        stdout = child.communicate()[0]
                        if child.wait():
                            log += "Error: tar exitted with %d: %s\n" % (child.returncode, stdout)
                            try:
                                os.unlink(targetfile)
                            except:
                                pass

                            continue

                        task.remove()

        if CONFIG["DeleteTaskAfter"] > 0:
            # clean up old tasks
            try:
                files = os.listdir(CONFIG["SaveDir"])
            except OSError, ex:
                files = []
                log.write("Error listing task directory: %s\n" % ex)

            for filename in files:
                try:
                    task = RetraceTask(filename)
                except:
                    continue

                if task.get_age() >= CONFIG["DeleteTaskAfter"]:
                    log.write("Deleting old task %s\n" % filename)
                    task.remove()
