MibooGram/core/dependence/assets.py

import asyncio
from pathlib import Path
from ssl import SSLZeroReturnError
from typing import Optional, List, Dict

from aiofiles import open as async_open
from httpx import AsyncClient, HTTPError

from core.base_service import BaseService
from modules.wiki.base import WikiModel
from modules.wiki.models.avatar import Avatar
from modules.wiki.models.weapon import Weapon
from modules.wiki.models.buddy import Buddy
from modules.wiki.models.equipment_suit import EquipmentSuit
from utils.const import PROJECT_ROOT
from utils.log import logger
from utils.typedefs import StrOrURL, StrOrInt

ASSETS_PATH = PROJECT_ROOT.joinpath("resources/assets")
ASSETS_PATH.mkdir(exist_ok=True, parents=True)
DATA_MAP = {
    "avatar": WikiModel.BASE_URL + "avatars.json",
    "weapon": WikiModel.BASE_URL + "weapons.json",
    "buddy": WikiModel.BASE_URL + "buddy.json",
    "equipment_suit": WikiModel.BASE_URL + "equipment_suits.json",
}


class AssetsServiceError(Exception):
    pass


class AssetsCouldNotFound(AssetsServiceError):
    def __init__(self, message: str, target: str):
        self.message = message
        self.target = target
        super().__init__(f"{message}: target={target}")


class _AssetsService:
    client: Optional[AsyncClient] = None

    def __init__(self, client: Optional[AsyncClient] = None) -> None:
        self.client = client

    async def _download(self, url: StrOrURL, path: Path, retry: int = 5) -> Optional[Path]:
        """从 url 下载图标至 path"""
        if not url:
            return None
        logger.debug("正在从 %s 下载图标至 %s", url, path)
        headers = None
        for time in range(retry):
            try:
                response = await self.client.get(url, follow_redirects=False, headers=headers)
            except Exception as error:  # pylint: disable=W0703
                if not isinstance(error, (HTTPError, SSLZeroReturnError)):
                    logger.error(error)  # 打印未知错误
                if time != retry - 1:  # 未达到重试次数
                    await asyncio.sleep(1)
                else:
                    raise error
                continue
            if response.status_code != 200:  # 判定页面是否正常
                return None
            async with async_open(path, "wb") as file:
                await file.write(response.content)  # 保存图标
            return path.resolve()


class _AvatarAssets(_AssetsService):
    path: Path
    data: List[Avatar]
    name_map: Dict[str, Avatar]
    id_map: Dict[int, Avatar]

    def __init__(self, client: Optional[AsyncClient] = None) -> None:
        super().__init__(client)
        self.path = ASSETS_PATH.joinpath("agent")
        self.path.mkdir(exist_ok=True, parents=True)

    async def initialize(self):
        logger.info("正在初始化角色素材图标")
        html = await self.client.get(DATA_MAP["avatar"])
        self.data = [Avatar(**data) for data in html.json()]
        self.name_map = {icon.name: icon for icon in self.data}
        self.id_map = {icon.id: icon for icon in self.data}
        tasks = []
        for icon in self.data:
            base_path = self.path / f"{icon.id}"
            base_path.mkdir(exist_ok=True, parents=True)
            gacha_path = base_path / "gacha.png"
            icon_path = base_path / "icon.png"
            square_path = base_path / "square.png"
            normal_path = base_path / "normal.png"
            if not gacha_path.exists():
                tasks.append(self._download(icon.gacha, gacha_path))
            if not icon_path.exists():
                tasks.append(self._download(icon.icon_, icon_path))
            if not square_path.exists():
                tasks.append(self._download(icon.square, square_path))
            if not normal_path.exists():
                tasks.append(self._download(icon.normal, normal_path))

            if len(tasks) >= 100:
                await asyncio.gather(*tasks)
                tasks = []
        if tasks:
            await asyncio.gather(*tasks)
        logger.info("角色素材图标初始化完成")

    def get_path(self, icon: Avatar, name: str, ext: str = "png") -> Path:
        path = self.path / f"{icon.id}"
        path.mkdir(exist_ok=True, parents=True)
        return path / f"{name}.{ext}"

    def get_by_id(self, id_: int) -> Optional[Avatar]:
        return self.id_map.get(id_, None)

    def get_by_name(self, name: str) -> Optional[Avatar]:
        return self.name_map.get(name, None)

    def get_target(self, target: StrOrInt, second_target: StrOrInt = None) -> Avatar:
        data = None
        if isinstance(target, int):
            data = self.get_by_id(target)
        elif isinstance(target, str):
            data = self.get_by_name(target)
        if data is None:
            if second_target:
                return self.get_target(second_target)
            raise AssetsCouldNotFound("角色素材图标不存在", target)
        return data

    def gacha(self, target: StrOrInt, second_target: StrOrInt = None) -> Path:
        icon = self.get_target(target, second_target)
        return self.get_path(icon, "gacha")

    def icon(self, target: StrOrInt, second_target: StrOrInt = None) -> Path:
        icon = self.get_target(target, second_target)
        return self.get_path(icon, "icon")

    def square(self, target: StrOrInt, second_target: StrOrInt = None) -> Path:
        icon = self.get_target(target, second_target)
        return self.get_path(icon, "square")

    def normal(self, target: StrOrInt, second_target: StrOrInt = None) -> Path:
        icon = self.get_target(target, second_target)
        return self.get_path(icon, "normal")


