API Versioning

import functools
import typing

from fastapi import FastAPI
from pydantic import BaseModel
from ratelimit.types import Scope
from starlette.datastructures import Headers
from starlette.routing import BaseRoute, Match, Route
from starlette.types import ASGIApp, Receive, Send

from polar.openapi import APITag
from polar.routing import APIRouter


class APIVersion:
    __slots__ = ("month", "year")
    year: int
    month: int

    def __init__(self, year: int, month: int):
        self.year = year
        self.month = month

    def __str__(self) -> str:
        return f"{self.year}-{self.month:02d}"

    def __eq__(self, value: object, /) -> bool:
        if not isinstance(value, APIVersion):
            return NotImplemented
        return value.year == self.year and value.month == self.month

    def __lt__(self, value: object, /) -> bool:
        if not isinstance(value, APIVersion):
            return NotImplemented
        if self.year != value.year:
            return self.year < value.year
        return self.month < value.month

    def __le__(self, value: object, /) -> bool:
        if not isinstance(value, APIVersion):
            return NotImplemented
        if self.year != value.year:
            return self.year < value.year
        return self.month <= value.month

    def __gt__(self, value: object, /) -> bool:
        if not isinstance(value, APIVersion):
            return NotImplemented
        if self.year != value.year:
            return self.year > value.year
        return self.month > value.month

    def __ge__(self, value: object, /) -> bool:
        if not isinstance(value, APIVersion):
            return NotImplemented
        if self.year != value.year:
            return self.year > value.year
        return self.month >= value.month

    @classmethod
    def parse(cls, version_str: str) -> "APIVersion":
        try:
            year_str, month_str = version_str.split("-")
            return cls(year=int(year_str), month=int(month_str))
        except ValueError as e:
            raise ValueError(f"Invalid version string: {version_str}") from e


def version(*versions: APIVersion):
    def decorator(func):
        @functools.wraps(func)
        async def wrapper(*args, **kwargs):
            return await func(*args, **kwargs)

        wrapper._api_versions = versions  # type: ignore[attr-defined]
        return wrapper

    return decorator


class VersionedAPIRouter(APIRouter):
    def __init__(self, version: APIVersion, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.version = version

    def add_api_route(self, path, endpoint, *, methods=None, **kwargs):
        endpoint_versions = getattr(endpoint, "_api_versions", None)
        conflicting_route = self._get_conflicting_route(path, methods)

        if endpoint_versions is None:
            if conflicting_route is not None:
                conflicting_route_versions = getattr(
                    conflicting_route.endpoint, "_api_versions", None
                )
                if conflicting_route_versions is not None:
                    return
        elif self.version not in endpoint_versions:
            return

        if conflicting_route is not None:
            self.routes.remove(conflicting_route)

        super().add_api_route(path, endpoint, methods=methods, **kwargs)

    def _get_conflicting_route(
        self, path: str, methods: set[str] | list[str] | None
    ) -> Route | None:
        for route in self.routes:
            if not isinstance(route, Route):
                continue
            if route.path == path and route.methods == set(methods or []):
                return route
        return None


class VersionRoute(BaseRoute):
    def __init__(
        self,
        version: APIVersion,
        current_version: APIVersion,
        app: ASGIApp,
        name: str | None = None,
    ) -> None:
        self.version = version
        self.current_version = current_version
        self.app = app
        self.name = name

    @property
    def routes(self) -> list[BaseRoute]:
        return getattr(self.app, "routes", [])

    def matches(self, scope: Scope) -> tuple[Match, Scope]:
        if scope["type"] in ("http", "websocket"):
            headers = Headers(scope=scope)
            raw_version = headers.get("polar-version")
            try:
                request_version = (
                    APIVersion.parse(raw_version)
                    if raw_version
                    else self.current_version
                )
            except ValueError:
                return Match.NONE, scope

            if request_version == self.version:
                child_scope = {"endpoint": self.app}
                return Match.FULL, child_scope
        return Match.NONE, {}

    async def handle(self, scope: Scope, receive: Receive, send: Send) -> None:
        await self.app(scope, receive, send)

    def __eq__(self, other: typing.Any) -> bool:
        if not isinstance(other, VersionRoute):
            return NotImplemented
        return self.version == other.version and self.app == other.app

    def __repr__(self) -> str:
        class_name = self.__class__.__name__
        name = self.name or ""
        return (
            f"{class_name}(version={self.version!r}, name={name!r}, app={self.app!r})"
        )