Added stage cutin

app/core/nice/quest.py

@@ -13,8 +13,11 @@
 from ...db.helpers.quest import get_questSelect_container
 from ...db.helpers.rayshift import (
     get_all_quest_hashes,
+    get_cutin_drops,
+    get_cutin_skills,
     get_rayshift_drops,
     get_war_board_quest_details,
+    quest_has_cutins,
 )
 from ...rayshift.quest import get_quest_detail
 from ...redis import Redis
@@ -47,6 +50,7 @@
     NiceQuestReleaseOverwrite,
     NiceRestriction,
     NiceStage,
+    NiceStageCutIn,
     NiceStageStartMovie,
     QuestEnemy,
     SupportServant,
@@ -78,6 +82,7 @@
 from .follower import get_nice_support_servants
 from .gift import get_nice_gift
 from .item import get_nice_item_amount, get_nice_item_from_raw
+from .stage_cutin import get_quest_stage_cutins
 
 
 settings = Settings()
@@ -169,6 +174,7 @@ def get_nice_stage(
     enemies: list[QuestEnemy],
     bgms: list[MstBgm],
     waveStartMovies: dict[int, list[NiceStageStartMovie]],
+    stage_cutins: dict[int, NiceStageCutIn],
     lang: Language,
 ) -> NiceStage:
     filtered_bgms = [bgm for bgm in bgms if bgm.id == raw_stage.bgmId]
@@ -190,6 +196,7 @@ def get_nice_stage(
         else None,
         NoEntryIds=raw_stage.script.get("NoEntryIds"),
         waveStartMovies=waveStartMovies.get(raw_stage.wave, []),
+        cutin=stage_cutins.get(raw_stage.wave, None),
         originalScript=raw_stage.script or {},
         enemies=enemies,
     )
@@ -509,6 +516,7 @@ def set_follower_data(followers: dict[int, QuestEnemy] | None) -> None:
     nice_quest_drops: list[EnemyDrop] = []
     quest_enemies = QuestEnemies(enemy_waves=[[]] * len(db_data.raw.mstStage))
     all_rayshift_hashes: list[str] = []
+    stage_cutins: dict[int, NiceStageCutIn] = {}
 
     if stages:
         rayshift_quest_id = quest_id
@@ -535,6 +543,36 @@ def set_follower_data(followers: dict[int, QuestEnemy] | None) -> None:
             elif region == Region.NA:
                 min_query_id = 1062363  # 2022-07-04 09:00:00 UTC
 
+        rayshift_query_questHash = (
+            None if db_data.raw.mstQuest.type == QuestType.WAR_BOARD else questHash
+        )
+
+        runs_with_cutin = await quest_has_cutins(
+            conn=conn,
+            quest_id=rayshift_quest_id,
+            phase=phase,
+            questSelect=questSelect,
+            questHash=rayshift_query_questHash,
+            min_query_id=min_query_id,
+        )
+
+        rayshift_kwargs = {
+            "conn": conn,
+            "quest_id": rayshift_quest_id,
+            "phase": phase,
+            "questSelect": questSelect,
+            "questHash": rayshift_query_questHash,
+            "min_query_id": min_query_id,
+        }
+        if runs_with_cutin:
+            cutin_skills, cutin_drops = await asyncio.gather(
+                get_cutin_skills(**rayshift_kwargs),  # type:ignore
+                get_cutin_drops(**rayshift_kwargs, runs=runs_with_cutin),  # type:ignore
+            )
+            stage_cutins = await get_quest_stage_cutins(
+                conn, region, runs_with_cutin, cutin_drops, cutin_skills, lang
+            )
+
         if db_data.raw.mstQuest.type == QuestType.WAR_BOARD:
             quest_enemy_coro = get_war_board_quest_details(conn, quest_id, phase)
         else:
@@ -554,14 +592,7 @@ def set_follower_data(followers: dict[int, QuestEnemy] | None) -> None:
             all_rayshift_hashes,
         ) = await asyncio.gather(
             quest_enemy_coro,
-            get_rayshift_drops(
-                conn,
-                rayshift_quest_id,
-                phase,
-                questSelect,
-                None if db_data.raw.mstQuest.type == QuestType.WAR_BOARD else questHash,
-                min_query_id,
-            ),
+            get_rayshift_drops(**rayshift_kwargs),  # type:ignore
             get_all_quest_hashes(conn, rayshift_quest_id, phase, questSelect),
         )
 
@@ -625,7 +656,13 @@ def set_follower_data(followers: dict[int, QuestEnemy] | None) -> None:
 
     new_nice_stages = [
         get_nice_stage(
-            region, stage, enemies, db_data.raw.mstBgm, waveStartMovies, lang
+            region,
+            stage,
+            enemies,
+            db_data.raw.mstBgm,
+            waveStartMovies,
+            stage_cutins,
+            lang,
         )
         for stage, enemies in zip(stages, quest_enemies.enemy_waves, strict=False)
     ]

app/core/nice/stage_cutin.py

@@ -0,0 +1,46 @@
+from sqlalchemy.ext.asyncio import AsyncConnection
+
+from ...schemas.common import Language, Region
+from ...schemas.nice import NiceStageCutIn, NiceStageCutInSkill
+from ...schemas.rayshift import CutInSkill, QuestDrop
+from .enemy import get_nice_drop
+from .skill import MultipleNiceSkills, SkillSvt, get_multiple_nice_skills
+
+
+def get_nice_stage_cut_in(
+    runs: int,
+    stage: int,
+    cutin_drops: list[QuestDrop],
+    cutin_skills: list[CutInSkill],
+    all_skills: MultipleNiceSkills,
+) -> NiceStageCutIn:
+    skills = [
+        NiceStageCutInSkill(
+            skill=all_skills[SkillSvt(skill.skill_id, 0)],
+            appearCount=skill.appear_count,
+        )
+        for skill in cutin_skills
+        if skill.stage == stage
+    ]
+
+    drops = [get_nice_drop(drop) for drop in cutin_drops if drop.stage == stage]
+
+    return NiceStageCutIn(runs=runs, skills=skills, drops=drops)
+
+
+async def get_quest_stage_cutins(
+    conn: AsyncConnection,
+    region: Region,
+    runs: int,
+    cutin_drops: list[QuestDrop],
+    cutin_skills: list[CutInSkill],
+    lang: Language = Language.jp,
+) -> dict[int, NiceStageCutIn]:
+    all_skill_ids = [SkillSvt(skill.skill_id, 0) for skill in cutin_skills]
+    all_skills = await get_multiple_nice_skills(conn, region, all_skill_ids, lang)
+
+    return {
+        stage: get_nice_stage_cut_in(runs, stage, cutin_drops, cutin_skills, all_skills)
+        for stage in {drop.stage for drop in cutin_drops}
+        | {skill.stage for skill in cutin_skills}
+    }