Support for motion type '0' (#4)

src/cysp2skel.py

@@ -7,26 +7,30 @@
 
 
 class Cysp2Skel:
-    _base: BufferedReader | None = None
-    _animation_count_index: int = 0
-    _base_animation_count: int = 0
+    _default_base: BufferedReader | None = None
+    _default_base_animation_count: int = 0
+    _default_animation_count_index: int = 0
+
     _unit_class_data: dict[str, str] | None = None
 
-    __slots__ = ("_cysp_dir", "mdb")
+    __slots__ = (
+        "_cysp_dir",
+        "mdb",
+        "_current_base",
+        "_current_base_animation_count",
+        "_current_animation_count_index",
+    )
 
     def __init__(self, cysp_dir: str | Path, mdb: MDB) -> None:
         self._cysp_dir = Path(cysp_dir)
         self.mdb = mdb
+        self._current_base: BufferedReader | None = None
+        self._current_base_animation_count: int = 0
+        self._current_animation_count_index: int = 0
 
-    @property
-    def base(self) -> BufferedReader:
-        if Cysp2Skel._base is not None:
-            Cysp2Skel._base.seek(0)
-            return Cysp2Skel._base
-
-        base_id = "000000"
+    def _get_base(self, str_id: str) -> RawIOBase:
         files: list[str] = [
-            f"{base_id}_{name}.cysp"
+            f"{str_id}_{name}.cysp"
             for name in [
                 "CHARA_BASE",
                 "DEAR",
@@ -40,48 +44,104 @@ def base(self) -> BufferedReader:
 
         paths: list[Path] = [self._cysp_dir / file for file in files]
 
+        self._current_base_animation_count = 0
         base = BytesIO()
         for file in paths:
-            with file.open("rb") as f:
-                f.seek(12)
-                count = _read_varint(f)
-                f.seek((count + 1) * 32)
-                base.write(f.read())
-                if file.name == "000000_CHARA_BASE.cysp":
-                    Cysp2Skel._animation_count_index = base.tell()
-                    base.write(b"\0")
-                else:
-                    Cysp2Skel._base_animation_count += count
-        base.seek(Cysp2Skel._animation_count_index)
-        base.write(bytes([Cysp2Skel._base_animation_count]))
+            try:
+                with file.open("rb") as f:
+                    f.seek(12)
+                    count = _read_varint(f)
+                    f.seek((count + 1) * 32)
+                    base.write(f.read())
+                    if "CHARA_BASE.cysp" in file.name:
+                        self._current_animation_count_index = base.tell()
+                        base.write(b"\0")
+                    else:
+                        self._current_base_animation_count += count
+            except FileNotFoundError as e:
+                print(f"Ignoring {e}")
+        base.seek(self._current_animation_count_index)
+        base.write(bytes([self._current_base_animation_count]))
         base.seek(0)
-        Cysp2Skel._base = BufferedReader(cast(RawIOBase, base))
-        return Cysp2Skel._base
+        return base
+
+    def get_default_base(self) -> BufferedReader:
+        if Cysp2Skel._default_base is not None:
+            Cysp2Skel._default_base.seek(0)
+            self._current_base_animation_count = Cysp2Skel._default_base_animation_count
+            self._current_animation_count_index = (
+                Cysp2Skel._default_animation_count_index
+            )
+            return Cysp2Skel._default_base
+
+        Cysp2Skel._default_base = BufferedReader(self._get_base("000000"))
+        Cysp2Skel._default_base_animation_count = self._current_base_animation_count
+        Cysp2Skel._default_animation_count_index = self._current_animation_count_index
+
+        return Cysp2Skel._default_base
+
+    def get_unit_data(self) -> dict[str | dict[str | str]]:
+        unit_data = self.mdb.c.execute(
+            "select unit_id,prefab_id,prefab_id_battle,motion_type from unit_data"
+        ).fetchall()
 
-    def get_unit_class_data(self) -> dict:
-        data = self.mdb.c.execute(
-            "select unit_id,motion_type from unit_data"
+        unit_enemy_data = self.mdb.c.execute(
+            "select unit_id,prefab_id,motion_type from unit_enemy_data"
         ).fetchall()
-        return {str(id_): str(type_) for id_, type_ in data}
+
+        unit_data_dict = {
+            str(unit_id): {
+                "motion_type": str(motion_type),
+                "prefab_id": str(prefab_id),
+                "prefab_id_battle": str(prefab_id_battle),
+            }
+            for unit_id, prefab_id, prefab_id_battle, motion_type in unit_data
+        }
+        enemy_data_dict = {
+            str(unit_id): {
+                "motion_type": str(motion_type),
+                "prefab_id": str(prefab_id),
+            }
+            for unit_id, prefab_id, motion_type in unit_enemy_data
+        }
+
+        return {**unit_data_dict, **enemy_data_dict}
 
     @property
-    def unit_class_data(self) -> dict[str, str]:
+    def unit_data(self) -> dict[str, str]:
         if Cysp2Skel._unit_class_data is None:
-            Cysp2Skel._unit_class_data = self.get_unit_class_data()
+            Cysp2Skel._unit_class_data = self.get_unit_data()
         return Cysp2Skel._unit_class_data
 
     def get_skeleton_buffer(self, unit_id: int) -> BufferedReader | None:
-        class_type = self.unit_class_data.get(str(unit_id // 100) + "01")
-        if class_type is None:
+        base_unit_id = str(unit_id // 100) + "01"
+
+        if (unit := self.unit_data.get(base_unit_id)) is None:
+            for unit in self.unit_data.values():
+                if unit.get("prefab_id") == base_unit_id:
+                    break
+                if unit.get("prefab_id_battle") == base_unit_id:
+                    break
+
+        motion_type: str = unit.get("motion_type")
+
+        if motion_type is None:
             return None
-        chara_class: str = class_type.rjust(2, "0")
+        if motion_type == "0":
+            motion_type = base_unit_id
+            self._current_base = self._get_base(base_unit_id)
+        else:
+            self._current_base = self.get_default_base()
+
+        chara_class: str = motion_type.rjust(2, "0")
         files: list[str] = []
         files.append(f"{chara_class}_COMMON_BATTLE.cysp")
-        files.append(f"{unit_id//100}01_BATTLE.cysp")
+        # files.append(f"{chara_class}_LOADING.cysp")
+        files.append(f"{unit.get('prefab_id')}_BATTLE.cysp")
         paths: list[Path] = [self._cysp_dir / file for file in files]
 
         skel = BytesIO()
-        skel.write(self.base.read())
+        skel.write(self._current_base.read())
         class_animation_count = 0
         for file in paths:
             with file.open("rb") as f:
@@ -90,8 +150,8 @@ def get_skeleton_buffer(self, unit_id: int) -> BufferedReader | None:
                 f.seek((count + 1) * 32)
                 skel.write(f.read())
                 class_animation_count += count
-        skel.seek(Cysp2Skel._animation_count_index)
-        anim_count = Cysp2Skel._base_animation_count + class_animation_count
+        skel.seek(self._current_animation_count_index)
+        anim_count = self._current_base_animation_count + class_animation_count
         skel.write(bytes([anim_count]))
         skel.seek(0)
         return BufferedReader(cast(RawIOBase, skel))