from __future__ import annotations

import json

try:
    from .hardcore_position_config import (
        build_hardcore_action_filter_json,
        build_hardcore_position_pool_json,
        hardcore_position_family_choices,
        hardcore_position_focus_choices,
        hardcore_position_key_choices,
    )
except ImportError:  # Allows local smoke tests from the repository root.
    from hardcore_position_config import (
        build_hardcore_action_filter_json,
        build_hardcore_position_pool_json,
        hardcore_position_family_choices,
        hardcore_position_focus_choices,
        hardcore_position_key_choices,
    )


SXCP_HARDCORE_POSITION_CONFIG = "SXCP_HARDCORE_POSITION_CONFIG"


def _choice_input_key(prefix, choice):
    key = "".join(char if char.isalnum() else "_" for char in str(choice).lower()).strip("_")
    while "__" in key:
        key = key.replace("__", "_")
    return f"{prefix}_{key}"


class SxCPHardcorePositionPool:
    @classmethod
    def INPUT_TYPES(cls):
        required = {
            "combine_mode": (["replace", "add"], {"default": "replace"}),
            "family": (hardcore_position_family_choices(), {"default": "any"}),
        }
        for choice in hardcore_position_key_choices():
            required[_choice_input_key("include", choice)] = ("BOOLEAN", {"default": False})
        return {
            "required": required,
            "optional": {
                "hardcore_position_config": (SXCP_HARDCORE_POSITION_CONFIG,),
            },
        }

    RETURN_TYPES = (SXCP_HARDCORE_POSITION_CONFIG, "STRING")
    RETURN_NAMES = ("hardcore_position_config", "summary")
    FUNCTION = "build"
    CATEGORY = "prompt_builder"

    def build(self, combine_mode="replace", family="any", hardcore_position_config="", **kwargs):
        selected = [
            choice
            for choice in hardcore_position_key_choices()
            if bool(kwargs.get(_choice_input_key("include", choice), False))
        ]
        config = build_hardcore_position_pool_json(
            hardcore_position_config=hardcore_position_config or "",
            combine_mode=combine_mode,
            family=family,
            selected_positions=selected,
        )
        return config, json.loads(config).get("summary", "")


class SxCPHardcoreActionFilter:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "focus": (hardcore_position_focus_choices(), {"default": "keep_pool"}),
                "allow_toys": ("BOOLEAN", {"default": False}),
                "allow_double": ("BOOLEAN", {"default": False}),
                "allow_penetration": ("BOOLEAN", {"default": True}),
                "allow_foreplay": ("BOOLEAN", {"default": True}),
                "allow_interaction": ("BOOLEAN", {"default": True}),
                "allow_manual": ("BOOLEAN", {"default": True}),
                "allow_oral": ("BOOLEAN", {"default": True}),
                "allow_outercourse": ("BOOLEAN", {"default": True}),
                "allow_anal": ("BOOLEAN", {"default": True}),
                "allow_climax": ("BOOLEAN", {"default": True}),
            },
            "optional": {
                "hardcore_position_config": (SXCP_HARDCORE_POSITION_CONFIG,),
            },
        }

    RETURN_TYPES = (SXCP_HARDCORE_POSITION_CONFIG, "STRING")
    RETURN_NAMES = ("hardcore_position_config", "summary")
    FUNCTION = "build"
    CATEGORY = "prompt_builder"

    def build(
        self,
        focus,
        allow_toys,
        allow_double,
        allow_penetration,
        allow_foreplay,
        allow_interaction,
        allow_manual,
        allow_oral,
        allow_outercourse,
        allow_anal,
        allow_climax,
        hardcore_position_config="",
    ):
        config = build_hardcore_action_filter_json(
            hardcore_position_config=hardcore_position_config or "",
            focus=focus,
            allow_toys=allow_toys,
            allow_double=allow_double,
            allow_penetration=allow_penetration,
            allow_foreplay=allow_foreplay,
            allow_interaction=allow_interaction,
            allow_manual=allow_manual,
            allow_oral=allow_oral,
            allow_outercourse=allow_outercourse,
            allow_anal=allow_anal,
            allow_climax=allow_climax,
        )
        return config, json.loads(config).get("summary", "")


NODE_CLASS_MAPPINGS = {
    "SxCPHardcorePositionPool": SxCPHardcorePositionPool,
    "SxCPHardcoreActionFilter": SxCPHardcoreActionFilter,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "SxCPHardcorePositionPool": "SxCP Hardcore Position Pool",
    "SxCPHardcoreActionFilter": "SxCP Hardcore Action Filter",
}