neonsr/database/inventory/inventory_data.py

from database.mongodb import get_collection
from database.base_database_data import BaseDatabaseData
from game_server.resource import ResourceManager
from game_server.resource.configdb.relic_main_affix_config import RelicMainAffixConfigData
from game_server.resource.configdb.equipment_config import EquipmentConfig
from game_server.resource.configdb.relic_config import RelicConfigData
from game_server.game.items.relic_manager import RelicManager
from game_server.game.enums.avatar.relic_type import RelicTypeEnum
from typing import List
from pydantic import BaseModel
from rail_proto.lib import (
    Equipment,
    Relic,
    BattleRelic,
    RelicAffix,
    PlayerSyncScNotify
)

items_collection = get_collection("items")

class SubAffix(BaseModel):
    id: int
    count: int
    step: int

class InventoryData(BaseDatabaseData):
    item_id: int
    count: int = 1
    level: int = 1
    exp: int = 0
    promotion: int = 0
    rank: int = 0
    locked: bool = False
    discarded: bool = False
    main_affix: int = 0
    sub_affix: list[SubAffix] = []
    equip_avatar: int = 0
    unique_id: int = 0
    main_type: int = 0
    sub_type: int = 0

    def add_lightcone(self):
        last_item = items_collection.find_one({"uid": self.uid},sort=[("unique_id", -1)])
        unique_id = last_item["unique_id"] + 1 if last_item else 1
        equipment = ResourceManager.instance().find_by_index(EquipmentConfig, self.item_id)
        if not equipment:
            return None
        
        self.level = 80
        self.promotion = 6
        self.main_type = 1
        self.unique_id = unique_id

        lightcone_data = self.model_dump()
        items_collection.insert_one(lightcone_data)
        return self

    def add_relic(self):
        last_item = items_collection.find_one({"uid": self.uid}, sort=[("unique_id", -1)])
        unique_id = last_item["unique_id"] + 1 if last_item else 1
        
        relic = ResourceManager.instance().find_by_index(RelicConfigData, self.item_id)
        if not relic:
            return None

        self.level = 15
        self.main_type = 2
        self.sub_type = RelicTypeEnum[relic.Type].value

        relic_manager = RelicManager()
        main_stat = ""
        get_main_affix = None

        if self.main_affix > 0:
            relics = ResourceManager.instance().find_by_index(RelicConfigData, self.item_id)
            main_affix_group = ResourceManager.instance().find_all_by_index(RelicMainAffixConfigData, relics.MainAffixGroup)
            main_stat = next(
                (affix.Property for affix in main_affix_group if affix.AffixID == int(self.main_affix))
            )
        elif self.main_affix == 0:
            get_main_affix = relic_manager.GetRandomRelicMainAffix(self.item_id)
            self.main_affix = get_main_affix.AffixID

        if not main_stat and get_main_affix:
            main_stat = get_main_affix.Property

        if not self.sub_affix:
            self.sub_affix = relic_manager.GetRandomRelicSubAffix(self.item_id, main_stat, 4)

        if self.sub_affix and len(self.sub_affix) < 4:
            self.sub_affix.extend(relic_manager.GetRandomRelicSubAffix(self.item_id, main_stat, 4 - len(self.sub_affix), self.sub_affix))

        self.unique_id = unique_id

        relic_data = self.model_dump()
        items_collection.insert_one(relic_data)

        return self

    def ToProto(self):
        if self.main_type == 1:
            proto = Equipment(
                tid=self.item_id,
                level=self.level,
                rank=self.rank,
                promotion=self.promotion,
                dress_avatar_id=self.equip_avatar,
                unique_id=self.unique_id,
            )
        if self.main_type == 2:
            proto = Relic(
                tid=self.item_id,
                level=self.level,
                dress_avatar_id=self.equip_avatar,
                unique_id=self.unique_id,
                main_affix_id=self.main_affix,
                sub_affix_list=[
                    RelicAffix(
                        affix_id=affix.id,
                        cnt=affix.count,
                        step=affix.step
                    )
                    for affix in self.sub_affix
                ]
            )
        return proto

    def RelicSyncProto(self):
        proto = PlayerSyncScNotify(
            relic_list=[self.ToProto()]
        )
        return proto

    def LightconeSyncProto(self):
        proto = PlayerSyncScNotify(
            equipment_list=[self.ToProto()]
        )
        return proto
    
    def save_item(self):
        item_data = self.model_dump()
        item_data["uid"] = self.uid
        query = {"uid": self.uid, "unique_id": self.unique_id}
        items_collection.update_one(query, {"$set": item_data})

    def RelicBattleProto(self):
        return BattleRelic(
            id=self.item_id,
            level=self.level,
            main_affix_id=self.main_affix,
            sub_affix_list=[
                RelicAffix(
                    affix_id=affix.id,
                    cnt=affix.count,
                    step=affix.step
                )
                for affix in self.sub_affix
            ],
            unique_id=self.unique_id,
            type=self.sub_type
        )
        


def get_all_items_by_uid(uid: int) -> List[InventoryData]:
    items_data = items_collection.find({"uid": uid})
    result = []
    for items in items_data:
        del items["_id"]
        result.append(InventoryData(**items))
    return result