# Copyright 2026 Marimo. All rights reserved.
from __future__ import annotations

from asyncio import iscoroutine
from collections.abc import Callable
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, TypeVar

from starlette.responses import (
    FileResponse,
    Response,
)
from starlette.routing import Mount, Router

from marimo import _loggers
from marimo._server.responses import StructResponse

if TYPE_CHECKING:
    from starlette.requests import Request

LOGGER = _loggers.marimo_logger()

DecoratedCallable = TypeVar("DecoratedCallable", bound=Callable[..., Any])


@dataclass
class APIRouter(Router):
    def __init__(self, prefix: str = "") -> None:
        self.prefix = prefix
        super().__init__()

    def __post_init__(self) -> None:
        if self.prefix:
            assert self.prefix.startswith("/"), (
                "Path prefix must start with '/'"
            )
            assert not self.prefix.endswith("/"), (
                "Path prefix must not end with '/'"
            )

    def post(
        self, path: str, include_in_schema: bool = True
    ) -> Callable[[DecoratedCallable], DecoratedCallable]:
        """Post method that returns a JSON response"""

        def decorator(func: DecoratedCallable) -> DecoratedCallable:
            async def wrapper_func(request: Request) -> Response:
                response = await func(request=request)
                if isinstance(response, FileResponse):
                    return response

                # If we already have a response, just return it
                if isinstance(response, Response):
                    return response

                # Otherwise it's one of our structs that we need to encode
                return StructResponse(response)

            # Set docstring of wrapper_func to the docstring of func
            wrapper_func.__doc__ = func.__doc__

            self.add_route(
                path=self.prefix + path,
                endpoint=wrapper_func,
                methods=["POST"],
                include_in_schema=include_in_schema,
            )

            return wrapper_func  # type: ignore[return-value]

        return decorator

    def get(
        self,
        path: str,
        include_in_schema: bool = True,
        name: str | None = None,
    ) -> Callable[[DecoratedCallable], DecoratedCallable]:
        """Get method."""

        def decorator(func: DecoratedCallable) -> DecoratedCallable:
            async def wrapper_func(request: Request) -> Response:
                response = func(request=request)
                if iscoroutine(response):
                    response = await response

                # If we already have a response, just return it
                if isinstance(response, Response):
                    return response

                # Otherwise it's one of our structs that we need to encode
                return StructResponse(response)

            # Set docstring of wrapper_func to the docstring of func
            wrapper_func.__doc__ = func.__doc__

            self.add_route(
                path=self.prefix + path,
                endpoint=wrapper_func,
                methods=["GET"],
                include_in_schema=include_in_schema,
                name=name,
            )
            return func

        return decorator

    def delete(
        self, path: str, include_in_schema: bool = True
    ) -> Callable[[DecoratedCallable], DecoratedCallable]:
        """Delete method."""

        def decorator(func: DecoratedCallable) -> DecoratedCallable:
            self.add_route(
                path=self.prefix + path,
                endpoint=func,
                methods=["DELETE"],
                include_in_schema=include_in_schema,
            )
            return func

        return decorator

    def websocket(
        self, path: str
    ) -> Callable[[DecoratedCallable], DecoratedCallable]:
        """Websocket method."""

        def decorator(func: DecoratedCallable) -> DecoratedCallable:
            self.add_websocket_route(path=self.prefix + path, endpoint=func)
            return func

        return decorator

    def include_router(
        self, router: APIRouter, prefix: str = "", name: str | None = None
    ) -> None:
        """Include another router in this one."""
        # Merge Mounts with the same path
        resolved_prefix = self.prefix + prefix
        for route in self.routes:
            if isinstance(route, Mount) and route.path == resolved_prefix:
                # NOTE: We don't merge middleware here, because it's not
                # clear what the correct behavior is.
                route.routes.extend(router.routes)
                return

        self.mount(path=resolved_prefix, app=router, name=name)