class _WeaponAssets(_AssetsService):
    path: Path
    data: List[Weapon]
    name_map: Dict[str, Weapon]
    id_map: Dict[int, Weapon]

    def __init__(self, client: Optional[AsyncClient] = None) -> None:
        super().__init__(client)
        self.path = ASSETS_PATH.joinpath("engines")
        self.path.mkdir(exist_ok=True, parents=True)

    async def initialize(self):
        logger.info("正在初始化武器素材图标")
        html = await self.client.get(DATA_MAP["weapon"])
        self.data = [Weapon(**data) for data in html.json()]
        self.name_map = {icon.name: icon for icon in self.data}
        self.id_map = {icon.id: icon for icon in self.data}
        tasks = []
        for icon in self.data:
            base_path = self.path / f"{icon.id}"
            base_path.mkdir(exist_ok=True, parents=True)
            icon_path = base_path / "icon.webp"
            if not icon_path.exists():
                tasks.append(self._download(icon.icon, icon_path))
            if len(tasks) >= 100:
                await asyncio.gather(*tasks)
                tasks = []
        if tasks:
            await asyncio.gather(*tasks)
        logger.info("武器素材图标初始化完成")

    def get_path(self, icon: Weapon, name: str) -> Path:
        path = self.path / f"{icon.id}"
        path.mkdir(exist_ok=True, parents=True)
        return path / f"{name}.webp"

    def get_by_id(self, id_: int) -> Optional[Weapon]:
        return self.id_map.get(id_, None)

    def get_by_name(self, name: str) -> Optional[Weapon]:
        return self.name_map.get(name, None)

    def get_target(self, target: StrOrInt, second_target: StrOrInt = None) -> Optional[Weapon]:
        if isinstance(target, int):
            return self.get_by_id(target)
        elif isinstance(target, str):
            return self.get_by_name(target)
        if second_target:
            return self.get_target(second_target)
        raise AssetsCouldNotFound("武器素材图标不存在", target)

    def icon(self, target: StrOrInt, second_target: StrOrInt = None) -> Path:
        icon = self.get_target(target, second_target)
        return self.get_path(icon, "icon")


class _BuddyAssets(_AssetsService):
    path: Path
    data: List[Buddy]
    id_map: Dict[int, Buddy]
    name_map: Dict[str, Buddy]

    def __init__(self, client: Optional[AsyncClient] = None) -> None:
        super().__init__(client)
        self.path = ASSETS_PATH.joinpath("buddy")
        self.path.mkdir(exist_ok=True, parents=True)

    async def initialize(self):
        logger.info("正在初始化邦布素材图标")
        html = await self.client.get(DATA_MAP["buddy"])
        self.data = [Buddy(**data) for data in html.json()]
        self.id_map = {icon.id: icon for icon in self.data}
        self.name_map = {icon.name: icon for icon in self.data}
        tasks = []
        for icon in self.data:
            webp_path = self.path / f"{icon.id}.webp"
            png_path = self.path / f"{icon.id}.png"
            square_path = self.path / f"{icon.id}_square.png"
            if not webp_path.exists() and icon.webp:
                tasks.append(self._download(icon.webp, webp_path))
            if not png_path.exists() and icon.png:
                tasks.append(self._download(icon.png, png_path))
            if not square_path.exists() and icon.square:
                tasks.append(self._download(icon.square, square_path))
            if len(tasks) >= 100:
                await asyncio.gather(*tasks)
                tasks = []
        if tasks:
            await asyncio.gather(*tasks)
        logger.info("邦布素材图标初始化完成")

    def get_path(self, icon: Buddy, ext: str, square: bool = False) -> Path:
        square_str = "_square" if square else ""
        path = self.path / f"{icon.id}{square_str}.{ext}"
        return path

    def get_by_id(self, id_: int) -> Optional[Buddy]:
        return self.id_map.get(id_, None)

    def get_by_name(self, name: str) -> Optional[Buddy]:
        return self.name_map.get(name, None)

    def get_target(self, target: StrOrInt, second_target: StrOrInt = None) -> Optional[Buddy]:
        if isinstance(target, int):
            return self.get_by_id(target)
        elif isinstance(target, str):
            return self.get_by_name(target)
        if second_target:
            return self.get_target(second_target)
        raise AssetsCouldNotFound("邦布素材图标不存在", target)

    def webp(self, target: StrOrInt, second_target: StrOrInt = None) -> Path:
        icon = self.get_target(target, second_target)
        return self.get_path(icon, "webp")

    def png(self, target: StrOrInt, second_target: StrOrInt = None) -> Path:
        icon = self.get_target(target, second_target)
        return self.get_path(icon, "png")

    def icon(self, target: StrOrInt, second_target: StrOrInt = None) -> Path:
        icon = self.get_target(target, second_target)
        webp_path = self.get_path(icon, "webp")
        png_path = self.get_path(icon, "png")
        if webp_path.exists():
            return webp_path
        if png_path.exists():
            return png_path
        raise AssetsCouldNotFound("邦布素材图标不存在", target)

    def square(self, target: StrOrInt, second_target: StrOrInt = None) -> Path:
        icon = self.get_target(target, second_target)
        return self.get_path(icon, "png", square=True)


