import os
import sys
import subprocess
import threading
from numba import cuda
from numba.cuda.testing import (unittest, CUDATestCase, skip_on_cudasim,
                                skip_under_cuda_memcheck)
from numba.tests.support import captured_stdout


class TestCudaDetect(CUDATestCase):
    def test_cuda_detect(self):
        # exercise the code path
        with captured_stdout() as out:
            cuda.detect()
        output = out.getvalue()
        self.assertIn('Found', output)
        self.assertIn('CUDA devices', output)


@skip_under_cuda_memcheck('Hangs cuda-memcheck')
class TestCUDAFindLibs(CUDATestCase):

    def run_cmd(self, cmdline, env):
        popen = subprocess.Popen(cmdline,
                                 stdout=subprocess.PIPE,
                                 stderr=subprocess.PIPE,
                                 env=env)

        # finish in 5 minutes or kill it
        timeout = threading.Timer(5 * 60., popen.kill)
        try:
            timeout.start()
            out, err = popen.communicate()
            # the process should exit with an error
            return out.decode(), err.decode()
        finally:
            timeout.cancel()
        return None, None

    def run_test_in_separate_process(self, envvar, envvar_value):
        env_copy = os.environ.copy()
        env_copy[envvar] = str(envvar_value)
        code = """if 1:
            from numba import cuda
            @cuda.jit('(int64,)')
            def kernel(x):
                pass
            kernel(1,)
            """
        cmdline = [sys.executable, "-c", code]
        return self.run_cmd(cmdline, env_copy)

    @skip_on_cudasim('Simulator does not hit device library search code path')
    @unittest.skipIf(not sys.platform.startswith('linux'), "linux only")
    def test_cuda_find_lib_errors(self):
        """
        This tests that the find_libs works as expected in the case of an
        environment variable being used to set the path.
        """
        # one of these is likely to exist on linux, it's also unlikely that
        # someone has extracted the contents of libdevice into here!
        locs = ['lib', 'lib64']

        looking_for = None
        for l in locs:
            looking_for = os.path.join(os.path.sep, l)
            if os.path.exists(looking_for):
                break

        # This is the testing part, the test will only run if there's a valid
        # path in which to look
        if looking_for is not None:
            out, err = self.run_test_in_separate_process("NUMBA_CUDA_DRIVER",
                                                         looking_for)
            self.assertTrue(out is not None)
            self.assertTrue(err is not None)


if __name__ == '__main__':
    unittest.main()