from database.mongodb import get_collection
from typing import Optional
from database.base_database_data import BaseDatabaseData
from game_server.resource import ResourceManager
from game_server.resource.configdb.avatar_config import AvatarConfig
from typing import List
from rail_proto.lib import (
    Avatar,
    AvatarSkillTree,
    BattleAvatar,
    AvatarType,
    SpBarInfo,
    AvatarSkillTree,
    BattleEquipment,
    EquipRelic
)

avatars_collection = get_collection("avatars")

class AvatarData(BaseDatabaseData):
    avatar_id: int
    level: int = 80
    exp: int = 0
    promotion: int = 6
    rank: int = 6
    lightcone_id: int = 0
    relic_ids: dict[str,int] = {}
    skills: dict[str,int] = {}
    current_hp: int = 10000
    current_sp: int = 0

    def add_avatar(self) -> "AvatarData":
        get_avatar = avatars_collection.find_one({"uid": self.uid, "avatar_id":self.avatar_id})
        if get_avatar:
            return False
        
        avatar = ResourceManager.instance().find_by_index(AvatarConfig, self.avatar_id)
        if not avatar:
            return False
        
        self.avatar_id=avatar.AvatarID

        for skill_id,skill_level in avatar.AvatarSkills.items():
            self.skills[str(skill_id)] = skill_level

        avatar_data = self.model_dump()
        avatars_collection.insert_one(avatar_data)
        return self
    
    def ToProto(self) -> Avatar:
        return Avatar(
            base_avatar_id=self.avatar_id,
            level=self.level,
            exp=self.exp,
            promotion=self.promotion,
            rank=self.rank,
            skilltree_list=[
                AvatarSkillTree(
                    point_id=int(skill_id),
                    level=skill_level
                )
                for skill_id,skill_level in self.skills.items()
            ],
            equipment_unique_id=self.lightcone_id,
            equip_relic_list=[
                EquipRelic(
                    relic_unique_id=relic_id,
                    type=int(type)
                )
                for type,relic_id in self.relic_ids.items()
            ],
            has_taken_promotion_reward_list=[1,2,3,4,5,6],
        )
    
    def ToBattleProto(self,index,player):
        get_equipment = None
        if self.lightcone_id > 0:
            get_equipment = player.inventory_manager.get(self.lightcone_id)

        return BattleAvatar(
            id=self.avatar_id,
            index=index,
            level=self.level,
            promotion=self.promotion,
            rank=self.rank,
            hp=self.current_hp,
            avatar_type=AvatarType.AVATAR_FORMAL_TYPE.value,
            sp_bar=SpBarInfo(
                cur_sp=5000, 
                max_sp=10000
            ),
            skilltree_list=[
                AvatarSkillTree(
                    point_id=int(skill_id),
                    level=level
                )
                for skill_id,level in self.skills.items()
            ],
            equipment_list=[
                BattleEquipment(
                    id=get_equipment.item_id if get_equipment else 0,
                    level=get_equipment.level if get_equipment else 0,
                    promotion=get_equipment.promotion if get_equipment else 0,
                    rank=get_equipment.rank if get_equipment else 0
                )
            ],
            relic_list = [
                player.inventory_manager.get(relic_id).RelicBattleProto()
                for type, relic_id in self.relic_ids.items()
                if relic_id > 0
            ]
        )
    
    def save_avatar(self):
        avatar_data = self.model_dump()
        avatar_data["uid"] = self.uid
        query = {"uid": self.uid, "avatar_id": self.avatar_id}
        avatars_collection.update_one(query, {"$set": avatar_data})


def find_avatar_by_avatar_id(uid: int, avatar_id: int) -> Optional[AvatarData]:
    avatar_data = avatars_collection.find_one({"uid": uid, "avatar_id":avatar_id})
    if avatar_data:
        del avatar_data["_id"]
        return AvatarData(**avatar_data)
    return None

def get_all_avatars_by_uid(uid: int) -> List[AvatarData]:
    avatars_data = avatars_collection.find({"uid": uid})
    result = []
    for avatar in avatars_data:
        del avatar["_id"]
        result.append(AvatarData(**avatar))
    return result