import base64
import json
from collections.abc import AsyncIterator
from typing import Annotated

from fastapi import Body, FastAPI, HTTPException, Path, Response
from fastapi.responses import StreamingResponse
from pydantic import BaseModel

from ipython_shell.jupyter import InputRequest, JupyterApp, ParsedJupyterMessage

app = FastAPI(title="IPython Jupyter Web App", version="0.1.0")
shell_app = JupyterApp()


class InputReply(BaseModel):
    """Value supplied by a client answering a kernel input request."""

    value: str


def serialize_event(event: ParsedJupyterMessage | InputRequest) -> dict[str, object]:
    """Turn a shell event into JSON-safe data for a web client."""
    if isinstance(event, InputRequest):
        return {
            "type": "input_request",
            "call_id": event.call_id,
            "prompt": event.prompt,
            "password": event.password,
        }

    # Some Jupyter messages keep binary payloads in message buffers. Base64
    # makes those buffers safe to carry beside JSON MIME data in the stream.
    content = event.content
    if isinstance(content, BaseModel):
        content = content.model_dump(mode="json")

    return {
        "type": "kernel_message",
        "msg_type": event.msg_type,
        "parent_id": event.parent_id,
        "content": content,
        "buffers": [
            base64.b64encode(buffer).decode("ascii") for buffer in event.buffers
        ],
    }


async def _stream_events(code: str, shell_name: str) -> AsyncIterator[str]:
    """Encode one execution's events as newline-delimited JSON."""
    async for event in shell_app.run_code_stream(code, shell_name):
        yield json.dumps(serialize_event(event), ensure_ascii=False) + "\n"


@app.post("/shells/{shell_name}/run")
async def run_code(
    shell_name: Annotated[
        str,
        Path(
            ...,
            description="The Jupyter shell to run code in. Existing shell guess. 'new' creates a new shell, 'last' uses the most recent shell.",
        ),
    ],
    code: Annotated[
        str,
        Body(..., description="The code to execute", media_type="text/plain"),
    ],
) -> StreamingResponse:
    """Run code and stream each Jupyter event as one JSON line."""
    return StreamingResponse(
        _stream_events(code, shell_name),
        media_type="application/x-ndjson",
    )


@app.post("/calls/{call_id}/input", status_code=204)
async def reply_to_input(call_id: str, reply: InputReply) -> Response:
    """Answer an input request belonging to a currently running call."""
    try:
        await shell_app.reply_to_input(call_id, reply.value)
    except KeyError as error:
        raise HTTPException(status_code=404, detail=str(error)) from error
    except RuntimeError as error:
        raise HTTPException(status_code=409, detail=str(error)) from error
    return Response(status_code=204)


def main() -> None:
    """Run the Jupyter-backed HTTP API with Uvicorn."""
    import uvicorn

    uvicorn.run(app, host="::", port=8001, log_level="info")


if __name__ == "__main__":
    main()