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

CONFIG = config.Config()

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

def check_config():
    if CONFIG["DeleteTaskAfter"] > 0 and CONFIG["ArchiveTaskAfter"] > 0:
        winner = "archiving"
        if CONFIG["DeleteTaskAfter"] < CONFIG["ArchiveTaskAfter"]:
            winner = "deleting"

        sys.stderr.write("WARNING: Both DeleteTaskAfter and ArchiveTaskAfter "
                         "are enabled, however they are mutually exclusive "
                         "settings. With your current settings only %s will "
                         "take place. Please check retrace-server.conf and "
                         "disable either DeleteTaskAfter or ArchiveTaskAfter "
                         "by setting its value to 0.\n" % winner)

    if (CONFIG["DeleteFailedTaskAfter"] > 0 and
        CONFIG["DeleteTaskAfter"] > 0 and
        CONFIG["DeleteTaskAfter"] < CONFIG["DeleteFailedTaskAfter"]):
        sys.stderr.write("ERROR: DeleteTaskAfter is set to a lower value than "
                         "DeleteFailedTaskAfter. Check retrace-server.conf, "
                         "did you accidentally reverse the values for "
                         "DeleteTaskAfter and DeleteFailedTaskAfter? "
                         "Not touching any tasks.\n")
        exit(1)

if __name__ == "__main__":
    check_config()

    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.create_worker().clean_task()
                task.set_log("Task was killed due to running too long or taking too many resources.\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.create_worker().remove_task()

        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.create_worker().remove_task()

        if CONFIG["DeleteFailedTaskAfter"] > 0:
            # clean up old failed 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["DeleteFailedTaskAfter"] and task.get_status() == STATUS_FAIL:
                    log.write("Deleting old failed task %s\n" % filename)
                    task.create_worker().remove_task()
