# Copyright (c) 2016-present, Facebook, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
##############################################################################

## @package timeout_guard
# Module caffe2.python.timeout_guard
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import contextlib
import threading
import os
import time
import signal
import logging
from future.utils import viewitems


'''
Sometimes CUDA devices can get stuck, 'deadlock'. In this case it is often
better just the kill the process automatically. Use this guard to set a
maximum timespan for a python call, such as RunNet(). If it does not complete
in time, process is killed.

Example usage:
    with timeout_guard.CompleteInTimeOrDie(10.0):
        core.RunNet(...)
'''


class WatcherThread(threading.Thread):

    def __init__(self, timeout_secs):
        threading.Thread.__init__(self)
        self.timeout_secs = timeout_secs
        self.completed = False
        self.condition = threading.Condition()
        self.daemon = True
        self.caller_thread = threading.current_thread()

    def run(self):
        started = time.time()
        self.condition.acquire()
        while time.time() - started < self.timeout_secs and not self.completed:
            self.condition.wait(self.timeout_secs - (time.time() - started))
        self.condition.release()
        if not self.completed:
            log = logging.getLogger("timeout_guard")
            log.error("Call did not finish in time. Timeout:{}s PID: {}".format(
                self.timeout_secs,
                os.getpid(),
            ))

            # First try dying cleanly, but in 10 secs, exit properly
            def forcequit():
                time.sleep(10.0)
                log.info("Prepared output, dumping threads. ")
                print("Caller thread was: {}".format(self.caller_thread))
                print("-----After force------")
                import sys
                import traceback
                code = []
                for threadId, stack in viewitems(sys._current_frames()):
                    if threadId == self.caller_thread.ident:
                        code.append("\n# ThreadID: %s" % threadId)
                        for filename, lineno, name, line in traceback.extract_stack(stack):
                            code.append('File: "%s", line %d, in %s' % (filename, lineno, name))
                            if line:
                                code.append("  %s" % (line.strip()))

                print("\n".join(code))
                log.error("Process did not terminate cleanly in 10 s, forcing")
                os._exit(1)

            forcet = threading.Thread(target=forcequit, args=())
            forcet.daemon = True
            forcet.start()
            print("Caller thread was: {}".format(self.caller_thread))
            print("-----Before forcing------")
            import sys
            import traceback
            code = []
            for threadId, stack in viewitems(sys._current_frames()):
                code.append("\n# ThreadID: %s" % threadId)
                for filename, lineno, name, line in traceback.extract_stack(stack):
                    code.append('File: "%s", line %d, in %s' % (filename, lineno, name))
                    if line:
                        code.append("  %s" % (line.strip()))

            print("\n".join(code))
            os.kill(os.getpid(), signal.SIGINT)


@contextlib.contextmanager
def CompleteInTimeOrDie(timeout_secs):
    watcher = WatcherThread(timeout_secs)
    watcher.start()
    yield
    watcher.completed = True
    watcher.condition.acquire()
    watcher.condition.notify()
    watcher.condition.release()


def EuthanizeIfNecessary(timeout_secs=120):
    '''
    Call this if you have problem with process getting stuck at shutdown.
    It will kill the process if it does not terminate in timeout_secs.
    '''
    watcher = WatcherThread(timeout_secs)
    watcher.start()