import asyncio
import json
import logging
from asyncio import Lock
from collections.abc import AsyncGenerator, Awaitable, Callable
from functools import partial
from http import HTTPStatus
from io import StringIO
from logging import Logger, LoggerAdapter
from typing import (
    Any,
    Literal,
    cast,
    get_args,
)

from anyio import (
    CancelScope,
    create_task_group,
    get_cancelled_exc_class,
    move_on_after,
    sleep,
)
from graphql import DocumentNode, GraphQLSchema, MiddlewareManager
from starlette.requests import Request
from starlette.responses import Response
from starlette.types import Receive, Scope, Send

from ..asgi.handlers import GraphQLHTTPHandler
from ..exceptions import HttpError
from ..format_error import format_error
from ..graphql import (
    ExecutionResult,
    GraphQLError,
    parse_query,
    subscribe,
    validate_data,
)
from ..logger import log_error
from ..subscription_handlers.events import SubscriptionEventType
from ..subscription_handlers.handlers import SubscriptionHandler
from ..types import (
    ErrorFormatter,
    ExtensionList,
    Extensions,
    MiddlewareList,
    Middlewares,
    QueryParser,
    QueryValidator,
    RootValue,
    ValidationRules,
)

EVENT_TYPES = Literal["next", "complete"]


class GraphQLServerSentEvent:
    """GraphQLServerSentEvent is a class that represents a single Server-Sent Event
    as defined in the GraphQL SSE Protocol specification
    (https://github.com/enisdenjo/graphql-sse/blob/master/PROTOCOL.md)
    """

    DEFAULT_SEPARATOR = "\r\n"

    def __init__(
        self,
        event: EVENT_TYPES,
        result: ExecutionResult | None = None,
    ):
        """Initializes the Server-Sent Event
        # Required arguments
        `event`: the type of the event. Either "next" or "complete"

        # Optional arguments
        `result`: an `ExecutionResult` or a `dict` that represents
        the result of the operation
        """
        assert event in get_args(EVENT_TYPES), f"Invalid event type: {event}"
        self.event = event
        self.result = result
        self.logger = logging.Logger("GraphQLServerSentEvent")

    def __str__(self) -> str:
        """Returns the string representation of the Server-Sent Event"""
        buffer = StringIO()
        buffer = self._write_to_buffer(buffer, "event", self.event)
        buffer = self._write_to_buffer(
            buffer,
            "data",
            (
                self.encode_execution_result()
                if self.event == "next" and self.result
                else ""
            ),
        )
        buffer.write(self.DEFAULT_SEPARATOR)

        return buffer.getvalue()

    def _write_to_buffer(
        self, buffer: StringIO, name: str, value: str | None
    ) -> StringIO:
        """Writes an SSE field to the buffered SSE event representation

        Returns the `StringIO` buffer with the field written to it

        # Required arguments
        `buffer`: the `StringIO` buffer to write to
        `name`: the name of the field
        `value`: the value of the field
        """
        if value is not None:
            buffer.write(f"{name}: {value}{self.DEFAULT_SEPARATOR}")
        return buffer

    def encode_execution_result(self) -> str:
        """Encodes the execution result into a single line JSON string

        Returns the JSON string representation of the execution result
        """
        payload: dict[str, Any] = {}
        if self.result is not None and self.result.data is not None:
            payload["data"] = self.result.data
        if self.result is not None and self.result.errors is not None:
            errors = []
            for error in self.result.errors:
                errors.append(format_error(error))
            payload["errors"] = errors

        return json.dumps(payload)


