#!/usr/bin/python

import os
import sys
from retrace import *

WORKER_RUNNING_PARSER = re.compile("^[ \t]*([0-9]+)[ \t]+[0-9]+[ \t]+([^ ^\t]+)[ \t]+.*retrace-server-worker ([0-9]+)$")

def run_ps():
    child = Popen(["ps", "-eo", "pid,ppid,etime,cmd"], stdout=PIPE)
    lines = child.communicate()[0].split("\n")

    return lines

def get_running_tasks(ps_output=None):
    if not ps_output:
        ps_output = run_ps()

    result = []

    for line in ps_output:
        match = WORKER_RUNNING_PARSER.match(line)
        if match:
            result.append((int(match.group(1)), int(match.group(3)), match.group(2)))

    return result

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:
            # 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", True)

        # 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(int(filename))
            except:
                continue

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