"""REST API Module with automatic OpenAPI generation."""
import asyncio
import os
from functools import wraps
from inspect import signature
from typing import (
    Any,
    Awaitable,
    Callable,
    Coroutine,
    List,
    Literal,
    Optional,
    TypeVar,
    cast,
)

from aiohttp.typedefs import Handler
from aiohttp.web import Application, HTTPFound, Request, Response, StreamResponse
from aiohttp.web_middlewares import middleware
from aiohttp.web_ws import WebSocketResponse
from aiohttp_sse import EventSourceResponse, sse_response
from typing_extensions import ParamSpec

from ..data.document import Document
from ..data.json import jsonable_encoder
from ..data.odm import FaunaModel
from ..helpers import do_response
from ..utils import setup_logging
from .docs import extract, html_string, load, transform
from .router import APIRouter

P = ParamSpec("P")
_T = TypeVar("_T")

Middleware = Callable[[Request, Handler], Awaitable[StreamResponse]]


class APIServer(Application):
    """
    AioFauna APIServer class for handling API operations.

    :param title: The title of the API.
    :type title: str
    :param servers: List of server URLs.
    :type servers: list[str]
    :param description: API description.
    :type description: str
    :param version: API version.
    :type version: str
    :param openapi_url: URL for OpenAPI documentation.
    :type openapi_url: str
    """

    def __init__(
        self,
        *args: Any,
        title: str = "AioFauna",
        servers: Optional[List[str]] = None,
        version: str = "0.0.1",
        openapi_url: str = "/openapi.json",
        **kwargs: Any,
    ):
        super().__init__(*args, logger=setup_logging(self.__class__.__name__), **kwargs)
        schemas = FaunaModel.Metadata.subclasses + Document.Metadata.subclasses
        self.openapi = {
            "openapi": "3.1.0",
            "info": {"title": title, "version": version},
            "paths": {},
            "servers": [{"url": url} for url in servers or []],
            "components": {
                "schemas": {schema.__name__: schema.schema() for schema in schemas}
            },
        }
        self._route_open_api_params = {}
        self.openapi_url: str = openapi_url

        @self.get("/openapi.json")
        async def openapi():
            response = jsonable_encoder(self.openapi)
            return response

        @self.get("/docs")
        async def docs():
            return Response(
                text=html_string(self.openapi_url), content_type="text/html"
            )

    def document(
        self, path: str, method: str
    ) -> Callable[
        [Callable[P, Coroutine[Any, Any, _T]]], Callable[P, Coroutine[Any, Any, _T]]
    ]:
        """

        Decorator to generate OpenAPI documentation for an endpoint handler.

        :param path: The URL path of the endpoint.
        :type path: str
        :param method: The HTTP method of the endpoint.
        :type method: str
        :return: Decorated function for handling the endpoint.
        :rtype: Callable

        """

        def decorator(
            func: Callable[P, Coroutine[Any, Any, _T]]
        ) -> Callable[P, Coroutine[Any, Any, _T]]:
            sig = signature(func)
            params = sig.parameters
            open_api_params = extract(params.copy(), path)
            self._route_open_api_params[(path, method)] = open_api_params
            transform(self.openapi, path, method, func, open_api_params)

            async def wrapper(*args: P.args, **kwargs: P.kwargs) -> StreamResponse:
                request = cast(Request, args[0])
                new_args = args[1:]
                args_to_apply = await load(request, params.copy())
                definitive_args = {}
                for name, param in params.items():
                    if name in args_to_apply:
                        definitive_args[name] = args_to_apply[name]
                    elif param.default is not param.empty:
                        definitive_args[name] = param.default
                    else:
                        raise ValueError(
                            f"Missing parameter {name} for {func.__name__}"
                        )
                response = await func(*new_args, **definitive_args, **kwargs)  # type: ignore
                return do_response(response)

            wrapper._handler = func
            wrapper.__name__ = func.__name__
            return cast(Callable[P, Coroutine[Any, Any, _T]], wrapper)

        return decorator

    def get(
        self, path: str, **kwargs: Any
    ) -> Callable[[Callable[P, Awaitable[_T]]], Callable[P, Awaitable[_T]]]:
        """
        Decorator for handling GET requests.

        :param path: The API path for the GET request.
        :type path: str
        :return: Decorated function for handling the GET request.
        :rtype: Callable
        """

        def decorator(
            func: Callable[P, Awaitable[StreamResponse]]
        ) -> Callable[P, Awaitable[StreamResponse]]:
            self.router.add_get(path, self.document(path, "GET")(func), **kwargs)
            return func

        return decorator

    def post(
        self, path: str, **kwargs: Any
    ) -> Callable[[Callable[P, Awaitable[_T]]], Callable[P, Awaitable[_T]]]:
        """
        Decorator for handling POST requests.

        :param path: The API path for the POST request.
        :type path: str
        :return: Decorated function for handling the POST request.
        :rtype: Callable
        """

        def decorator(
            func: Callable[P, Awaitable[StreamResponse]]
        ) -> Callable[P, Awaitable[StreamResponse]]:
            self.router.add_post(path, self.document(path, "POST")(func), **kwargs)
            return func

        return decorator

    def put(
        self, path: str, **kwargs: Any
    ) -> Callable[[Callable[P, Awaitable[_T]]], Callable[P, Awaitable[_T]]]:
        """
        Decorator for handling PUT requests.

        :param path: The API path for the PUT request.
        :type path: str
        :return: Decorated function for handling the PUT request.
        :rtype: Callable

        """

        def decorator(
            func: Callable[P, Awaitable[StreamResponse]]
        ) -> Callable[P, Awaitable[StreamResponse]]:
            self.router.add_put(path, self.document(path, "PUT")(func), **kwargs)
            return func

        return decorator

    def delete(
        self, path: str, **kwargs: Any
    ) -> Callable[[Callable[P, Awaitable[_T]]], Callable[P, Awaitable[_T]]]:
        """
        Decorator for handling DELETE requests.

        :param path: The API path for the DELETE request.
        :type path: str
        :return: Decorated function for handling the DELETE request.
        :rtype: Callable

        """

        def decorator(
            func: Callable[P, Awaitable[StreamResponse]]
        ) -> Callable[P, Awaitable[StreamResponse]]:
            self.router.add_delete(path, self.document(path, "DELETE")(func), **kwargs)
            return func

        return decorator

    def patch(
        self, path: str, **kwargs: Any
    ) -> Callable[[Callable[P, Awaitable[_T]]], Callable[P, Awaitable[_T]]]:
        """

        Decorator for handling PATCH requests.

        :param path: The API path for the PATCH request.
        :type path: str
        :return: Decorated function for handling the PATCH request.
        :rtype: Callable

        """

        def decorator(
            func: Callable[P, Awaitable[StreamResponse]]
        ) -> Callable[P, Awaitable[StreamResponse]]:
            self.router.add_patch(path, self.document(path, "PATCH")(func), **kwargs)
            return func

        return decorator

    def head(
        self, path: str, **kwargs: Any
    ) -> Callable[[Callable[P, Awaitable[_T]]], Callable[P, Awaitable[_T]]]:
        """

        Decorator for handling HEAD requests.

        :param path: The API path for the HEAD request.
        :type path: str
        :return: Decorated function for handling the HEAD request.
        :rtype: Callable

        """

        def decorator(
            func: Callable[P, Awaitable[StreamResponse]]
        ) -> Callable[P, Awaitable[StreamResponse]]:
            self.router.add_head(path, self.document(path, "HEAD")(func), **kwargs)
            return func

        return decorator

    def options(
        self, path: str, **kwargs: Any
    ) -> Callable[[Callable[P, Awaitable[_T]]], Callable[P, Awaitable[_T]]]:
        """

        Decorator for handling OPTIONS requests.

        :param path: The API path for the OPTIONS request.
        :type path: str
        :return: Decorated function for handling the OPTIONS request.
        :rtype: Callable

        """

        def decorator(
            func: Callable[P, Awaitable[StreamResponse]]
        ) -> Callable[P, Awaitable[StreamResponse]]:
            self.router.add_options(
                path, self.document(path, "OPTIONS")(func), **kwargs
            )
            return func

        return decorator

    def on_event(
        self, event: Literal["startup", "shutdown"]
    ) -> Callable[
        [Callable[[Application], Awaitable[None]]],
        Callable[[Application], Awaitable[None]],
    ]:
        """

        Decorator for handling startup and shutdown lifecycle events.

        :param event: The event to handle (startup or shutdown).

        """

        def decorator(
            func: Callable[[Application], Awaitable[None]]
        ) -> Callable[[Application], Awaitable[None]]:
            if event not in ("startup", "shutdown"):
                raise ValueError("Event must be startup or shutdown")
            elif event == "startup":
                self.on_startup.append(func)
            else:
                self.on_shutdown.append(func)
            return func

        return decorator

    def sse(
        self, path: str
    ) -> Callable[
        [Callable[P, Awaitable[EventSourceResponse]]],
        Callable[P, Awaitable[EventSourceResponse]],
    ]:
        """

        Server-Sent Events decorator.

        :param path: The API path for the SSE request.
        :type path: str
        :return: Decorated function for handling the SSE request.

        This decorator will handle SSE connections and pass the response object to the decorated function, be sure to include a parameter with the type annotation `EventSourceResponse` in your function signature.
        You can use the `send` method of the response object to send events to the client.

        Example:

        .. code-block:: python

        import asyncio
        from aiofauna import FaunaClient, APIServer, EventSourceResponse, to_json,q

        app = APIServer()
        fauna = FaunaClient()

        @app.sse("/stream")
        async def stream(sse: EventSourceResponse):
                        while True:
                                        now = await fauna.query(q.now())
                                        await sse.send(to_json(now["@ts"]))
                                        await asyncio.sleep(1)

        """

        def decorator(
            func: Callable[P, Awaitable[EventSourceResponse]]
        ) -> Callable[P, Awaitable[EventSourceResponse]]:
            @wraps(func)
            async def wrapper(request: Request) -> EventSourceResponse:
                async with sse_response(request) as resp:
                    args_to_apply = await load(
                        request, signature(func).parameters.copy()
                    )
                    definitive_args = {}
                    for name, param in signature(func).parameters.items():
                        if param.annotation == EventSourceResponse:
                            definitive_args[name] = resp
                        elif name in args_to_apply:
                            definitive_args[name] = args_to_apply[name]
                            args_to_apply.pop(name)
                        elif param.default is not param.empty:
                            definitive_args[name] = param.default
                        else:
                            raise ValueError(
                                f"Missing parameter {name} for {func.__name__}"
                            )
                    await func(**definitive_args)
                    return resp

            self.router.add_get(path, self.document(path, "OPTIONS")(wrapper))
            return wrapper

        return decorator

    def websocket(
        self, path: str
    ) -> Callable[
        [Callable[P, Awaitable[WebSocketResponse]]],
        Callable[P, Awaitable[WebSocketResponse]],
    ]:
        """

        Websocket decorator.

        :param path: The API path for the websocket request.
        :type path: str
        :return: Decorated function for handling the websocket request.

        This decorator will handle websocket connections and pass the websocket object to the decorated function, be sure to include a parameter with the type annotation `WebSocketResponse` in your function signature and also setup proper error handling, since `PING` and `PONG` messages are handled automatically by the server and will result on a `MessageTypeError` if not handled.

        """

        def decorator(
            func: Callable[P, Awaitable[WebSocketResponse]]
        ) -> Callable[P, Awaitable[WebSocketResponse]]:
            @wraps(func)
            async def wrapper(request: Request):
                args_to_apply = await load(request, signature(func).parameters.copy())
                ws = WebSocketResponse()
                await ws.prepare(request)
                definitive_args = {}
                for name, param in signature(func).parameters.items():
                    if param.annotation == WebSocketResponse:
                        definitive_args[name] = ws
                    elif name in args_to_apply:
                        definitive_args[name] = args_to_apply[name]
                        args_to_apply.pop(name)
                    elif param.default is not param.empty:
                        definitive_args[name] = param.default
                    else:
                        raise ValueError(
                            f"Missing parameter {name} for {func.__name__}"
                        )
                await func(**definitive_args)
                return ws

            self.router.add_get(path, wrapper)
            return wrapper

        return decorator

    def static(self, path: str = "", dir: str = "static") -> "APIServer":
        """

        Static files decorator.

        This decorator will serve static files from the `static` directory.
        It will also serve the `index.html` file from the `static` directory when the root path is requested.
        Useful for serving a frontend single-page application.

        :return: The server instance.

        """

        @self.get(f"/{path}")
        async def index(request: Request):
            base_url = request.url.path
            return HTTPFound(location=f"{base_url}/{path}/index.html")

        try:
            os.makedirs(dir, exist_ok=True)
        except OSError:
            pass
        self.router.add_static(f"/{path}", dir)

        return self

    def middleware(self, func: Middleware) -> Middleware:
        @wraps(func)
        @middleware
        async def wrapper(request: Request, handler: Handler) -> Response:
            response = await func(request, handler)
            if isinstance(response, Response):
                return response
            return do_response(response)

        self.middlewares.append(wrapper)
        return wrapper

    def use(self, router: APIRouter, prefix: str = "") -> "APIServer":
        """Chainable method to add a router to the server"""
        for route in router.routes:
            path = prefix + route.path
            self.router.add_route(
                route.method,
                prefix + router.prefix + path,
                route.handler,
                **route.kwargs,
            )

        for path, details in router.openapi["paths"].items():
            path = prefix + path
            self.openapi["paths"][path] = details

        self.openapi["components"]["schemas"].update(
            router.openapi["components"]["schemas"]
        )
        return self