class ServerSentEventResponse(Response):
    """Sends GraphQL SSE events using the EventSource protocol using
    Starlette's Response class based on the implementation https://github.com/sysid/sse-starlette/
    """

    # Sends a ping event to the client every 15 seconds to overcome proxy timeout issues
    DEFAULT_PING_INTERVAL = 15

    def __init__(
        self,
        *args,
        generator: AsyncGenerator[GraphQLServerSentEvent, Any],
        send_timeout: int | None = None,
        ping_interval: int | None = None,
        headers: dict[str, str] | None = None,
        **kwargs,
    ):
        """Initializes an SSE Response that sends events generated by an async generator

        # Required arguments
        `generator`: an async generator that yields `GraphQLServerSentEvent` objects

        # Optional arguments
        `send_timeout`: the timeout in seconds to send an event to the client
        `ping_interval`: the interval in seconds to send a ping event to the client,
        overrides the DEFAULT_PING_INTERVAL of 15 seconds
        `headers`: a dictionary of headers to be sent with the response
        `encoding`: the encoding to use for the response
        """
        super().__init__(*args, **kwargs)
        self.generator = generator
        self.status_code = HTTPStatus.OK
        self.send_timeout = send_timeout
        self.ping_interval = ping_interval or self.DEFAULT_PING_INTERVAL
        self.body = None

        _headers: dict[str, str] = {}
        if headers is not None:
            _headers.update(headers)
        # mandatory for servers-sent events headers
        # allow cache control header to be set by user to support fan out proxies
        # https://www.fastly.com/blog/server-sent-events-fastly
        _headers.setdefault("Cache-Control", "no-cache")
        _headers.setdefault("Connection", "keep-alive")
        _headers.setdefault("X-Accel-Buffering", "no")
        _headers.setdefault("Transfer-Encoding", "chunked")
        self.media_type = "text/event-stream"
        self.init_headers(_headers)

        self._send_lock = Lock()

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        """The main entrypoint for the ServerSentEventResponse
        which is called by starlette

        # Required arguments
        `scope`: the starlette Scope object

        `receive`: the starlette Receive object

        `send`: the starlette Send object

        """
        async with create_task_group() as task_group:

            async def wrap_cancelling(func: Callable[[], Awaitable[None]]) -> None:
                await func()
                task_group.cancel_scope.cancel()

            task_group.start_soon(wrap_cancelling, partial(self._ping, send))
            task_group.start_soon(wrap_cancelling, partial(self.send_events, send))
            # this will cancel the task group when the client disconnects
            await wrap_cancelling(partial(self.listen_for_disconnect, receive))

    async def _ping(self, send: Send) -> None:
        """Sends a ping event to the client every `ping_interval` seconds gets
        cancelled if the client disconnects through the anyio
        CancelScope of the TaskGroup

        # Required arguments
        `send`: the starlette Send object
        """
        while True:
            await sleep(self.ping_interval)
            async with self._send_lock:
                await send(
                    {
                        "type": "http.response.body",
                        # always encode as utf-8 as per
                        # https://html.spec.whatwg.org/multipage/server-sent-events.html#sse-processing-model
                        "body": b":\r\n\r\n",
                        "more_body": True,
                    }
                )

    async def send_events(self, send: Send) -> None:
        """Sends the events generated by the async generator to the client

        # Required arguments
        `send`: the starlette Send object

        """
        async with self._send_lock:
            await send(
                {
                    "type": "http.response.start",
                    "status": self.status_code,
                    "headers": self.raw_headers,
                }
            )

        try:
            async for event in self.generator:
                async with self._send_lock:
                    with move_on_after(self.send_timeout) as timeout:
                        await send(
                            {
                                "type": "http.response.body",
                                "body": self.encode_event(event),
                                "more_body": True,
                            }
                        )

                    if timeout.cancel_called:
                        raise asyncio.TimeoutError()

        except (get_cancelled_exc_class(),) as e:
            logging.warning(e)
        finally:
            with CancelScope(shield=True):
                async with self._send_lock:
                    await send(
                        {"type": "http.response.body", "body": b"", "more_body": False}
                    )

    @staticmethod
    async def listen_for_disconnect(receive: Receive) -> None:
        """Listens for the client disconnect event and stops the streaming by exiting
        the infinite loop. This triggers the anyio CancelScope to cancel the TaskGroup

        # Required arguments
        `receive`: the starlette Receive object
        """
        while True:
            message = await receive()
            if message["type"] == "http.disconnect":
                logging.debug("Got event: http.disconnect. Stop streaming...")
                break

    @staticmethod
    def encode_event(event: GraphQLServerSentEvent) -> bytes:
        """Encodes the GraphQLServerSentEvent into a bytes object

        # Required arguments
        `event`: the GraphQLServerSentEvent object
        """
        # always encode as utf-8 as per
        # https://html.spec.whatwg.org/multipage/server-sent-events.html#sse-processing-model
        return str(event).encode("utf-8")


