Extract category cast config policy

2026-06-27 00:22:17 +02:00
parent f3f9929df5
commit 4c31553409
6 changed files with 159 additions and 80 deletions
@@ -0,0 +1,114 @@
+from __future__ import annotations
+
+import json
+from typing import Any
+
+
+RANDOM_SUBCATEGORY = "random"
+
+CATEGORY_PRESETS = {
+    "auto_weighted": ("auto_weighted", RANDOM_SUBCATEGORY),
+    "auto_full": ("auto_full", RANDOM_SUBCATEGORY),
+    "women_casual": ("Casual clothes", RANDOM_SUBCATEGORY),
+    "men_casual": ("Men casual clothes", RANDOM_SUBCATEGORY),
+    "couple_casual": ("Couple casual clothes", RANDOM_SUBCATEGORY),
+    "provocative_erotic": ("Provocative erotic clothes", RANDOM_SUBCATEGORY),
+    "hardcore_pose": ("Hardcore sexual poses", RANDOM_SUBCATEGORY),
+    "custom_random": ("custom_random", RANDOM_SUBCATEGORY),
+}
+
+CAST_PRESETS = {
+    "solo_woman": (1, 0),
+    "solo_man": (0, 1),
+    "mixed_couple": (1, 1),
+    "two_women": (2, 0),
+    "two_men": (0, 2),
+    "threesome_2w1m": (2, 1),
+    "small_group_3w2m": (3, 2),
+}
+
+
+def category_preset_choices() -> list[str]:
+    return list(CATEGORY_PRESETS)
+
+
+def cast_preset_choices() -> list[str]:
+    return list(CAST_PRESETS) + ["custom_counts"]
+
+
+def build_category_config_json(preset: str = "auto_weighted", subcategory: str = RANDOM_SUBCATEGORY) -> str:
+    category, default_subcategory = CATEGORY_PRESETS.get(preset, CATEGORY_PRESETS["auto_weighted"])
+    chosen_subcategory = subcategory if subcategory and subcategory != RANDOM_SUBCATEGORY else default_subcategory
+    return json.dumps(
+        {
+            "preset": preset if preset in CATEGORY_PRESETS else "auto_weighted",
+            "category": category,
+            "subcategory": chosen_subcategory,
+        },
+        ensure_ascii=True,
+        sort_keys=True,
+    )
+
+
+def parse_category_config(category_config: str | dict[str, Any] | None) -> tuple[str, str]:
+    if not category_config:
+        return CATEGORY_PRESETS["auto_weighted"]
+    if isinstance(category_config, dict):
+        raw = category_config
+    else:
+        try:
+            raw = json.loads(str(category_config))
+        except json.JSONDecodeError as exc:
+            raise ValueError(f"Invalid category_config JSON: {exc}") from exc
+    if not isinstance(raw, dict):
+        raise ValueError("category_config must be a JSON object")
+    preset = str(raw.get("preset") or "auto_weighted")
+    category, subcategory = CATEGORY_PRESETS.get(preset, CATEGORY_PRESETS["auto_weighted"])
+    category = str(raw.get("category") or category)
+    subcategory = str(raw.get("subcategory") or subcategory or RANDOM_SUBCATEGORY)
+    return category, subcategory
+
+
+def build_cast_config_json(cast_mode: str = "mixed_couple", women_count: int = 1, men_count: int = 1) -> str:
+    if cast_mode in CAST_PRESETS:
+        women_count, men_count = CAST_PRESETS[cast_mode]
+    else:
+        women_count = max(0, min(12, int(women_count)))
+        men_count = max(0, min(12, int(men_count)))
+        if women_count + men_count == 0:
+            women_count = 1
+        cast_mode = "custom_counts"
+    return json.dumps(
+        {
+            "cast_mode": cast_mode,
+            "women_count": int(women_count),
+            "men_count": int(men_count),
+        },
+        ensure_ascii=True,
+        sort_keys=True,
+    )
+
+
+def parse_cast_config(cast_config: str | dict[str, Any] | None) -> dict[str, int | str]:
+    if not cast_config:
+        return {"cast_mode": "mixed_couple", "women_count": 1, "men_count": 1}
+    if isinstance(cast_config, dict):
+        raw = cast_config
+    else:
+        try:
+            raw = json.loads(str(cast_config))
+        except json.JSONDecodeError as exc:
+            raise ValueError(f"Invalid cast_config JSON: {exc}") from exc
+    if not isinstance(raw, dict):
+        raise ValueError("cast_config must be a JSON object")
+    return json.loads(
+        build_cast_config_json(
+            str(raw.get("cast_mode") or "custom_counts"),
+            raw.get("women_count", 1),
+            raw.get("men_count", 1),
+        )
+    )
+
+
+_parse_category_config = parse_category_config
+_parse_cast_config = parse_cast_config