AstrBot/astrbot/api/web.py

from __future__ import annotations

import contextvars
from collections.abc import Callable, KeysView
from contextlib import contextmanager
from pathlib import Path
from typing import Any, Generic, TypeVar, overload

from fastapi.encoders import jsonable_encoder
from fastapi.responses import FileResponse, JSONResponse
from starlette.datastructures import Headers
from starlette.datastructures import UploadFile as StarletteUploadFile
from starlette.responses import StreamingResponse

ValueT = TypeVar("ValueT")
DefaultT = TypeVar("DefaultT")
ConvertedT = TypeVar("ConvertedT")


class PluginMultiDict(Generic[ValueT]):
    """Dictionary-like request values that preserves duplicate keys."""

    def __init__(self, pairs: list[tuple[str, ValueT]]) -> None:
        self._pairs = pairs

    @overload
    def get(self, key: str) -> ValueT | None: ...

    @overload
    def get(self, key: str, default: DefaultT) -> ValueT | DefaultT: ...

    @overload
    def get(
        self,
        key: str,
        default: DefaultT,
        type: Callable[[ValueT], ConvertedT],
    ) -> ConvertedT | DefaultT: ...

    def get(self, key: str, default: Any = None, type: Callable | None = None):
        """Return the last value for a key.

        Args:
            key: Value key to read.
            default: Value returned when the key is missing or conversion fails.
            type: Optional callable used to convert the value.

        Returns:
            The matching value, converted value, or default.
        """
        for item_key, item_value in reversed(self._pairs):
            if item_key != key:
                continue
            if type is None:
                return item_value
            try:
                return type(item_value)
            except (TypeError, ValueError):
                return default
        return default

    def getlist(self, key: str) -> list[ValueT]:
        """Return all values for a key.

        Args:
            key: Value key to read.

        Returns:
            Values in request order.
        """
        return [item_value for item_key, item_value in self._pairs if item_key == key]

    def keys(self) -> KeysView[str]:
        return dict.fromkeys(item_key for item_key, _ in self._pairs).keys()

    def values(self) -> list[ValueT]:
        return [self[key] for key in self.keys()]

    def items(self) -> list[tuple[str, ValueT]]:
        return [(key, self[key]) for key in self.keys()]

    def __contains__(self, key: str) -> bool:
        return any(item_key == key for item_key, _ in self._pairs)

    def __getitem__(self, key: str) -> ValueT:
        value = self.get(key)
        if value is None and key not in self:
            raise KeyError(key)
        return value

    def __bool__(self) -> bool:
        return bool(self._pairs)


class PluginUploadFile:
    """Uploaded file wrapper exposed to plugin Web API handlers."""

    def __init__(self, upload_file: StarletteUploadFile) -> None:
        self._upload_file: StarletteUploadFile = upload_file
        self.filename: str | None = upload_file.filename
        self.content_type: str | None = upload_file.content_type
        self.headers: Headers = upload_file.headers
        self.content_length: int | None = self._resolve_content_length()

    def _resolve_content_length(self) -> int | None:
        try:
            raw = self.headers.get("content-length")
            return int(raw) if raw else None
        except (TypeError, ValueError):
            return None

    async def save(self, destination: str | Path) -> None:
        """Save the uploaded file to disk.

        Args:
            destination: Destination file path.
        """
        path = Path(destination)
        try:
            await self._upload_file.seek(0)
        except Exception:
            pass
        with path.open("wb") as output:
            while True:
                chunk = await self._upload_file.read(1024 * 1024)
                if not chunk:
                    break
                output.write(chunk)

    async def read(self, size: int = -1) -> bytes:
        """Read bytes from the uploaded file.

        Args:
            size: Maximum number of bytes to read. Use -1 to read all bytes.

        Returns:
            File bytes.
        """
        return await self._upload_file.read(size)

    async def write(self, data: bytes) -> None:
        """Write bytes to the uploaded file object.

        Args:
            data: Bytes to write.
        """
        await self._upload_file.write(data)

    async def seek(self, offset: int) -> None:
        """Move the uploaded file cursor.

        Args:
            offset: Absolute byte offset.
        """
        await self._upload_file.seek(offset)

    async def close(self) -> None:
        """Close the uploaded file."""
        await self._upload_file.close()

    def __getattr__(self, key: str) -> Any:
        return getattr(self._upload_file, key)


class PluginRequest:
    """Request object exposed to plugin Web API handlers."""

    def __init__(
        self,
        request_: Any,
        *,
        path_params: dict[str, Any] | None = None,
        plugin_name: str | None = None,
        username: str | None = None,
    ) -> None:
        self._request: Any = request_
        self.method: str = request_.method
        self.path: str = request_.url.path
        self.headers: Headers = request_.headers
        self.cookies: dict[str, str] = request_.cookies
        self.content_type: str | None = request_.headers.get("content-type")
        self.client_host: str | None = request_.client.host if request_.client else None
        self.path_params: dict[str, Any] = path_params or {}
        self.plugin_name: str | None = plugin_name
        self.username: str | None = username
        self.query: PluginMultiDict[str] = PluginMultiDict[str](
            list(request_.query_params.multi_items())
        )
        self._form_cache: PluginMultiDict[str] | None = None
        self._files_cache: PluginMultiDict[PluginUploadFile] | None = None

    async def body(self) -> bytes:
        """Read the raw request body.

        Returns:
            Raw request body bytes.
        """
        return await self._request.body()

    async def json(self, default: DefaultT | None = None) -> Any | DefaultT | None:
        """Read the JSON request body.

        Args:
            default: Value returned when the request body cannot be parsed as JSON.

        Returns:
            Parsed JSON data or default.
        """
        try:
            return await self._request.json()
        except Exception:
            return default

    async def _load_form_parts(self) -> None:
        if self._form_cache is not None and self._files_cache is not None:
            return
        form = await self._request.form()
        form_pairs: list[tuple[str, str]] = []
        file_pairs: list[tuple[str, PluginUploadFile]] = []
        for key, value in form.multi_items():
            if isinstance(value, StarletteUploadFile):
                file_pairs.append((key, PluginUploadFile(value)))
            else:
                form_pairs.append((key, value))
        self._form_cache = PluginMultiDict(form_pairs)
        self._files_cache = PluginMultiDict(file_pairs)

    async def form(self) -> PluginMultiDict[str]:
        """Read form fields from a multipart or form-urlencoded request.

        Returns:
            Form values without uploaded files.
        """
        await self._load_form_parts()
        assert self._form_cache is not None
        return self._form_cache

    async def files(self) -> PluginMultiDict[PluginUploadFile]:
        """Read uploaded files from a multipart request.

        Returns:
            Uploaded file values.
        """
        await self._load_form_parts()
        assert self._files_cache is not None
        return self._files_cache


