from __future__ import annotations

import asyncio
from collections.abc import AsyncGenerator
from contextlib import suppress
from typing import (
    TYPE_CHECKING,
    Any,
    Generic,
    cast,
)

from strawberry.exceptions import ConnectionRejectionError
from strawberry.http.exceptions import NonTextMessageReceived, WebSocketDisconnected
from strawberry.http.typevars import Context, RootValue
from strawberry.schema.exceptions import CannotGetOperationTypeError
from strawberry.subscriptions.protocols.graphql_ws.types import (
    CompleteMessage,
    ConnectionInitMessage,
    ConnectionTerminateMessage,
    DataMessage,
    ErrorMessage,
    OperationMessage,
    StartMessage,
    StopMessage,
)
from strawberry.types.execution import ExecutionResult, PreExecutionError
from strawberry.types.unset import UnsetType

if TYPE_CHECKING:
    from collections.abc import AsyncGenerator

    from strawberry.http.async_base_view import AsyncBaseHTTPView, AsyncWebSocketAdapter
    from strawberry.schema import BaseSchema


class BaseGraphQLWSHandler(Generic[Context, RootValue]):
    def __init__(
        self,
        view: AsyncBaseHTTPView[Any, Any, Any, Any, Any, Context, RootValue],
        websocket: AsyncWebSocketAdapter,
        context: Context,
        root_value: RootValue | None,
        schema: BaseSchema,
        keep_alive: bool,
        keep_alive_interval: float | None,
        max_subscriptions_per_connection: int | None = None,
    ) -> None:
        self.view = view
        self.websocket = websocket
        self.context = context
        self.root_value = root_value
        self.schema = schema
        self.keep_alive = keep_alive
        self.keep_alive_interval = keep_alive_interval
        self.max_subscriptions_per_connection = max_subscriptions_per_connection
        self.keep_alive_task: asyncio.Task | None = None
        self.subscriptions: dict[str, AsyncGenerator] = {}
        self.tasks: dict[str, asyncio.Task] = {}
        self.connection_acknowledged: bool = False

    async def handle(self) -> None:
        try:
            try:
                async for message in self.websocket.iter_json(
                    ignore_parsing_errors=True
                ):
                    await self.handle_message(cast("OperationMessage", message))
            except NonTextMessageReceived:
                await self.websocket.close(
                    code=1002, reason="WebSocket message type must be text"
                )
        except WebSocketDisconnected:
            pass
        finally:
            if self.keep_alive_task:
                self.keep_alive_task.cancel()
                with suppress(BaseException):
                    await self.keep_alive_task

            await self.cleanup()

    async def handle_message(
        self,
        message: OperationMessage,
    ) -> None:
        if message["type"] == "connection_init":
            await self.handle_connection_init(message)
        elif message["type"] == "connection_terminate":
            await self.handle_connection_terminate(message)
        elif message["type"] == "start":
            await self.handle_start(message)
        elif message["type"] == "stop":
            await self.handle_stop(message)

    async def handle_connection_init(self, message: ConnectionInitMessage) -> None:
        payload = message.get("payload")
        if payload is not None and not isinstance(payload, dict):
            await self.send_message({"type": "connection_error"})
            await self.websocket.close(code=1000, reason="")
            return

        if isinstance(self.context, dict):
            self.context["connection_params"] = payload
        elif hasattr(self.context, "connection_params"):
            self.context.connection_params = payload

        try:
            connection_ack_payload = await self.view.on_ws_connect(self.context)
        except ConnectionRejectionError as e:
            await self.send_message({"type": "connection_error", "payload": e.payload})
            await self.websocket.close(code=1011, reason="")
            return

        if (
            isinstance(connection_ack_payload, UnsetType)
            or connection_ack_payload is None
        ):
            await self.send_message({"type": "connection_ack"})
        else:
            await self.send_message(
                {"type": "connection_ack", "payload": connection_ack_payload}
            )

        self.connection_acknowledged = True

        if self.keep_alive:
            keep_alive_handler = self.handle_keep_alive()
            self.keep_alive_task = asyncio.create_task(keep_alive_handler)

    async def handle_connection_terminate(
        self, message: ConnectionTerminateMessage
    ) -> None:
        await self.websocket.close(code=1000, reason="")

    async def handle_start(self, message: StartMessage) -> None:
        if not self.connection_acknowledged:
            await self.websocket.close(code=4401, reason="Unauthorized")
            return

        operation_id = message["id"]

        # Clean up existing operation with same ID to prevent task leaks
        if operation_id in self.tasks:
            await self.cleanup_operation(operation_id)

        if (
            self.max_subscriptions_per_connection is not None
            and len(self.tasks) >= self.max_subscriptions_per_connection
        ):
            await self.send_message(
                ErrorMessage(
                    type="error",
                    id=operation_id,
                    payload={"message": "Subscription limit reached"},
                )
            )
            return

        payload = message["payload"]
        query = payload["query"]
        operation_name = payload.get("operationName")
        variables = payload.get("variables")

        result_handler = self.handle_async_results(
            operation_id, query, operation_name, variables
        )
        self.tasks[operation_id] = asyncio.create_task(result_handler)

    async def handle_stop(self, message: StopMessage) -> None:
        operation_id = message["id"]
        await self.cleanup_operation(operation_id)

    async def handle_keep_alive(self) -> None:
        assert self.keep_alive_interval
        while True:
            await self.send_message({"type": "ka"})
            await asyncio.sleep(self.keep_alive_interval)

    async def handle_async_results(
        self,
        operation_id: str,
        query: str,
        operation_name: str | None,
        variables: dict[str, object] | None,
    ) -> None:
        try:
            result_source = await self.schema.subscribe(
                query=query,
                variable_values=variables,
                operation_name=operation_name,
                context_value=self.context,
                root_value=self.root_value,
            )
            self.subscriptions[operation_id] = result_source

            is_first_result = True

            async for result in result_source:
                if is_first_result and isinstance(result, PreExecutionError):
                    assert result.errors

                    await self.send_message(
                        ErrorMessage(
                            type="error",
                            id=operation_id,
                            payload=result.errors[0].formatted,
                        )
                    )
                    return

                await self.send_data_message(result, operation_id)
                is_first_result = False

            await self.send_message(CompleteMessage(type="complete", id=operation_id))

        except CannotGetOperationTypeError as e:
            await self.send_message(
                ErrorMessage(
                    type="error",
                    id=operation_id,
                    payload={"message": e.as_http_error_reason()},
                )
            )
        except asyncio.CancelledError:
            await self.send_message(CompleteMessage(type="complete", id=operation_id))

    async def cleanup_operation(self, operation_id: str) -> None:
        if operation_id in self.subscriptions:
            with suppress(RuntimeError):
                await self.subscriptions[operation_id].aclose()
            del self.subscriptions[operation_id]

        self.tasks[operation_id].cancel()
        with suppress(BaseException):
            await self.tasks[operation_id]
        del self.tasks[operation_id]

    async def cleanup(self) -> None:
        for operation_id in list(self.tasks.keys()):
            await self.cleanup_operation(operation_id)

    async def send_data_message(
        self, execution_result: ExecutionResult, operation_id: str
    ) -> None:
        data_message: DataMessage = {
            "type": "data",
            "id": operation_id,
            "payload": {"data": execution_result.data},
        }

        if execution_result.errors:
            data_message["payload"]["errors"] = [
                err.formatted for err in execution_result.errors
            ]

        if execution_result.extensions:
            data_message["payload"]["extensions"] = execution_result.extensions

        await self.send_message(data_message)

    async def send_message(self, message: OperationMessage) -> None:
        await self.websocket.send_json(message)


__all__ = ["BaseGraphQLWSHandler"]
