"""Python entrypoint of router."""

from collections.abc import AsyncGenerator
from http import HTTPStatus
from typing import List, Literal, Optional, Type  # noqa: UP035

import fastapi
import uvicorn
from fastapi.middleware.cors import CORSMiddleware

from mlc_llm.protocol import error_protocol
from mlc_llm.protocol.openai_api_protocol import CompletionLogProbs, CompletionRequest
from mlc_llm.router import Router
from mlc_llm.serve import engine_base, engine_utils


def serve(
    model: str,
    model_lib: Optional[str],
    router_host: str,
    router_port: int,
    endpoint_hosts: List[str],  # noqa: UP006
    endpoint_ports: List[int],  # noqa: UP006
    endpoint_num_gpus: List[int],  # noqa: UP006
    enable_prefix_cache: bool,
    router_mode: Literal["disagg", "round-robin"] = "round-robin",
    pd_balance_factor: float = 0.0,
    router_type: Type[Router] = Router,  # noqa: UP006
):
    """Start the router with the specified configuration."""
    # 1. Instantiate router
    router = router_type(
        model=model,
        model_lib=model_lib,
        hosts=endpoint_hosts,
        ports=endpoint_ports,
        num_gpus=endpoint_num_gpus,
        enable_prefix_cache=enable_prefix_cache,
        router_mode=router_mode,
        pd_balance_factor=pd_balance_factor,
    )

    router_app = fastapi.APIRouter()

    @router_app.post("/v1/completions")
    async def request_completion(request: CompletionRequest, raw_request: fastapi.Request):
        """OpenAI-compatible completion API.
        API reference: https://platform.openai.com/docs/api-reference/completions/create
        """
        if router is None:
            return error_protocol.create_error_response(
                HTTPStatus.BAD_REQUEST, message="Router is not initialized."
            )
        request_id = f"cmpl-{engine_utils.random_uuid()}"

        # Streaming response.
        if request.stream:
            # We manually get the first response from generator to
            # capture potential exceptions in this scope, rather then
            # the StreamingResponse scope.
            stream_generator = router.handle_completion(request, request_id)
            first_response = await anext(  # noqa: F821
                stream_generator
            )

            async def completion_stream_generator() -> AsyncGenerator[str, None]:
                if isinstance(first_response, StopAsyncIteration):
                    yield "data: [DONE]\n\n"
                    return
                yield f"data: {first_response.model_dump_json(by_alias=True)}\n\n"
                async for response in stream_generator:
                    yield f"data: {response.model_dump_json(by_alias=True)}\n\n"
                yield "data: [DONE]\n\n"

            return fastapi.responses.StreamingResponse(
                completion_stream_generator(), media_type="text/event-stream"
            )

        # FIXME: Non-streaming response not fully implemented
        request_final_usage = None
        output_texts = [""] * request.n
        finish_reasons: List[Optional[str]] = [None] * request.n  # noqa: UP006
        logprob_results: List[Optional[CompletionLogProbs]] = [None] * request.n  # noqa: UP006

        async for response in router.handle_completion(request, request_id):
            if await raw_request.is_disconnected():
                # In non-streaming cases, the engine will not be notified
                # when the request is disconnected.
                # Therefore, we check if it is disconnected each time,
                # and explicitly return.
                # Note that requesta abort is triggered when the async for and funciton scope ends.
                return error_protocol.create_error_response(
                    HTTPStatus.BAD_REQUEST, message="The request has disconnected"
                )
            # TODO(Charlie): This is copied from engine.py --
            # why is it here? Non-streaming only has a single chunk right?
            # this is the final chunk
            # if response.usage is not None:
            #     request_final_usage = response.usage
            #     continue
            for choice in response.choices:
                output_texts[choice.index] += choice.text
                if choice.finish_reason is not None and finish_reasons[choice.index] is None:
                    finish_reasons[choice.index] = choice.finish_reason
                if choice.logprobs is not None:
                    logprob_results[choice.index] = choice.logprobs

        assert all(finish_reason is not None for finish_reason in finish_reasons)
        return engine_base.wrap_completion_response(
            request_id=request_id,
            model=request.model,
            output_texts=output_texts,
            finish_reasons=finish_reasons,
            logprob_results=logprob_results,
            usage=request_final_usage,
        )

    # 2. Set up app
    app = fastapi.FastAPI()
    app.add_middleware(CORSMiddleware)
    app.include_router(router_app)
    app.exception_handler(error_protocol.BadRequestError)(error_protocol.bad_request_error_handler)

    # 3. Run
    uvicorn.run(app, host=router_host, port=router_port, log_level="info")