_request_var: contextvars.ContextVar[PluginRequest] = contextvars.ContextVar(
    "astrbot_plugin_web_request"
)


class PluginRequestProxy:
    """Typed proxy for the request bound to the current plugin Web handler."""

    def _get_current(self) -> PluginRequest:
        try:
            return _request_var.get()
        except LookupError as exc:
            raise RuntimeError(
                "astrbot.api.web.request is only available inside a plugin Web API "
                "handler."
            ) from exc

    @property
    def method(self) -> str:
        return self._get_current().method

    @property
    def path(self) -> str:
        return self._get_current().path

    @property
    def headers(self) -> Headers:
        return self._get_current().headers

    @property
    def cookies(self) -> dict[str, str]:
        return self._get_current().cookies

    @property
    def content_type(self) -> str | None:
        return self._get_current().content_type

    @property
    def client_host(self) -> str | None:
        return self._get_current().client_host

    @property
    def path_params(self) -> dict[str, Any]:
        return self._get_current().path_params

    @property
    def plugin_name(self) -> str | None:
        return self._get_current().plugin_name

    @property
    def username(self) -> str | None:
        return self._get_current().username

    @property
    def query(self) -> PluginMultiDict[str]:
        return self._get_current().query

    async def body(self) -> bytes:
        return await self._get_current().body()

    async def json(self, default: DefaultT | None = None) -> Any | DefaultT | None:
        return await self._get_current().json(default=default)

    async def form(self) -> PluginMultiDict[str]:
        return await self._get_current().form()

    async def files(self) -> PluginMultiDict[PluginUploadFile]:
        return await self._get_current().files()

    def __getattr__(self, key: str) -> Any:
        return getattr(self._get_current(), key)


request: PluginRequestProxy = PluginRequestProxy()


@contextmanager
def bind_request_context(request_: PluginRequest):
    """Bind a plugin Web request for the current async context.

    Args:
        request_: Request object exposed through the module-level request proxy.

    Yields:
        The bound request object.
    """
    token = _request_var.set(request_)
    try:
        yield request_
    finally:
        _request_var.reset(token)


def json_response(
    data: Any = None,
    *,
    status_code: int = 200,
    headers: dict[str, str] | None = None,
) -> JSONResponse:
    """Build a JSON response for plugin Web API handlers.

    Args:
        data: JSON-serializable response body.
        status_code: HTTP status code.
        headers: Optional response headers.

    Returns:
        A Starlette/FastAPI JSON response.
    """
    return JSONResponse(
        jsonable_encoder({} if data is None else data),
        status_code=status_code,
        headers=headers,
    )


def error_response(
    message: str,
    *,
    status_code: int = 400,
    data: Any = None,
    headers: dict[str, str] | None = None,
) -> JSONResponse:
    """Build a standard error response for plugin bridge calls.

    Args:
        message: Public error message.
        status_code: HTTP status code.
        data: Optional error details that are safe to expose.
        headers: Optional response headers.

    Returns:
        A JSON response with the AstrBot error envelope.
    """
    return json_response(
        {"status": "error", "message": message, "data": data},
        status_code=status_code,
        headers=headers,
    )


def file_response(
    path: str | Path,
    *,
    filename: str | None = None,
    content_type: str | None = None,
    headers: dict[str, str] | None = None,
) -> FileResponse:
    """Build a file download response for plugin Web API handlers.

    Args:
        path: File path to send.
        filename: Optional download filename.
        content_type: Optional response media type.
        headers: Optional response headers.

    Returns:
        A Starlette/FastAPI file response.
    """
    return FileResponse(
        path,
        filename=filename,
        media_type=content_type,
        headers=headers,
    )


def stream_response(
    content: Any,
    *,
    content_type: str = "text/event-stream",
    status_code: int = 200,
    headers: dict[str, str] | None = None,
) -> StreamingResponse:
    """Build a streaming response for plugin Web API handlers.

    Args:
        content: Sync or async iterable that yields response chunks.
        content_type: Response media type.
        status_code: HTTP status code.
        headers: Optional response headers.

    Returns:
        A Starlette/FastAPI streaming response.
    """
    return StreamingResponse(
        content,
        media_type=content_type,
        status_code=status_code,
        headers=headers,
    )


__all__ = [
    "PluginMultiDict",
    "PluginRequest",
    "PluginRequestProxy",
    "PluginUploadFile",
    "bind_request_context",
    "error_response",
    "file_response",
    "json_response",
    "request",
    "stream_response",
]