import asyncio

from aiohttp.test_utils import unittest_run_loop

from ddtrace.contrib.aiohttp.middlewares import trace_app, trace_middleware, CONFIG_KEY
from ddtrace.ext import http
from ddtrace.sampler import RateSampler
from ddtrace.constants import SAMPLING_PRIORITY_KEY, ANALYTICS_SAMPLE_RATE_KEY

from opentracing.scope_managers.asyncio import AsyncioScopeManager
from tests.opentracer.utils import init_tracer
from .utils import TraceTestCase
from .app.web import setup_app, noop_middleware
from ...utils import assert_span_http_status_code


class TestTraceMiddleware(TraceTestCase):
    """
    Ensures that the trace Middleware creates root spans at
    the beginning of a request.
    """
    def enable_tracing(self):
        trace_app(self.app, self.tracer)

    @unittest_run_loop
    @asyncio.coroutine
    def test_handler(self):
        # it should create a root span when there is a handler hit
        # with the proper tags
        request = yield from self.client.request('GET', '/')
        assert 200 == request.status
        text = yield from request.text()
        assert 'What\'s tracing?' == text
        # the trace is created
        traces = self.tracer.writer.pop_traces()
        assert 1 == len(traces)
        assert 1 == len(traces[0])
        span = traces[0][0]
        # with the right fields
        assert 'aiohttp.request' == span.name
        assert 'aiohttp-web' == span.service
        assert 'web' == span.span_type
        assert 'GET /' == span.resource
        assert str(self.client.make_url('/')) == span.get_tag(http.URL)
        assert 'GET' == span.get_tag('http.method')
        assert_span_http_status_code(span, 200)
        assert 0 == span.error

    @asyncio.coroutine
    def _test_param_handler(self, query_string=''):
        if query_string:
            fqs = '?' + query_string
        else:
            fqs = ''
        # it should manage properly handlers with params
        request = yield from self.client.request('GET', '/echo/team' + fqs)
        assert 200 == request.status
        text = yield from request.text()
        assert 'Hello team' == text
        # the trace is created
        traces = self.tracer.writer.pop_traces()
        assert 1 == len(traces)
        assert 1 == len(traces[0])
        span = traces[0][0]
        # with the right fields
        assert 'GET /echo/{name}' == span.resource
        assert str(self.client.make_url('/echo/team')) == span.get_tag(http.URL)
        assert_span_http_status_code(span, 200)
        if self.app[CONFIG_KEY].get('trace_query_string'):
            assert query_string == span.get_tag(http.QUERY_STRING)
        else:
            assert http.QUERY_STRING not in span.meta

    @unittest_run_loop
    def test_param_handler(self):
        return self._test_param_handler()

    @unittest_run_loop
    def test_query_string(self):
        return self._test_param_handler('foo=bar')

    @unittest_run_loop
    def test_query_string_duplicate_keys(self):
        return self._test_param_handler('foo=bar&foo=baz&x=y')

    @unittest_run_loop
    def test_param_handler_trace(self):
        self.app[CONFIG_KEY]['trace_query_string'] = True
        return self._test_param_handler()

    @unittest_run_loop
    def test_query_string_trace(self):
        self.app[CONFIG_KEY]['trace_query_string'] = True
        return self._test_param_handler('foo=bar')

    @unittest_run_loop
    def test_query_string_duplicate_keys_trace(self):
        self.app[CONFIG_KEY]['trace_query_string'] = True
        return self._test_param_handler('foo=bar&foo=baz&x=y')

    @unittest_run_loop
    @asyncio.coroutine
    def test_404_handler(self):
        # it should not pollute the resource space
        request = yield from self.client.request('GET', '/404/not_found')
        assert 404 == request.status
        # the trace is created
        traces = self.tracer.writer.pop_traces()
        assert 1 == len(traces)
        assert 1 == len(traces[0])
        span = traces[0][0]
        # with the right fields
        assert '404' == span.resource
        assert str(self.client.make_url('/404/not_found')) == span.get_tag(http.URL)
        assert 'GET' == span.get_tag('http.method')
        assert_span_http_status_code(span, 404)

    @unittest_run_loop
    @asyncio.coroutine
    def test_server_error(self):
        """
        When a server error occurs (uncaught exception)
            The span should be flagged as an error
        """
        request = yield from self.client.request('GET', '/uncaught_server_error')
        assert request.status == 500
        traces = self.tracer.writer.pop_traces()
        assert len(traces) == 1
        assert len(traces[0]) == 1
        span = traces[0][0]
        assert span.get_tag('http.method') == 'GET'
        assert_span_http_status_code(span, 500)
        assert span.error == 1

    @unittest_run_loop
    @asyncio.coroutine
    def test_500_response_code(self):
        """
        When a 5XX response code is returned
            The span should be flagged as an error
        """
        request = yield from self.client.request('GET', '/caught_server_error')
        assert request.status == 503
        traces = self.tracer.writer.pop_traces()
        assert len(traces) == 1
        assert len(traces[0]) == 1
        span = traces[0][0]
        assert span.get_tag('http.method') == 'GET'
        assert_span_http_status_code(span, 503)
        assert span.error == 1

    @unittest_run_loop
    @asyncio.coroutine
    def test_coroutine_chaining(self):
        # it should create a trace with multiple spans
        request = yield from self.client.request('GET', '/chaining/')
        assert 200 == request.status
        text = yield from request.text()
        assert 'OK' == text
        # the trace is created
        traces = self.tracer.writer.pop_traces()
        assert 1 == len(traces)
        assert 3 == len(traces[0])
        root = traces[0][0]
        handler = traces[0][1]
        coroutine = traces[0][2]
        # root span created in the middleware
        assert 'aiohttp.request' == root.name
        assert 'GET /chaining/' == root.resource
        assert str(self.client.make_url('/chaining/')) == root.get_tag(http.URL)
        assert 'GET' == root.get_tag('http.method')
        assert_span_http_status_code(root, 200)
        # span created in the coroutine_chaining handler
        assert 'aiohttp.coro_1' == handler.name
        assert root.span_id == handler.parent_id
        assert root.trace_id == handler.trace_id
        # span created in the coro_2 handler
        assert 'aiohttp.coro_2' == coroutine.name
        assert handler.span_id == coroutine.parent_id
        assert root.trace_id == coroutine.trace_id

    @unittest_run_loop
    @asyncio.coroutine
    def test_static_handler(self):
        # it should create a trace with multiple spans
        request = yield from self.client.request('GET', '/statics/empty.txt')
        assert 200 == request.status
        text = yield from request.text()
        assert 'Static file\n' == text
        # the trace is created
        traces = self.tracer.writer.pop_traces()
        assert 1 == len(traces)
        assert 1 == len(traces[0])
        span = traces[0][0]
        # root span created in the middleware
        assert 'aiohttp.request' == span.name
        assert 'GET /statics' == span.resource
        assert str(self.client.make_url('/statics/empty.txt')) == span.get_tag(http.URL)
        assert 'GET' == span.get_tag('http.method')
        assert_span_http_status_code(span, 200)

    @unittest_run_loop
    @asyncio.coroutine
    def test_middleware_applied_twice(self):
        # it should be idempotent
        app = setup_app(self.app.loop)
        # the middleware is not present
        assert 1 == len(app.middlewares)
        assert noop_middleware == app.middlewares[0]
        # the middleware is present (with the noop middleware)
        trace_app(app, self.tracer)
        assert 2 == len(app.middlewares)
        # applying the middleware twice doesn't add it again
        trace_app(app, self.tracer)
        assert 2 == len(app.middlewares)
        # and the middleware is always the first
        assert trace_middleware == app.middlewares[0]
        assert noop_middleware == app.middlewares[1]

    @unittest_run_loop
    @asyncio.coroutine
    def test_exception(self):
        request = yield from self.client.request('GET', '/exception')
        assert 500 == request.status
        yield from request.text()

        traces = self.tracer.writer.pop_traces()
        assert 1 == len(traces)
        spans = traces[0]
        assert 1 == len(spans)
        span = spans[0]
        assert 1 == span.error
        assert 'GET /exception' == span.resource
        assert 'error' == span.get_tag('error.msg')
        assert 'Exception: error' in span.get_tag('error.stack')

    @unittest_run_loop
    @asyncio.coroutine
    def test_async_exception(self):
        request = yield from self.client.request('GET', '/async_exception')
        assert 500 == request.status
        yield from request.text()

        traces = self.tracer.writer.pop_traces()
        assert 1 == len(traces)
        spans = traces[0]
        assert 1 == len(spans)
        span = spans[0]
        assert 1 == span.error
        assert 'GET /async_exception' == span.resource
        assert 'error' == span.get_tag('error.msg')
        assert 'Exception: error' in span.get_tag('error.stack')

    @unittest_run_loop
    @asyncio.coroutine
    def test_wrapped_coroutine(self):
        request = yield from self.client.request('GET', '/wrapped_coroutine')
        assert 200 == request.status
        text = yield from request.text()
        assert 'OK' == text

        traces = self.tracer.writer.pop_traces()
        assert 1 == len(traces)
        spans = traces[0]
        assert 2 == len(spans)
        span = spans[0]
        assert 'GET /wrapped_coroutine' == span.resource
        span = spans[1]
        assert 'nested' == span.name
        assert span.duration > 0.25, 'span.duration={0}'.format(span.duration)

    @unittest_run_loop
    @asyncio.coroutine
    def test_distributed_tracing(self):
        # distributed tracing is enabled by default
        tracing_headers = {
            'x-datadog-trace-id': '100',
            'x-datadog-parent-id': '42',
        }

        request = yield from self.client.request('GET', '/', headers=tracing_headers)
        assert 200 == request.status
        text = yield from request.text()
        assert "What's tracing?" == text
        # the trace is created
        traces = self.tracer.writer.pop_traces()
        assert 1 == len(traces)
        assert 1 == len(traces[0])
        span = traces[0][0]
        # with the right trace_id and parent_id
        assert span.trace_id == 100
        assert span.parent_id == 42
        assert span.get_metric(SAMPLING_PRIORITY_KEY) is None

    @unittest_run_loop
    @asyncio.coroutine
    def test_distributed_tracing_with_sampling_true(self):
        self.tracer.priority_sampler = RateSampler(0.1)

        tracing_headers = {
            'x-datadog-trace-id': '100',
            'x-datadog-parent-id': '42',
            'x-datadog-sampling-priority': '1',
        }

        request = yield from self.client.request('GET', '/', headers=tracing_headers)
        assert 200 == request.status
        text = yield from request.text()
        assert "What's tracing?" == text
        # the trace is created
        traces = self.tracer.writer.pop_traces()
        assert 1 == len(traces)
        assert 1 == len(traces[0])
        span = traces[0][0]
        # with the right trace_id and parent_id
        assert 100 == span.trace_id
        assert 42 == span.parent_id
        assert 1 == span.get_metric(SAMPLING_PRIORITY_KEY)

    @unittest_run_loop
    @asyncio.coroutine
    def test_distributed_tracing_with_sampling_false(self):
        self.tracer.priority_sampler = RateSampler(0.9)

        tracing_headers = {
            'x-datadog-trace-id': '100',
            'x-datadog-parent-id': '42',
            'x-datadog-sampling-priority': '0',
        }

        request = yield from self.client.request('GET', '/', headers=tracing_headers)
        assert 200 == request.status
        text = yield from request.text()
        assert "What's tracing?" == text
        # the trace is created
        traces = self.tracer.writer.pop_traces()
        assert 1 == len(traces)
        assert 1 == len(traces[0])
        span = traces[0][0]
        # with the right trace_id and parent_id
        assert 100 == span.trace_id
        assert 42 == span.parent_id
        assert 0 == span.get_metric(SAMPLING_PRIORITY_KEY)

    @unittest_run_loop
    @asyncio.coroutine
    def test_distributed_tracing_disabled(self):
        # pass headers for distributed tracing
        self.app['datadog_trace']['distributed_tracing_enabled'] = False
        tracing_headers = {
            'x-datadog-trace-id': '100',
            'x-datadog-parent-id': '42',
        }

        request = yield from self.client.request('GET', '/', headers=tracing_headers)
        assert 200 == request.status
        text = yield from request.text()
        assert "What's tracing?" == text
        # the trace is created
        traces = self.tracer.writer.pop_traces()
        assert 1 == len(traces)
        assert 1 == len(traces[0])
        span = traces[0][0]
        # distributed tracing must be ignored by default
        assert span.trace_id != 100
        assert span.parent_id != 42

    @unittest_run_loop
    @asyncio.coroutine
    def test_distributed_tracing_sub_span(self):
        self.tracer.priority_sampler = RateSampler(1.0)

        # activate distributed tracing
        tracing_headers = {
            'x-datadog-trace-id': '100',
            'x-datadog-parent-id': '42',
            'x-datadog-sampling-priority': '0',
        }

        request = yield from self.client.request('GET', '/sub_span', headers=tracing_headers)
        assert 200 == request.status
        text = yield from request.text()
        assert 'OK' == text
        # the trace is created
        traces = self.tracer.writer.pop_traces()
        assert 1 == len(traces)
        assert 2 == len(traces[0])
        span, sub_span = traces[0][0], traces[0][1]
        # with the right trace_id and parent_id
        assert 100 == span.trace_id
        assert 42 == span.parent_id
        assert 0 == span.get_metric(SAMPLING_PRIORITY_KEY)
        # check parenting is OK with custom sub-span created within server code
        assert 100 == sub_span.trace_id
        assert span.span_id == sub_span.parent_id
        assert sub_span.get_metric(SAMPLING_PRIORITY_KEY) is None

    def _assert_200_parenting(self, traces):
        """Helper to assert parenting when handling aiohttp requests.

        This is used to ensure that parenting is consistent between Datadog
        and OpenTracing implementations of tracing.
        """
        assert 2 == len(traces)
        assert 1 == len(traces[0])

        # the inner span will be the first trace since it completes before the
        # outer span does
        inner_span = traces[0][0]
        outer_span = traces[1][0]

        # confirm the parenting
        assert outer_span.parent_id is None
        assert inner_span.parent_id is None

        assert outer_span.name == 'aiohttp_op'

        # with the right fields
        assert 'aiohttp.request' == inner_span.name
        assert 'aiohttp-web' == inner_span.service
        assert 'web' == inner_span.span_type
        assert 'GET /' == inner_span.resource
        assert str(self.client.make_url('/')) == inner_span.get_tag(http.URL)
        assert 'GET' == inner_span.get_tag('http.method')
        assert_span_http_status_code(inner_span, 200)
        assert 0 == inner_span.error

    @unittest_run_loop
    @asyncio.coroutine
    def test_parenting_200_dd(self):
        with self.tracer.trace('aiohttp_op'):
            request = yield from self.client.request('GET', '/')
            assert 200 == request.status
            text = yield from request.text()

        assert "What's tracing?" == text
        traces = self.tracer.writer.pop_traces()
        self._assert_200_parenting(traces)

    @unittest_run_loop
    @asyncio.coroutine
    def test_parenting_200_ot(self):
        """OpenTracing version of test_handler."""
        ot_tracer = init_tracer('aiohttp_svc', self.tracer, scope_manager=AsyncioScopeManager())

        with ot_tracer.start_active_span('aiohttp_op'):
            request = yield from self.client.request('GET', '/')
            assert 200 == request.status
            text = yield from request.text()

        assert "What's tracing?" == text
        traces = self.tracer.writer.pop_traces()
        self._assert_200_parenting(traces)

    @unittest_run_loop
    @asyncio.coroutine
    def test_analytics_integration_enabled(self):
        """ Check trace has analytics sample rate set """
        self.app['datadog_trace']['analytics_enabled'] = True
        self.app['datadog_trace']['analytics_sample_rate'] = 0.5
        request = yield from self.client.request('GET', '/template/')
        yield from request.text()

        # Assert root span sets the appropriate metric
        self.assert_structure(
            dict(name='aiohttp.request', metrics={ANALYTICS_SAMPLE_RATE_KEY: 0.5})
        )

    @unittest_run_loop
    @asyncio.coroutine
    def test_analytics_integration_default(self):
        """ Check trace has analytics sample rate set """
        request = yield from self.client.request('GET', '/template/')
        yield from request.text()

        # Assert root span does not have the appropriate metric
        root = self.get_root_span()
        self.assertIsNone(root.get_metric(ANALYTICS_SAMPLE_RATE_KEY))

    @unittest_run_loop
    @asyncio.coroutine
    def test_analytics_integration_disabled(self):
        """ Check trace has analytics sample rate set """
        self.app['datadog_trace']['analytics_enabled'] = False
        request = yield from self.client.request('GET', '/template/')
        yield from request.text()

        # Assert root span does not have the appropriate metric
        root = self.get_root_span()
        self.assertIsNone(root.get_metric(ANALYTICS_SAMPLE_RATE_KEY))