class _EquipmentSuitAssets(_AssetsService):
    path: Path
    data: List[EquipmentSuit]
    id_map: Dict[int, EquipmentSuit]
    name_map: Dict[str, EquipmentSuit]

    def __init__(self, client: Optional[AsyncClient] = None) -> None:
        super().__init__(client)
        self.path = ASSETS_PATH.joinpath("equipment_suit")
        self.path.mkdir(exist_ok=True, parents=True)

    async def initialize(self):
        logger.info("正在初始化驱动盘素材图标")
        html = await self.client.get(DATA_MAP["equipment_suit"])
        self.data = [EquipmentSuit(**data) for data in html.json()]
        self.id_map = {theme.id: theme for theme in self.data}
        self.name_map = {theme.name: theme for theme in self.data}
        tasks = []
        for theme in self.data:
            path = self.path / f"{theme.id}.webp"
            if not path.exists():
                tasks.append(self._download(theme.icon, path))
            if len(tasks) >= 100:
                await asyncio.gather(*tasks)
                tasks = []
        if tasks:
            await asyncio.gather(*tasks)
        logger.info("驱动盘素材图标初始化完成")

    def get_path(self, theme: EquipmentSuit, ext: str) -> Path:
        path = self.path / f"{theme.id}.{ext}"
        return path

    def get_by_id(self, id_: int) -> Optional[EquipmentSuit]:
        return self.id_map.get(id_, None)

    def get_by_name(self, name_: str) -> Optional[EquipmentSuit]:
        return self.name_map.get(name_, None)

    def get_target(self, target: StrOrInt, second_target: StrOrInt = None) -> Optional[EquipmentSuit]:
        if isinstance(target, int):
            return self.get_by_id(target)
        elif isinstance(target, str):
            return self.get_by_name(target)
        if second_target:
            return self.get_target(second_target)
        raise AssetsCouldNotFound("驱动盘素材图标不存在", target)

    def icon(self, target: StrOrInt, second_target: StrOrInt = None) -> Path:
        theme = self.get_target(target, second_target)
        webp_path = self.get_path(theme, "webp")
        if webp_path.exists():
            return webp_path
        raise AssetsCouldNotFound("驱动盘素材图标不存在", target)


class AssetsService(BaseService.Dependence):
    """asset服务

    用于储存和管理 asset :
        当对应的 asset (如某角色图标)不存在时，该服务会先查找本地。
        若本地不存在，则从网络上下载；若存在，则返回其路径
    """

    client: Optional[AsyncClient] = None

    avatar: _AvatarAssets
    """角色"""

    weapon: _WeaponAssets
    """武器"""

    buddy: _BuddyAssets
    """邦布"""

    equipment_suit: _EquipmentSuitAssets
    """驱动盘"""

    def __init__(self):
        self.client = AsyncClient(timeout=60.0)
        self.avatar = _AvatarAssets(self.client)
        self.weapon = _WeaponAssets(self.client)
        self.buddy = _BuddyAssets(self.client)
        self.equipment_suit = _EquipmentSuitAssets(self.client)

    async def initialize(self):  # pylint: disable=W0221
        await self.avatar.initialize()
        await self.weapon.initialize()
        await self.buddy.initialize()
        await self.equipment_suit.initialize()