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), "woman": ("woman", RANDOM_SUBCATEGORY), "man": ("man", RANDOM_SUBCATEGORY), "couple": ("couple", RANDOM_SUBCATEGORY), "group_or_layout": ("group_or_layout", 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