class SSESubscriptionHandler(SubscriptionHandler):
    """Subscription handler that delivers events via Server-Sent Events.

    Implements the `SubscriptionHandler` interface for use with
    `GraphQLHTTPHandler`'s `subscription_handlers` parameter.

    Detects SSE requests by checking for the `Accept: text/event-stream` header
    on POST requests. Uses the GraphQL SSE Protocol specification
    (https://github.com/enisdenjo/graphql-sse/blob/master/PROTOCOL.md) in
    distinct connections mode.

    # Example usage

    ```python
    from ariadne.asgi import GraphQL
    from ariadne.asgi.handlers import GraphQLHTTPHandler
    from ariadne.contrib.sse import SSESubscriptionHandler

    app = GraphQL(
        schema,
        http_handler=GraphQLHTTPHandler(
            subscription_handlers=[SSESubscriptionHandler()],
        ),
    )
    ```

    # Combining with other handlers

    ```python
    app = GraphQL(
        schema,
        http_handler=GraphQLHTTPHandler(
            subscription_handlers=[
                CallbackSubscriptionHandler(),  # Gateway access
                SSESubscriptionHandler(),  # Direct browser access
            ],
        ),
    )
    ```
    """

    def __init__(
        self,
        send_timeout: int | None = None,
        ping_interval: int | None = None,
        default_response_headers: dict[str, str] | None = None,
    ) -> None:
        """Initialize the SSE subscription handler.

        # Optional arguments

        `send_timeout`: the timeout in seconds to send an event to the client.

        `ping_interval`: the interval in seconds to send a ping event to the
        client. Defaults to `ServerSentEventResponse.DEFAULT_PING_INTERVAL`
        (15 seconds).

        `default_response_headers`: a dictionary of additional headers to be
        sent with the SSE response.
        """
        self.send_timeout = send_timeout
        self.ping_interval = ping_interval
        self.default_response_headers = default_response_headers

    def supports(self, request: Request, data: dict) -> bool:
        """Determine if this handler supports the given request.

        Returns `True` if the request is a POST with
        `Accept: text/event-stream` header.

        # Required arguments

        `request`: the Starlette/FastAPI request

        `data`: the parsed GraphQL request data
        """
        if request.method != "POST":
            return False
        accept = request.headers.get("Accept", "").split(",")
        accept = [a.strip() for a in accept]
        return "text/event-stream" in accept

    async def handle(
        self,
        request: Request,
        data: dict,
        *,
        schema: GraphQLSchema,
        context_value: Any,
        root_value: RootValue | None,
        query_parser: QueryParser | None,
        query_validator: QueryValidator | None,
        validation_rules: ValidationRules | None,
        debug: bool,
        introspection: bool,
        logger: None | str | Logger | LoggerAdapter,
        error_formatter: ErrorFormatter,
        middleware: MiddlewareList = None,
        middleware_manager_class: type[MiddlewareManager] | None = None,
        extensions: ExtensionList | None = None,
    ) -> Response:
        """Handle the subscription request via Server-Sent Events.

        Returns a `ServerSentEventResponse` that streams SSE events to the
        client with periodic ping/keepalive messages.
        """
        return ServerSentEventResponse(
            generator=self._sse_event_generator(
                data,
                schema=schema,
                context_value=context_value,
                root_value=root_value,
                query_parser=query_parser,
                query_validator=query_validator,
                validation_rules=validation_rules,
                debug=debug,
                introspection=introspection,
                logger=logger,
                error_formatter=error_formatter,
                middleware=middleware,
                middleware_manager_class=middleware_manager_class,
                extensions=extensions,
            ),
            ping_interval=self.ping_interval,
            send_timeout=self.send_timeout,
            headers=self.default_response_headers,
        )

    async def _sse_event_generator(
        self,
        data: dict,
        **kwargs: Any,
    ) -> AsyncGenerator[GraphQLServerSentEvent, None]:
        """Convert subscription events to SSE events.

        Calls `generate_events()` and maps each `SubscriptionEvent` to
        a `GraphQLServerSentEvent` for the SSE protocol.
        """
        async for event in self.generate_events(data, query_document=None, **kwargs):
            if event.event_type == SubscriptionEventType.NEXT:
                yield GraphQLServerSentEvent(event="next", result=event.result)
            elif event.event_type == SubscriptionEventType.ERROR:
                # Per GraphQL SSE Protocol: errors are sent as "next" events
                yield GraphQLServerSentEvent(event="next", result=event.result)
            elif event.event_type == SubscriptionEventType.COMPLETE:
                yield GraphQLServerSentEvent(event="complete")
            # KEEP_ALIVE is handled by ServerSentEventResponse._ping()


