from pydantic import BaseModel
from database.player.player_data import PlayerData,players_collection
from database.avatar.avatar_data import AvatarData,get_all_avatars_by_uid,avatars_collection
from database.lineup.lineup_data import LineupData,get_all_lineup_by_uid,lineups_collection
from database.inventory.inventory_data import InventoryData,get_all_items_by_uid,items_collection
from game_server.game.scene.scene_manager import SceneManager
from typing import Optional
from rail_proto.lib import (
    PlayerSyncScNotify,
    AvatarSync
)
from pymongo import UpdateOne
from utils.logger import Error


class PlayerManager(BaseModel):
    data : PlayerData = PlayerData()
    avatar_mananger: dict[int,AvatarData] = {}
    lineup_manager: dict[int,LineupData] = {}
    inventory_manager: dict[int,InventoryData] = {}
    scene_manager: SceneManager = None

    def init_default(self):
        self.add_all_avatars()
        self.add_all_lineups()
        self.add_all_items()
        self.scene_manager = SceneManager(
            entry_id=self.data.entry_id
        )

    def save_all(self):
        self.save_player_data_bulk()
        self.save_all_avatars_bulk()
        self.save_all_lineups_bulk()
        self.save_all_items_bulk()

    def add_all_avatars(self):
        avatars = get_all_avatars_by_uid(uid=self.data.uid)
        if not avatars:
            avatar = AvatarData(
                uid=self.data.uid,
                avatar_id=self.data.cur_basic_type
            ).add_avatar()
            self.avatar_mananger[self.data.cur_basic_type] = avatar
            
        for avatar in avatars:
            self.avatar_mananger[avatar.avatar_id] = avatar

    def add_all_lineups(self):
        lineups = get_all_lineup_by_uid(uid=self.data.uid)
        if not lineups:
            lineup = LineupData(
                uid=self.data.uid,
                index=0,
                name="",
                avatar_list=[self.data.cur_basic_type]
            ).add_lineup()
            self.lineup_manager[lineup.index] = lineup
        for lineup in lineups:
            self.lineup_manager[lineup.index] = lineup

    def add_all_items(self):
        items = get_all_items_by_uid(uid=self.data.uid)
        for item in items:
            self.inventory_manager[item.unique_id] = item

    def add_avatar(self,avatar_id:int) -> Optional[AvatarData]:
        avatar = self.avatar_mananger.get(avatar_id)
        if avatar:
            return None
        
        avatar = AvatarData(uid=self.data.uid,avatar_id=avatar_id).add_avatar()
        if avatar:
            return avatar
        return None

    def add_lineup(self,avatar_ids: list[int],index=0,name=""):
        LineupData(
            uid=self.data.uid,
            index=index,
            name=name,
            avatar_list=avatar_ids
        ).add_lineup()

    def add_lightcone(self,lightcone_id: int, rank=1) -> Optional[InventoryData]:
        item = InventoryData(
            uid=self.data.uid,
            item_id=lightcone_id,
            rank=rank
        ).add_lightcone()
        if item:
            return item
        return None

    def add_relic(self,relic_id: int, main_affix=0, sub_affix=[]) -> Optional[InventoryData]:
        item = InventoryData(
            uid=self.data.uid,
            item_id=relic_id,
            main_affix=main_affix,
            sub_affix=sub_affix
        ).add_relic()
        if item:
            return item
        return None


    def PlayerSyncProto(self) -> PlayerSyncScNotify:
        avatars = []
        for avatar_id,avatar in self.avatar_mananger.items():
            avatars.append(avatar.ToProto())

        lightcones = []
        for unique_id,item in self.inventory_manager.items():
            if item.main_type == 1:
                lightcones.append(item.ToProto())

        relics = []
        for unique_id,item in self.inventory_manager.items():
            if item.main_type == 2:
                relics.append(item.ToProto())
        return PlayerSyncScNotify(
                equipment_list=lightcones,
                relic_list=relics,
                avatar_sync=
                    AvatarSync(
                        avatar_list=avatars
                    ),
                basic_info=self.data.ToProto()
            )

    def save_player_data_bulk(self):
        try:
            player_data = self.data.model_dump()
            player_data["uid"] = self.data.uid
            query = {"uid": self.data.uid}
            players_collection.update_one(query, {"$set": player_data}, upsert=True)
        except Exception as e:
            Error(f"Error save player data: {e}")

    def save_all_avatars_bulk(self):
        operations = []
        for avatar in self.avatar_mananger.values():
            avatar_data = avatar.model_dump()
            avatar_data["uid"] = avatar.uid
            query = {"uid": avatar.uid, "avatar_id": avatar.avatar_id}
            operations.append(UpdateOne(query, {"$set": avatar_data}, upsert=True))

        if operations:
            avatars_collection.bulk_write(operations)

    def save_all_lineups_bulk(self):
        operations = []
        for lineup in self.lineup_manager.values():
            lineup_data = lineup.model_dump()
            lineup_data["uid"] = lineup.uid
            query = {"uid": lineup.uid, "index": lineup.index}
            operations.append(UpdateOne(query, {"$set": lineup_data}, upsert=True))

        if operations:
            lineups_collection.bulk_write(operations)

    def save_all_items_bulk(self):
        operations = []
        for item in self.inventory_manager.values():
            item_data = item.model_dump()
            item_data["uid"] = item.uid
            query = {"uid": item.uid, "unique_id": item.unique_id}
            operations.append(UpdateOne(query, {"$set": item_data}, upsert=True))

        if operations:
            items_collection.bulk_write(operations)