from database.mongodb import get_collection
from database.base_database_data import BaseDatabaseData
from typing import List
from rail_proto.lib import (
    LineupInfo,
    ExtraLineupType,
    LineupAvatar,
    SpBarInfo,
    AvatarType
)

lineups_collection = get_collection("lineups")

class LineupData(BaseDatabaseData):
    name: str = ""
    index: int = 0
    mp: int = 5
    extra_lineup: ExtraLineupType = ExtraLineupType.LINEUP_NONE
    leader_slot: int = 0
    avatar_list: list[int] = []
    
    def ToProto(self) -> LineupInfo:
        return LineupInfo(
            name=self.name,
            index=self.index,
            mp=self.mp,
            max_mp=5,
            extra_lineup_type=self.extra_lineup,
            leader_slot=self.leader_slot,
            avatar_list=[
                LineupAvatar(
                    id=avatar_id,
                    slot=index,
                    hp=10000,
                    sp_bar=SpBarInfo(cur_sp=5000, max_sp=10000),
                    satiety=100,
                    avatar_type=AvatarType.AVATAR_FORMAL_TYPE.value
                )
                for index, avatar_id in enumerate(self.avatar_list)
            ]
        )
    def add_lineup(self):
        lineup_data = self.model_dump()
        lineups_collection.insert_one(lineup_data)
        return self

    def save_lineup(self):
        lineup_data = self.model_dump()
        lineup_data["uid"] = self.uid
        query = {"uid": self.uid, "index": self.index}
        existing_lineup = lineups_collection.find_one(query)

        if existing_lineup:
            lineups_collection.update_one(query, {"$set": lineup_data})
        else:
            lineups_collection.insert_one(lineup_data)

        return True


def get_all_lineup_by_uid(uid: int) -> List[LineupData]:
    lineups_data = lineups_collection.find({"uid": uid})
    result = []
    for lineup in lineups_data:
        del lineup["_id"]
        result.append(LineupData(**lineup))
    return result