# coding=utf-8
from __future__ import unicode_literals, print_function, absolute_import


import os
import sys
import threading
import traceback
import signal
try:
    from Queue import Queue
except ImportError:
    from queue import Queue


class WatchDog(object):

    def __init__(self, timeout):
        self._timeout = float(timeout)
        self._q = Queue()
        self._th = threading.Thread(target=self._worker)
        self._th.daemon = True
        self._th.start()

    def _worker(self):
        try:
            while True:
                job = self._q.get()
                job()
        except StopIteration:
            pass

    def _dispatch(self, job):
        self._q.put(job)

    def _terminate(self):
        try:
            chunks = []
            for thread_id, frame in sys._current_frames().items():
                chunks.append('Thread {}\n'.format(thread_id))
                chunks.extend(traceback.format_stack(frame))
                chunks.append('\n')
            data = ''.join(chunks)
            sys.stdout.write(data)
            sys.stdout.flush()
            sys.stderr.write(data)
            sys.stderr.flush()
        finally:
            os.kill(os.getpid(), signal.SIGKILL)

    def guard(self):
        ev = threading.Event()

        @self._dispatch
        def wait_timeout():
            is_set = ev.wait(self._timeout)
            if not is_set:
                self._terminate()

        return WatchDogGuard(self, ev)


class WatchDogGuard(object):

    def __init__(self, watch_dog, event):
        assert isinstance(watch_dog, WatchDog)
        self._watch_dog = watch_dog
        self._ev = event

    def __enter__(self):
        pass

    def __exit__(self, exc_type, exc_val, exc_tb):
        self._ev.set()


def test():
    w = WatchDog(5)

    for d in range(20):
        import time
        print(d)
        with w.guard():
            time.sleep(d)


if __name__ == '__main__':
    test()

