import functools
import logging
from collections.abc import AsyncIterable, Awaitable
from typing import Any, Callable, Union

from fastapi import FastAPI, APIRouter
from sse_starlette import EventSourceResponse
from starlette.requests import Request
from starlette.responses import Response, JSONResponse

from aduib_rpc.server.context import ServerContext
from aduib_rpc.server.protocols.rpc import DefaultServerContextBuilder
from aduib_rpc.server.protocols.rpc.jsonrpc_app import ServerContextBuilder
from aduib_rpc.server.request_handlers import RequestHandler, RESTHandler
from aduib_rpc.utils.constant import DEFAULT_RPC_PATH

logger = logging.getLogger(__name__)

class  AduibRpcRestFastAPIApp:
    def __init__(self,
                 request_handler: RequestHandler,
                 context_builder: ServerContextBuilder | None = None
                    ):
        """Initializes the AduibRpcRestFastAPIApp.
        Args:
            request_handler (RequestHandler): The request handler to process incoming requests.
            context_builder (ServerContextBuilder | None): Optional context builder for request processing.
        """
        self.context_builder = context_builder or DefaultServerContextBuilder()
        self.request_handler = RESTHandler(request_handler)

    def build(self,
              rpc_path: str=DEFAULT_RPC_PATH,
              **kwargs: Any,
              ) -> FastAPI:
        """Builds and returns the FastAPI application with the necessary routes.
        Args:
            rpc_path (str): The path for the RPC endpoint. Defaults to '/'.
            **kwargs: Additional keyword arguments.
        Returns:
            FastAPI: The configured FastAPI application.
        """
        app = FastAPI(**kwargs)
        router = APIRouter()
        for (route, handler) in self.routes().items():
            router.add_api_route(
                f'{rpc_path}{route[0]}', handler, methods=[route[1]]
            )
        app.include_router(router)
        return app

    def routes(self) -> dict[tuple[str, str], Callable[[Request], Any]]:
        """Defines the routes for the FastAPI application.
        Returns:
            dict[tuple[str, str], Callable[[Request], Any]]: A mapping of (method, path) to handler functions.
        """
        routes: dict[tuple[str, str], Callable[[Request], Any]] = {
            ('/v1/message/completion', 'POST'): functools.partial(
                self._handle_requests,self.request_handler.on_message
            ),
            ('/v1/message/completion/stream', 'POST'): functools.partial(
                self._handle_streaming_requests, self.request_handler.on_stream_message
            ),
            # ('/v1/embeddings', 'POST'): functools.partial(
            #     self._handle_requests, self.request_handler.on_message
            # ),
        }
        return routes

    def _validate_rpc_path(self, request: Request, rpc_path: str) -> Response:
        """Validates the RPC path of the incoming request.
        Args:
            request (Request): The incoming FastAPI request.
            rpc_path (str): The expected RPC path.
        Returns:
            Response: A 404 response if the path is invalid, otherwise None.
        """
        if not request.url.path.startswith(rpc_path):
            return JSONResponse(
                status_code=404,
                content={"error": "Invalid RPC path."}
            )
        return None

    async def _handle_requests(self,
                        method: Union[Callable[[Request,ServerContext],Awaitable[Any]],Callable[[Request,ServerContext],AsyncIterable[Any]]],
                        request: Request
                        ) -> Response:
        """Handles incoming requests and delegates to the appropriate request handler method.
        Args:
            method (Callable[[Request,ServerContext],AsyncIterable[Any]]): The request handler method to invoke.
            request (Request): The incoming FastAPI request.
        Returns:
            Response: The response generated by the request handler.
        """
        if (error_response := self._validate_rpc_path(request, DEFAULT_RPC_PATH)) is not None:
            return error_response
        context = self.context_builder.build_context(request)
        response = await method(request, context)
        return JSONResponse(content=response)

    async def _handle_streaming_requests(self,
                        method: Union[Callable[[Request,ServerContext],Awaitable[Any]],Callable[[Request,ServerContext],AsyncIterable[Any]]],
                        request: Request
                        ) -> EventSourceResponse:
        """Handles incoming streaming requests and delegates to the appropriate request handler method.
        Args:
            method (Callable[[Request,ServerContext],AsyncIterable[Any]]): The request handler method to invoke.
            request (Request): The incoming FastAPI request.
        Returns:
            Response: The response generated by the request handler.
        """
        if (error_response := self._validate_rpc_path(request, DEFAULT_RPC_PATH)) is not None:
            return error_response
        context = self.context_builder.build_context(request)
        async def event_generator(
            stream: AsyncIterable[Any],
        ) -> AsyncIterable[dict[str, dict[str, Any]]]:
            async for item in stream:
                yield {'data': item}
        return EventSourceResponse(
            event_generator(method(request, context))
        )

