from __future__ import annotations

import typing
from enum import Enum

import strawberry
from opentelemetry import trace
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import (
    ConsoleSpanExporter,
    BatchSpanProcessor
)
from starlette.requests import Request
from starlette.websockets import WebSocket
from strawberry.asgi import GraphQL
from strawberry.extensions.tracing import OpenTelemetryExtension
from strawberry.schema import BaseSchema
from strawberry.types import Info

trace.set_tracer_provider(TracerProvider())
trace.get_tracer_provider().add_span_processor(
        BatchSpanProcessor(ConsoleSpanExporter())
)


class API(GraphQL):
    def __init__(self, config, schema: BaseSchema, *args, **kwargs):
        super().__init__(schema, *args, **kwargs)
        self.config = config

    async def get_context(self, request: typing.Union[Request, WebSocket], *args, **kwargs) -> typing.Any:
        return {
            'config': self.config,
        }


@strawberry.type
class Greeting:
    message: str


@strawberry.enum
class Mood(Enum):
    Happy = "Happy"
    Sad = "Sad"


@strawberry.input
class GreetInput:
    name: str
    mood: Mood


@strawberry.federation.type(extend=True)
class Query:
    @strawberry.field
    async def greet(self, info: Info, input: GreetInput) -> Greeting:
        config = info.context['config']

        return Greeting(
                message="Hi {}. Glad to hear you're glad!".format(input.name)
        ) if input.mood == Mood.Happy else Greeting(
                message="Hi {}, I hope you feel better !".format(input.name))

    @strawberry.field
    async def secret_greet(self, info: Info, input: GreetInput) -> Greeting:
        config = info.context['config']

        return Greeting(message="{}, {}".format(input.name, config['secret_message']))


def get_schema():
    return strawberry.federation.Schema(query=Query, extensions=[OpenTelemetryExtension])