SIMNet/simnet/client/components/calculator/genshin.py

from typing import Optional, Any, Dict, List

from simnet.client.base import BaseClient
from simnet.client.routes import CALCULATOR_URL
from simnet.errors import BadRequest
from simnet.models.genshin.calculator import (
    CalculatorResult,
    CalculatorCharacter,
    CalculatorWeapon,
    CalculatorArtifact,
    CalculatorFurnishing,
    CalculatorCharacterDetails,
    CalculatorTalent,
)
from simnet.utils.enums import Region
from simnet.utils.player import recognize_genshin_server


class CalculatorClient(BaseClient):
    """A client for retrieving data from Genshin's calculator component."""

    async def request_calculator(
        self,
        endpoint: str,
        *,
        method: str = "POST",
        lang: Optional[str] = None,
        params: Optional[Dict[str, Any]] = None,
        data: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        """Make a request towards the calculator endpoint.

        Args:
            endpoint (str): The calculator endpoint to send the request to.
            method (str): The HTTP method to use for the request (default "POST").
            lang (str): The language to use for the request (default None).
            params (dict): The parameters to include in the request URL (default None).
            data (dict): The data to include in the request body (default None).

        Returns:
            dict: The data returned by the calculator endpoint.
        """
        params = dict(params or {})

        base_url = CALCULATOR_URL.get_url(self.region)
        url = base_url / endpoint

        if method == "GET":
            params["lang"] = lang or self.lang
            data = None
        else:
            data = dict(data or {})
            data["lang"] = lang or self.lang

        headers = {}
        if self.region == Region.CHINESE:
            headers["Referer"] = (
                "https://webstatic.mihoyo.com/ys/event/e20200923adopt_calculator/index.html?"
                "bbs_presentation_style=fullscreen&bbs_auth_required=true&"
                "utm_source=bbs&utm_medium=mys&utm_campaign=icon#/"
            )

        data = await self.request_lab(url, method=method, params=params, data=data, headers=headers)

        return data

    async def _execute_calculator(
        self,
        data: Dict[str, Any],
        *,
        lang: Optional[str] = None,
    ) -> CalculatorResult:
        """Calculate the results of a builder.

        Args:
            data (dict): The data used to calculate the results.
            lang (str): The language to use for the request (default None).

        Returns:
            CalculatorResult: The calculated results.
        """
        data = await self.request_calculator("compute", lang=lang, data=data)
        return CalculatorResult(**data)

    async def _enable_calculator_sync(self, enabled: bool = True) -> None:
        """Enable data syncing in calculator.

        Args:
            enabled (bool): Whether to enable syncing (default True).
        """
        await self.request_calculator("avatar/auth", method="POST", data=dict(avatar_auth=int(enabled)))

    async def _get_calculator_items(
        self,
        slug: str,
        filters: Dict[str, Any],
        query: Optional[str] = None,
        *,
        player_id: Optional[int] = None,
        is_all: bool = False,
        sync: bool = False,
        lang: Optional[str] = None,
        autoauth: bool = True,
    ) -> List[Dict[str, Any]]:
        """Get all items of a specific slug from a calculator.

        Args:
            slug (str): The slug to get the items for.
            filters (dict): The filters to apply to the items.
            query (str): The query to search for (default None).
            player_id (int): The player ID to use for syncing (default None).
            is_all (bool): Whether to include Traveler items (default False).
            sync (bool): Whether to sync data from the calculator (default False).
            lang (str): The language to use for the request (default None).
            autoauth (bool): Whether to enable syncing if it is not already enabled (default True).

        Returns:
            list: A list of dictionaries representing the items retrieved from the calculator.
        """
        endpoint = f"sync/{slug}/list" if sync else f"{slug}/list"

        if query:
            if any(filters.values()):
                raise TypeError("Cannot specify a query and filter at the same time")

            filters = dict(keywords=query, **filters)

        payload: Dict[str, Any] = dict(page=1, size=69420, is_all=is_all, **filters)

        if sync:
            player_id = player_id or self.player_id
            payload["uid"] = player_id
            payload["region"] = recognize_genshin_server(player_id)

        try:
            data = await self.request_calculator(endpoint, lang=lang, data=payload)
        except BadRequest as e:
            if e.ret_code != -502002:  # Sync not enabled
                raise
            if not autoauth:
                raise BadRequest(e.response, "Calculator sync is not enabled") from e

            await self._enable_calculator_sync()
            data = await self.request_calculator(endpoint, lang=lang, data=payload)

        return data["list"]

    async def get_calculator_characters(
        self,
        *,
        query: Optional[str] = None,
        elements: Optional[List[int]] = None,
        weapon_types: Optional[List[int]] = None,
        include_traveler: bool = False,
        sync: bool = False,
        player_id: Optional[int] = None,
        lang: Optional[str] = None,
    ) -> List[CalculatorCharacter]:
        """Get all characters provided by the Enhancement Progression Calculator.

        Args:
            query (str): The query to search for (default None).
            elements (list): A list of element IDs to filter by (default None).
            weapon_types (list): A list of weapon type IDs to filter by (default None).
            include_traveler (bool): Whether to include Traveler characters (default False).
            sync (bool): Whether to sync data from the calculator (default False).
            player_id (int): The player ID to use for syncing (default None).
            lang (str): The language to use for the request (default None).

        Returns:
            list: A list of CalculatorCharacter objects representing the characters retrieved from the calculator.
        """
        data = await self._get_calculator_items(
            "avatar",
            lang=lang,
            is_all=include_traveler,
            sync=sync,
            player_id=player_id,
            query=query,
            filters=dict(
                element_attr_ids=elements or [],
                weapon_cat_ids=weapon_types or [],
            ),
        )
        return [CalculatorCharacter(**i) for i in data]

    async def get_calculator_weapons(
        self,
        *,
        query: Optional[str] = None,
        types: Optional[List[int]] = None,
        rarities: Optional[List[int]] = None,
        lang: Optional[str] = None,
    ) -> List[CalculatorWeapon]:
        """Get all weapons provided by the Enhancement Progression Calculator.

        Args:
            query (Optional[str], optional): A query string to filter the results. Defaults to None.
            types (Optional[List[int]], optional): A list of weapon types to include in the results. Defaults to None.
            rarities (Optional[List[int]], optional): A list of weapon rarities to include in the results.
                Defaults to None.
            lang (Optional[str], optional): The language to use for the calculator. Defaults to None.

        Returns:
            List[CalculatorWeapon]: A list of weapons provided by the Enhancement Progression Calculator.
        """
        data = await self._get_calculator_items(
            "weapon",
            lang=lang,
            query=query,
            filters=dict(
                weapon_cat_ids=types or [],
                weapon_levels=rarities or [],
            ),
        )
        return [CalculatorWeapon(**i) for i in data]

    async def get_calculator_furnishings(
        self,
        *,
        types: Optional[int] = None,
        rarities: Optional[int] = None,
        lang: Optional[str] = None,
    ) -> List[CalculatorFurnishing]:
        """
        Get all furnishings provided by the Enhancement Progression Calculator.

        Args:
            types (Optional[int], optional): The type of furnishings to retrieve. Defaults to None.
            rarities (Optional[int], optional): The rarity of the furnishings to retrieve. Defaults to None.
            lang (Optional[str], optional): The language to use for the calculator. Defaults to None.

        Returns:
            List[CalculatorFurnishing]: A list of furnishings provided by the Enhancement Progression Calculator.
        """
        data = await self._get_calculator_items(
            "furniture",
            lang=lang,
            filters=dict(
                cat_id=types or 0,
                weapon_levels=rarities or 0,
            ),
        )
        return [CalculatorFurnishing(**i) for i in data]

    async def get_character_details(
        self,
        character: int,
        *,
        player_id: Optional[int] = None,
        lang: Optional[str] = None,
    ) -> CalculatorCharacterDetails:
        """
        Get the weapon, artifacts and talents of a character.

        Args:
            character (int): The ID of the character to retrieve details for.
            player_id (Optional[int], optional): The player ID to use for the request. Defaults to None.
            lang (Optional[str], optional): The language to use for the calculator. Defaults to None.

        Returns:
            CalculatorCharacterDetails: The details of the character.
        """
        player_id = player_id or self.player_id

        data = await self.request_calculator(
            "sync/avatar/detail",
            method="GET",
            lang=lang,
            params=dict(
                avatar_id=int(character),
                uid=player_id,
                region=recognize_genshin_server(player_id),
            ),
        )
        return CalculatorCharacterDetails(**data)

    async def get_character_talents(
        self,
        character: int,
        *,
        lang: Optional[str] = None,
    ) -> List[CalculatorTalent]:
        """Get the talents of a character.

        Args:
            character (int): The ID of the character to retrieve talents for.
            lang (Optional[str], optional): The language to use for the calculator. Defaults to None.

        Returns:
            List[CalculatorTalent]: A list of talents for the specified character.
        """
        data = await self.request_calculator(
            "avatar/skill_list",
            method="GET",
            lang=lang,
            params=dict(avatar_id=int(character)),
        )
        return [CalculatorTalent(**i) for i in data["list"]]

    async def get_complete_artifact_set(
        self,
        artifact: int,
        *,
        lang: Optional[str] = None,
    ) -> List[CalculatorArtifact]:
        """
        Get all artifacts that share a set with a specified artifact.

        Args:
            artifact (int): The ID of the artifact to retrieve the set for.
            lang (Optional[str], optional): The language to use for the calculator. Defaults to None.

        Returns:
            List[CalculatorArtifact]: A list of artifacts that share a set with the specified artifact.
        """
        data = await self.request_calculator(
            "reliquary/set",
            method="GET",
            lang=lang,
            params=dict(reliquary_id=int(artifact)),
        )
        return [CalculatorArtifact(**i) for i in data["reliquary_list"]]

    async def _get_all_artifact_ids(self, artifact_id: int) -> List[int]:
        """Get all artifact IDs in the same set as a given artifact ID.

        Args:
            artifact_id (int): The ID of the artifact to retrieve the set for.

        Returns:
            List[int]: A list of artifact IDs that share a set with the specified artifact ID.
        """
        others = await self.get_complete_artifact_set(artifact_id)
        return [artifact_id] + [other.id for other in others]

    async def get_teapot_replica_blueprint(
        self,
        share_code: int,
        *,
        region: Optional[str] = None,
        player_id: Optional[int] = None,
        lang: Optional[str] = None,
    ) -> List[CalculatorFurnishing]:
        """Get the furnishings used by a teapot replica blueprint.

        Args:
            share_code (int): The share code of the teapot replica blueprint to retrieve furnishings for.
            region (Optional[str], optional): The region to use for the request. Defaults to None.
            player_id (Optional[int], optional): The player ID to use for the request. Defaults to None.
            lang (Optional[str], optional): The language to use for the calculator. Defaults to None.

        Returns:
            List[CalculatorFurnishing]: A list of furnishings used by the specified teapot replica blueprint.
        """
        if not region:
            player_id = player_id or self.player_id
            region = recognize_genshin_server(player_id)

        data = await self.request_calculator(
            "furniture/blueprint",
            method="GET",
            lang=lang,
            params=dict(share_code=share_code, region=region),
        )
        return [CalculatorFurnishing(**i) for i in data["list"]]