# Copyright 2014 MongoDB, 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.

"""Utilities for testing Motor with asyncio."""

import asyncio
import functools
import gc
import inspect
import os
import unittest
from unittest import SkipTest

try:
    from asyncio import ensure_future
except ImportError:
    from asyncio import async as ensure_future

from mockupdb import MockupDB

from motor import motor_asyncio
from test.assert_logs_backport import AssertLogsMixin
from test.test_environment import env, CA_PEM, CLIENT_PEM


class _TestMethodWrapper(object):
    """Wraps a test method to raise an error if it returns a value.

    This is mainly used to detect undecorated generators (if a test
    method yields it must use a decorator to consume the generator),
    but will also detect other kinds of return values (these are not
    necessarily errors, but we alert anyway since there is no good
    reason to return a value from a test).

    Adapted from Tornado's test framework.
    """
    def __init__(self, orig_method):
        self.orig_method = orig_method

    def __call__(self):
        result = self.orig_method()
        if inspect.isgenerator(result):
            raise TypeError("Generator test methods should be decorated with "
                            "@asyncio_test")
        elif result is not None:
            raise ValueError("Return value from test method ignored: %r" %
                             result)

    def __getattr__(self, name):
        """Proxy all unknown attributes to the original method.

        This is important for some of the decorators in the `unittest`
        module, such as `unittest.skipIf`.
        """
        return getattr(self.orig_method, name)


class AsyncIOTestCase(AssertLogsMixin, unittest.TestCase):
    longMessage = True  # Used by unittest.TestCase
    ssl = False  # If True, connect with SSL, skip if mongod isn't SSL

    def __init__(self, methodName='runTest'):
        super().__init__(methodName)

        # It's easy to forget the @asyncio_test decorator, but if you do
        # the test will silently be ignored because nothing will consume
        # the generator. Replace the test method with a wrapper that will
        # make sure it's not an undecorated generator.
        # (Adapted from Tornado's AsyncTestCase.)
        setattr(self, methodName, _TestMethodWrapper(
            getattr(self, methodName)))

    def setUp(self):
        super(AsyncIOTestCase, self).setUp()

        # Ensure that the event loop is passed explicitly in Motor.
        asyncio.set_event_loop(None)
        self.loop = asyncio.new_event_loop()
        asyncio.set_event_loop(self.loop)

        if self.ssl and not env.mongod_started_with_ssl:
            raise SkipTest("mongod doesn't support SSL, or is down")

        self.cx = self.asyncio_client()
        self.db = self.cx.motor_test
        self.collection = self.db.test_collection
        self.loop.run_until_complete(self.collection.drop())

    def get_client_kwargs(self, **kwargs):
        if env.mongod_started_with_ssl:
            kwargs.setdefault('ssl_ca_certs', CA_PEM)
            kwargs.setdefault('ssl_certfile', CLIENT_PEM)

        kwargs.setdefault('ssl', env.mongod_started_with_ssl)
        kwargs.setdefault('io_loop', self.loop)

        return kwargs

    def asyncio_client(self, uri=None, *args, **kwargs):
        """Get an AsyncIOMotorClient.

        Ignores self.ssl, you must pass 'ssl' argument.
        """
        return motor_asyncio.AsyncIOMotorClient(
            uri or env.uri,
            *args,
            **self.get_client_kwargs(**kwargs))

    def asyncio_rsc(self, uri=None, *args, **kwargs):
        """Get an open MotorClient for replica set.

        Ignores self.ssl, you must pass 'ssl' argument.
        """
        return motor_asyncio.AsyncIOMotorClient(
            uri or env.rs_uri,
            *args,
            **self.get_client_kwargs(**kwargs))

    @asyncio.coroutine
    def make_test_data(self):
        yield from self.collection.delete_many({})
        yield from self.collection.insert_many([{'_id': i} for i in range(200)])

    make_test_data.__test__ = False

    def tearDown(self):
        self.cx.close()
        self.loop.stop()
        self.loop.run_forever()
        self.loop.close()
        gc.collect()


