from typing import Any
from typing import Callable
from typing import List
from typing import Optional
from typing import Union

import pytest
import tornado
from graphql import parse
from opencensus.trace import execution_context
from opencensus.trace import tracer as tracer_module
from opencensus.trace.base_exporter import Exporter
from opencensus.trace.propagation.google_cloud_format import GoogleCloudFormatPropagator
from opencensus.trace.samplers import AlwaysOnSampler

from graphene_tornado.apollo_tooling.operation_id import (
    default_engine_reporting_signature,
)
from graphene_tornado.ext.apollo_engine_reporting.tests.schema import schema
from graphene_tornado.ext.apollo_engine_reporting.tests.test_engine_extension import (
    QUERY,
)
from graphene_tornado.ext.opencensus.opencensus_tracing_extension import (
    OpenCensusExtension,
)
from graphene_tornado.graphql_extension import GraphQLExtension
from graphene_tornado.tests.http_helper import HttpHelper
from graphene_tornado.tests.test_graphql import GRAPHQL_HEADER
from graphene_tornado.tests.test_graphql import response_json
from graphene_tornado.tests.test_graphql import url_string
from graphene_tornado.tornado_graphql_handler import TornadoGraphQLHandler


class GQLHandler(TornadoGraphQLHandler):
    def initialize(
        self,
        schema=None,
        middleware: Optional[Any] = None,
        root_value: Any = None,
        graphiql: bool = False,
        pretty: bool = False,
        batch: bool = False,
        extensions: List[
            Union[Callable[[], GraphQLExtension], GraphQLExtension]
        ] = None,
        exporter=None,
    ):
        super().initialize(
            schema, middleware, root_value, graphiql, pretty, batch, extensions
        )
        execution_context.set_opencensus_tracer(
            tracer_module.Tracer(
                sampler=AlwaysOnSampler(),
                exporter=exporter,
                propagator=GoogleCloudFormatPropagator(),
            )
        )

    def on_finish(self) -> None:
        tracer = execution_context.get_opencensus_tracer()
        tracer.finish()


class ExampleOpenCensusApplication(tornado.web.Application):
    def __init__(self, exporter):
        extension = lambda: OpenCensusExtension()
        handlers = [
            (
                r"/graphql",
                GQLHandler,
                dict(
                    graphiql=True,
                    schema=schema,
                    extensions=[extension],
                    exporter=exporter,
                ),
            ),
        ]
        tornado.web.Application.__init__(self, handlers)


@pytest.fixture
def app(exporter):
    return ExampleOpenCensusApplication(exporter)


@pytest.fixture
def app(exporter):
    return ExampleOpenCensusApplication(exporter)


@pytest.fixture
def http_helper(http_client, base_url):
    return HttpHelper(http_client, base_url)


@pytest.fixture
def exporter():
    return CapturingExporter()


@pytest.mark.gen_test()
def test_traces_match_query(http_helper, exporter):
    response = yield http_helper.get(url_string(query=QUERY), headers=GRAPHQL_HEADER)
    assert response.code == 200
    assert "data" in response_json(response)

    parent = exporter.spans.pop()[0]

    assert parent.name == "gql[b5c7307ba564]"
    assert parent.parent_span_id is None
    assert parent.attributes.get(
        "signature", None
    ) == default_engine_reporting_signature(parse(QUERY), "")

    spans = [span for span_list in exporter.spans for span in span_list]

    expected = [
        "gql_parsing",
        "gql_validation",
        "author",
        "aBoolean",
        "author.name",
        "author.posts",
        "author.posts.0.id",
        "author.posts.1.id",
    ]

    for span, exp in zip(spans, expected):
        assert span.name == exp
        assert span.parent_span_id == parent.span_id


class CapturingExporter(Exporter):
    def __init__(self):
        super(CapturingExporter, self).__init__()
        self._spans = []

    @property
    def spans(self):
        return self._spans

    def emit(self, span_datas):
        pass

    def export(self, span_datas):
        self._spans.append(span_datas)
