g0v0-server/app/calculators/performance/performance_server.py

import asyncio
import datetime
from typing import TYPE_CHECKING, TypedDict, cast

from app.models.mods import APIMod
from app.models.performance import (
    DifficultyAttributes,
    DifficultyAttributesUnion,
    PerformanceAttributes,
    PerformanceAttributesUnion,
)
from app.models.score import GameMode

from ._base import (
    AvailableModes,
    CalculateError,
    DifficultyError,
    PerformanceCalculator as BasePerformanceCalculator,
    PerformanceError,
)

from httpx import AsyncClient, HTTPError
from pydantic import TypeAdapter

if TYPE_CHECKING:
    from app.database.score import Score


class AvailableRulesetResp(TypedDict):
    has_performance_calculator: list[str]
    has_difficulty_calculator: list[str]
    loaded_rulesets: list[str]


class PerformanceServerPerformanceCalculator(BasePerformanceCalculator):
    def __init__(self, server_url: str = "http://localhost:5225") -> None:
        self.server_url = server_url

        self._available_modes: AvailableModes | None = None
        self._modes_lock = asyncio.Lock()
        self._today = datetime.date.today()

    async def init(self):
        await self.get_available_modes()

    def _process_modes(self, modes: AvailableRulesetResp) -> AvailableModes:
        performance_modes = {
            m for mode in modes["has_performance_calculator"] if (m := GameMode.parse(mode)) is not None
        }
        difficulty_modes = {m for mode in modes["has_difficulty_calculator"] if (m := GameMode.parse(mode)) is not None}
        if GameMode.OSU in performance_modes:
            performance_modes.add(GameMode.OSURX)
            performance_modes.add(GameMode.OSUAP)
        if GameMode.TAIKO in performance_modes:
            performance_modes.add(GameMode.TAIKORX)
        if GameMode.FRUITS in performance_modes:
            performance_modes.add(GameMode.FRUITSRX)

        return AvailableModes(
            has_performance_calculator=performance_modes,
            has_difficulty_calculator=difficulty_modes,
        )

    async def get_available_modes(self) -> AvailableModes:
        # https://github.com/GooGuTeam/osu-performance-server#get-available_rulesets
        if self._available_modes is not None and self._today == datetime.date.today():
            return self._available_modes
        async with self._modes_lock, AsyncClient() as client:
            try:
                resp = await client.get(f"{self.server_url}/available_rulesets")
                if resp.status_code != 200:
                    raise CalculateError(f"Failed to get available modes: {resp.text}")
                modes = cast(AvailableRulesetResp, resp.json())
                result = self._process_modes(modes)

                self._available_modes = result
                self._today = datetime.date.today()
                return result
            except HTTPError as e:
                raise CalculateError(f"Failed to get available modes: {e}") from e
            except Exception as e:
                raise CalculateError(f"Unknown error: {e}") from e

    async def calculate_performance(self, beatmap_raw: str, score: "Score") -> PerformanceAttributes:
        # https://github.com/GooGuTeam/osu-performance-server#post-performance
        async with AsyncClient(timeout=15) as client:
            try:
                resp = await client.post(
                    f"{self.server_url}/performance",
                    json={
                        "beatmap_id": score.beatmap_id,
                        "beatmap_file": beatmap_raw,
                        "checksum": score.map_md5,
                        "accuracy": score.accuracy,
                        "combo": score.max_combo,
                        "mods": score.mods,
                        "statistics": {
                            "great": score.n300,
                            "ok": score.n100,
                            "meh": score.n50,
                            "miss": score.nmiss,
                            "perfect": score.ngeki,
                            "good": score.nkatu,
                            "large_tick_hit": score.nlarge_tick_hit or 0,
                            "large_tick_miss": score.nlarge_tick_miss or 0,
                            "small_tick_hit": score.nsmall_tick_hit or 0,
                            "slider_tail_hit": score.nslider_tail_hit or 0,
                        },
                        "ruleset": score.gamemode.to_base_ruleset().value,
                    },
                )
                if resp.status_code != 200:
                    raise PerformanceError(f"Failed to calculate performance: {resp.text}")
                return TypeAdapter(PerformanceAttributesUnion).validate_json(resp.text)
            except HTTPError as e:
                raise PerformanceError(f"Failed to calculate performance: {e}") from e
            except Exception as e:
                raise CalculateError(f"Unknown error: {e}") from e

    async def calculate_difficulty(
        self, beatmap_raw: str, mods: list[APIMod] | None = None, gamemode: GameMode | None = None
    ) -> DifficultyAttributes:
        # https://github.com/GooGuTeam/osu-performance-server#post-difficulty
        async with AsyncClient(timeout=15) as client:
            try:
                resp = await client.post(
                    f"{self.server_url}/difficulty",
                    json={
                        "beatmap_file": beatmap_raw,
                        "mods": mods or [],
                        "ruleset": gamemode.value if gamemode else None,
                    },
                )
                if resp.status_code != 200:
                    raise DifficultyError(f"Failed to calculate difficulty: {resp.text}")
                return TypeAdapter(DifficultyAttributesUnion).validate_json(resp.text)
            except HTTPError as e:
                raise DifficultyError(f"Failed to calculate difficulty: {e}") from e
            except Exception as e:
                raise DifficultyError(f"Unknown error: {e}") from e


PerformanceCalculator = PerformanceServerPerformanceCalculator