from __future__ import annotations

import http
import ssl as ssl_module
import urllib.parse
from typing import Any, Awaitable, Callable, Literal

from werkzeug.exceptions import NotFound
from werkzeug.routing import Map, RequestRedirect

from ..http11 import Request, Response
from .server import Server, ServerConnection, serve


__all__ = ["route", "unix_route", "Router"]


class Router:
    """WebSocket router supporting :func:`route`."""

    def __init__(
        self,
        url_map: Map,
        server_name: str | None = None,
        url_scheme: str = "ws",
    ) -> None:
        self.url_map = url_map
        self.server_name = server_name
        self.url_scheme = url_scheme
        for rule in self.url_map.iter_rules():
            rule.websocket = True

    def get_server_name(self, connection: ServerConnection, request: Request) -> str:
        if self.server_name is None:
            return request.headers["Host"]
        else:
            return self.server_name

    def redirect(self, connection: ServerConnection, url: str) -> Response:
        response = connection.respond(http.HTTPStatus.FOUND, f"Found at {url}")
        response.headers["Location"] = url
        return response

    def not_found(self, connection: ServerConnection) -> Response:
        return connection.respond(http.HTTPStatus.NOT_FOUND, "Not Found")

    def route_request(
        self, connection: ServerConnection, request: Request
    ) -> Response | None:
        """Route incoming request."""
        url_map_adapter = self.url_map.bind(
            server_name=self.get_server_name(connection, request),
            url_scheme=self.url_scheme,
        )
        try:
            parsed = urllib.parse.urlparse(request.path)
            handler, kwargs = url_map_adapter.match(
                path_info=parsed.path,
                query_args=parsed.query,
            )
        except RequestRedirect as redirect:
            return self.redirect(connection, redirect.new_url)
        except NotFound:
            return self.not_found(connection)
        connection.handler, connection.handler_kwargs = handler, kwargs
        return None

    async def handler(self, connection: ServerConnection) -> None:
        """Handle a connection."""
        return await connection.handler(connection, **connection.handler_kwargs)


def route(
    url_map: Map,
    *args: Any,
    server_name: str | None = None,
    ssl: ssl_module.SSLContext | Literal[True] | None = None,
    create_router: type[Router] | None = None,
    **kwargs: Any,
) -> Awaitable[Server]:
    """
    Create a WebSocket server dispatching connections to different handlers.

    This feature requires the third-party library `werkzeug`_:

    .. code-block:: console

        $ pip install werkzeug

    .. _werkzeug: https://werkzeug.palletsprojects.com/

    :func:`route` accepts the same arguments as
    :func:`~websockets.sync.server.serve`, except as described below.

    The first argument is a :class:`werkzeug.routing.Map` that maps URL patterns
    to connection handlers. In addition to the connection, handlers receive
    parameters captured in the URL as keyword arguments.

    Here's an example::


        from websockets.asyncio.router import route
        from werkzeug.routing import Map, Rule

        async def channel_handler(websocket, channel_id):
            ...

        url_map = Map([
            Rule("/channel/<uuid:channel_id>", endpoint=channel_handler),
            ...
        ])

        # set this future to exit the server
        stop = asyncio.get_running_loop().create_future()

        async with route(url_map, ...) as server:
            await stop


    Refer to the documentation of :mod:`werkzeug.routing` for details.

    If you define redirects with ``Rule(..., redirect_to=...)`` in the URL map,
    when the server runs behind a reverse proxy that modifies the ``Host``
    header or terminates TLS, you need additional configuration:

    * Set ``server_name`` to the name of the server as seen by clients. When not
      provided, websockets uses the value of the ``Host`` header.

    * Set ``ssl=True`` to generate ``wss://`` URIs without actually enabling
      TLS. Under the hood, this bind the URL map with a ``url_scheme`` of
      ``wss://`` instead of ``ws://``.

    There is no need to specify ``websocket=True`` in each rule. It is added
    automatically.

    Args:
        url_map: Mapping of URL patterns to connection handlers.
        server_name: Name of the server as seen by clients. If :obj:`None`,
            websockets uses the value of the ``Host`` header.
        ssl: Configuration for enabling TLS on the connection. Set it to
            :obj:`True` if a reverse proxy terminates TLS connections.
        create_router: Factory for the :class:`Router` dispatching requests to
            handlers. Set it to a wrapper or a subclass to customize routing.

    """
    url_scheme = "ws" if ssl is None else "wss"
    if ssl is not True and ssl is not None:
        kwargs["ssl"] = ssl

    if create_router is None:
        create_router = Router

    router = create_router(url_map, server_name, url_scheme)

    _process_request: (
        Callable[
            [ServerConnection, Request],
            Awaitable[Response | None] | Response | None,
        ]
        | None
    ) = kwargs.pop("process_request", None)
    if _process_request is None:
        process_request: Callable[
            [ServerConnection, Request],
            Awaitable[Response | None] | Response | None,
        ] = router.route_request
    else:

        async def process_request(
            connection: ServerConnection, request: Request
        ) -> Response | None:
            response = _process_request(connection, request)
            if isinstance(response, Awaitable):
                response = await response
            if response is not None:
                return response
            return router.route_request(connection, request)

    return serve(router.handler, *args, process_request=process_request, **kwargs)


def unix_route(
    url_map: Map,
    path: str | None = None,
    **kwargs: Any,
) -> Awaitable[Server]:
    """
    Create a WebSocket Unix server dispatching connections to different handlers.

    :func:`unix_route` combines the behaviors of :func:`route` and
    :func:`~websockets.asyncio.server.unix_serve`.

    Args:
        url_map: Mapping of URL patterns to connection handlers.
        path: File system path to the Unix socket.

    """
    return route(url_map, unix=True, path=path, **kwargs)