class GraphQLHTTPSSEHandler(GraphQLHTTPHandler):
    """Extension to the default GraphQLHTTPHandler to also handle Server-Sent Events
    as per the GraphQL SSE Protocol specification. This handler only supports the
    defined `Distinct connections mode` due to its statelessness. This implementation
    is based on the specification as of commit 80cf75b5952d1a065c95bdbd6a74304c90dbe2c5.
    For more information see the specification
    (https://github.com/enisdenjo/graphql-sse/blob/master/PROTOCOL.md)
    """

    def __init__(
        self,
        extensions: Extensions | None = None,
        middleware: Middlewares | None = None,
        middleware_manager_class: type[MiddlewareManager] | None = None,
        send_timeout: int | None = None,
        ping_interval: int | None = None,
        default_response_headers: dict[str, str] | None = None,
    ):
        super().__init__(extensions, middleware, middleware_manager_class)
        self.send_timeout = send_timeout
        self.ping_interval = ping_interval
        self.default_response_headers = default_response_headers

    async def handle_request_override(self, request: Request) -> Response | None:
        """Overrides the handle_request_override method to handle Server-Sent Events

        # Required arguments
        `request`: the starlette or FastAPI `Request` instance

        """

        if request.method == "POST":
            accept = request.headers.get("Accept", "").split(",")
            accept = [a.strip() for a in accept]
            if "text/event-stream" in accept:
                return await self.handle_sse_request(request)
        return None

    async def handle_sse_request(self, request: Request) -> Response:
        """Handles the HTTP request with GraphQL Subscription query using Server-Sent
        Events.

        # Required arguments

        `request`: the starlette `Request` instance
        """

        try:
            data = await self.extract_data_from_request(request)
            query = await self.get_query_from_sse_request(request, data)

            if self.schema is None:
                raise TypeError(
                    "schema is not set, call configure method to initialize it"
                )

            validate_data(data)
            context_value = await self.get_context_for_request(request, data)
            return ServerSentEventResponse(
                generator=self.sse_subscribe_to_graphql(query, data, context_value),
                ping_interval=self.ping_interval,
                send_timeout=self.send_timeout,
                headers=self.default_response_headers,
            )
        except (HttpError, TypeError, GraphQLError) as error:
            log_error(error, self.logger)
            if not isinstance(error, GraphQLError):
                error_message = (
                    (error.message or error.status)
                    if isinstance(error, HttpError)
                    else str(error)
                )
                error = GraphQLError(error_message, original_error=error)
            return ServerSentEventResponse(
                generator=self.sse_generate_error_response([error]),
                ping_interval=self.ping_interval,
                send_timeout=self.send_timeout,
                headers=self.default_response_headers,
            )

    async def get_query_from_sse_request(
        self, request: Request, data: Any
    ) -> DocumentNode:
        """Extracts GraphQL query from SSE request.

        Returns a `DocumentNode` with parsed query.

        # Required arguments

        `request`: the starlette `Request` instance

        `data`: an additional data parameter to potentially extract the query from
        """

        context_value = await self.get_context_for_request(request, data)
        return parse_query(context_value, self.query_parser, data)

    async def sse_subscribe_to_graphql(
        self, query_document: DocumentNode, data: Any, context_value: Any
    ):
        """Main SSE subscription generator for the GraphQL query.
        Yields `GraphQLServerSentEvent` instances and is to be consumed by a
        `ServerSentEventResponse` instance

        # Required arguments

        `query_document`: an already parsed GraphQL query.

        `data`: a `dict` with query data (`query` string, optionally `operationName`
        string and `variables` dictionary).

        `context_value`: a context value to make accessible as 'context' attribute
        of second argument (`info`) passed to resolvers and source functions.
        """

        success, results = await subscribe(
            self.schema,  # type: ignore
            data,
            context_value=context_value,
            root_value=self.root_value,
            query_document=query_document,
            query_validator=self.query_validator,
            validation_rules=self.validation_rules,
            debug=self.debug,
            introspection=self.introspection,
            logger=self.logger,
            error_formatter=self.error_formatter,
        )

        if not success:
            if not isinstance(results, list):
                error_payload = cast(list[dict[str, Any]], [results])
            else:
                error_payload = cast(list[dict[str, Any]], results)

            # This needs to be handled better, subscribe returns preformatted errors
            yield GraphQLServerSentEvent(
                event="next",
                result=ExecutionResult(
                    errors=[
                        GraphQLError(message=cast(str, error.get("message", "")))
                        for error in error_payload
                    ]
                ),
            )
        else:
            try:
                async for result in cast(AsyncGenerator, results):
                    yield GraphQLServerSentEvent(event="next", result=result)
            except (Exception, GraphQLError) as error:
                if not isinstance(error, GraphQLError):
                    error = GraphQLError(str(error), original_error=error)
                    log_error(error, self.logger)
                yield GraphQLServerSentEvent(
                    event="next", result=ExecutionResult(errors=[error])
                )

        yield GraphQLServerSentEvent(event="complete")

    @staticmethod
    async def sse_generate_error_response(
        errors: list[GraphQLError],
    ) -> AsyncGenerator[GraphQLServerSentEvent, Any]:
        """A Server-Sent Event response generator for the errors
        To be passed to a ServerSentEventResponse instance

        # Required arguments

        `errors`: a list of `GraphQLError` instances
        """

        yield GraphQLServerSentEvent(
            event="next", result=ExecutionResult(errors=errors)
        )
        yield GraphQLServerSentEvent(event="complete")