class AsyncIOMockServerTestCase(AsyncIOTestCase):
    def server(self, *args, **kwargs):
        server = MockupDB(*args, **kwargs)
        server.run()
        self.addCleanup(server.stop)
        return server

    def client_server(self, *args, **kwargs):
        server = self.server(*args, **kwargs)
        client = motor_asyncio.AsyncIOMotorClient(server.uri, io_loop=self.loop)
        self.addCleanup(client.close)

        return client, server

    def run_thread(self, fn, *args, **kwargs):
        return self.loop.run_in_executor(None,
                                         functools.partial(fn, *args, **kwargs))

    def ensure_future(self, coro):
        return ensure_future(coro, loop=self.loop)

    def fetch_next(self, cursor):
        @asyncio.coroutine
        def fetch_next():
            return (yield from cursor.fetch_next)

        return self.ensure_future(fetch_next())


def get_async_test_timeout(default=5):
    """Get the global timeout setting for async tests.

    Returns a float, the timeout in seconds.
    """
    try:
        timeout = float(os.environ.get('ASYNC_TEST_TIMEOUT'))
        return max(timeout, default)
    except (ValueError, TypeError):
        return default


# TODO: Spin off to a PyPI package.
def asyncio_test(func=None, timeout=None):
    """Decorator for coroutine methods of AsyncIOTestCase::

        class MyTestCase(AsyncIOTestCase):
            @asyncio_test
            def test(self):
                # Your test code here....
                pass

    Default timeout is 5 seconds. Override like::

        class MyTestCase(AsyncIOTestCase):
            @asyncio_test(timeout=10)
            def test(self):
                # Your test code here....
                pass

    You can also set the ASYNC_TEST_TIMEOUT environment variable to a number
    of seconds. The final timeout is the ASYNC_TEST_TIMEOUT or the timeout
    in the test (5 seconds or the passed-in timeout), whichever is longest.
    """
    def wrap(f):
        @functools.wraps(f)
        def wrapped(self, *args, **kwargs):
            if timeout is None:
                actual_timeout = get_async_test_timeout()
            else:
                actual_timeout = get_async_test_timeout(timeout)

            coro_exc = None

            def exc_handler(loop, context):
                nonlocal coro_exc
                coro_exc = context['exception']

                # Raise CancelledError from run_until_complete below.
                task.cancel()

            self.loop.set_exception_handler(exc_handler)
            coro = asyncio.coroutine(f)(self, *args, **kwargs)
            coro = asyncio.wait_for(coro, actual_timeout, loop=self.loop)
            task = ensure_future(coro, loop=self.loop)
            try:
                self.loop.run_until_complete(task)
            except:
                if coro_exc:
                    # Raise the error thrown in on_timeout, with only the
                    # traceback from the coroutine itself, not from
                    # run_until_complete.
                    raise coro_exc from None

                raise

        return wrapped

    if func is not None:
        # Used like:
        #     @gen_test
        #     def f(self):
        #         pass
        if not inspect.isfunction(func):
            msg = ("%r is not a test method. Pass a timeout as"
                   " a keyword argument, like @asyncio_test(timeout=7)")
            raise TypeError(msg % func)
        return wrap(func)
    else:
        # Used like @gen_test(timeout=10)
        return wrap


@asyncio.coroutine
def get_command_line(client):
    command_line = yield from client.admin.command('getCmdLineOpts')
    assert command_line['ok'] == 1, "getCmdLineOpts() failed"
    return command_line


@asyncio.coroutine
def server_is_mongos(client):
    ismaster_response = yield from client.admin.command('ismaster')
    return ismaster_response.get('msg') == 'isdbgrid'


@asyncio.coroutine
def skip_if_mongos(client):
    is_mongos = yield from server_is_mongos(client)
    if is_mongos:
        raise unittest.SkipTest("connected to mongos")


@asyncio.coroutine
def remove_all_users(db):
    yield from db.command({"dropAllUsersFromDatabase": 1})
