from __future__ import annotations

import json

try:
    from .prompt_builder import (
        build_prompt,
        build_prompt_from_configs,
        category_choices,
        ethnicity_choices,
        subcategory_choices,
    )
except ImportError:  # Allows local smoke tests from the repository root.
    from prompt_builder import (
        build_prompt,
        build_prompt_from_configs,
        category_choices,
        ethnicity_choices,
        subcategory_choices,
    )


SXCP_ETHNICITY_LIST = "SXCP_ETHNICITY_LIST"
SXCP_FILTER_CONFIG = "SXCP_FILTER_CONFIG"
SXCP_SEED_CONFIG = "SXCP_SEED_CONFIG"
SXCP_CAMERA_CONFIG = "SXCP_CAMERA_CONFIG"
SXCP_LOCATION_CONFIG = "SXCP_LOCATION_CONFIG"
SXCP_COMPOSITION_CONFIG = "SXCP_COMPOSITION_CONFIG"
SXCP_CATEGORY_CONFIG = "SXCP_CATEGORY_CONFIG"
SXCP_CAST_CONFIG = "SXCP_CAST_CONFIG"
SXCP_GENERATION_PROFILE = "SXCP_GENERATION_PROFILE"
SXCP_HARDCORE_POSITION_CONFIG = "SXCP_HARDCORE_POSITION_CONFIG"
SXCP_CHARACTER_CAST = "SXCP_CHARACTER_CAST"
SXCP_CHARACTER_PROFILE = "SXCP_CHARACTER_PROFILE"


class SxCPPromptBuilder:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "category": (category_choices(), {"default": "auto_weighted"}),
                "subcategory": (subcategory_choices(), {"default": "random"}),
                "row_number": ("INT", {"default": 1, "min": 1, "max": 1000000, "step": 1}),
                "start_index": ("INT", {"default": 41, "min": 1, "max": 1000000, "step": 1}),
                "seed": ("INT", {"default": 20260614, "min": 0, "max": 0xFFFFFFFF, "step": 1}),
                "clothing": (["random", "full", "minimal"], {"default": "random"}),
                "ethnicity": (ethnicity_choices(), {"default": "any"}),
                "poses": (["random", "standard", "evocative"], {"default": "random"}),
                "expression_enabled": ("BOOLEAN", {"default": True}),
                "expression_intensity": ("FLOAT", {"default": -1.0, "min": -1.0, "max": 1.0, "step": 0.01}),
                "backside_bias": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}),
                "figure": (["random", "curvy", "balanced", "bombshell"], {"default": "random"}),
                "women_count": ("INT", {"default": 1, "min": 0, "max": 12, "step": 1}),
                "men_count": ("INT", {"default": 1, "min": 0, "max": 12, "step": 1}),
                "minimal_clothing_ratio": ("FLOAT", {"default": -1.0, "min": -1.0, "max": 1.0, "step": 0.01}),
                "standard_pose_ratio": ("FLOAT", {"default": -1.0, "min": -1.0, "max": 1.0, "step": 0.01}),
                "trigger": ("STRING", {"default": "sxcpinup_coloredpencil"}),
                "prepend_trigger_to_prompt": ("BOOLEAN", {"default": True}),
            },
            "optional": {
                "ethnicity_list": (SXCP_ETHNICITY_LIST,),
                "seed_config": (SXCP_SEED_CONFIG,),
                "camera_config": (SXCP_CAMERA_CONFIG,),
                "location_config": (SXCP_LOCATION_CONFIG,),
                "composition_config": (SXCP_COMPOSITION_CONFIG,),
                "character_profile": (SXCP_CHARACTER_PROFILE,),
                "character_cast": (SXCP_CHARACTER_CAST,),
                "hardcore_position_config": (SXCP_HARDCORE_POSITION_CONFIG,),
                "extra_positive": ("STRING", {"default": "", "multiline": True}),
                "extra_negative": ("STRING", {"default": "", "multiline": True}),
            },
        }

    RETURN_TYPES = ("STRING", "STRING", "STRING", "STRING", "STRING", "STRING")
    RETURN_NAMES = ("prompt", "negative_prompt", "caption", "metadata_json", "category", "subcategory")
    FUNCTION = "build"
    CATEGORY = "prompt_builder"

    def build(
        self,
        category,
        subcategory,
        row_number,
        start_index,
        seed,
        clothing,
        ethnicity,
        poses,
        expression_enabled,
        expression_intensity,
        backside_bias,
        figure,
        women_count,
        men_count,
        minimal_clothing_ratio,
        standard_pose_ratio,
        trigger,
        prepend_trigger_to_prompt,
        seed_config="",
        camera_config="",
        location_config="",
        composition_config="",
        character_profile="",
        character_cast="",
        hardcore_position_config="",
        extra_positive="",
        extra_negative="",
        no_plus_women=False,
        no_black=False,
        ethnicity_list="",
    ):
        row = build_prompt(
            category=category,
            subcategory=subcategory,
            row_number=row_number,
            start_index=start_index,
            seed=seed,
            clothing=clothing,
            ethnicity=ethnicity_list or ethnicity,
            poses=poses,
            expression_enabled=expression_enabled,
            expression_intensity=expression_intensity,
            backside_bias=backside_bias,
            figure=figure,
            no_plus_women=no_plus_women,
            no_black=no_black,
            women_count=women_count,
            men_count=men_count,
            minimal_clothing_ratio=minimal_clothing_ratio,
            standard_pose_ratio=standard_pose_ratio,
            trigger=trigger,
            prepend_trigger_to_prompt=prepend_trigger_to_prompt,
            extra_positive=extra_positive or "",
            extra_negative=extra_negative or "",
            seed_config=seed_config or "",
            camera_config=camera_config or "",
            location_config=location_config or "",
            composition_config=composition_config or "",
            character_profile=character_profile or "",
            character_cast=character_cast or "",
            hardcore_position_config=hardcore_position_config or "",
        )
        return (
            row["prompt"],
            row["negative_prompt"],
            row["caption"],
            json.dumps(row, ensure_ascii=True, sort_keys=True),
            row.get("main_category", category),
            row.get("subcategory", subcategory),
        )


class SxCPPromptBuilderFromConfigs:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "row_number": ("INT", {"default": 1, "min": 1, "max": 1000000, "step": 1}),
                "start_index": ("INT", {"default": 41, "min": 1, "max": 1000000, "step": 1}),
                "seed": ("INT", {"default": 20260614, "min": 0, "max": 0xFFFFFFFF, "step": 1}),
            },
            "optional": {
                "category_config": (SXCP_CATEGORY_CONFIG,),
                "cast_config": (SXCP_CAST_CONFIG,),
                "generation_profile": (SXCP_GENERATION_PROFILE,),
                "filter_config": (SXCP_FILTER_CONFIG,),
                "ethnicity_list": (SXCP_ETHNICITY_LIST,),
                "seed_config": (SXCP_SEED_CONFIG,),
                "camera_config": (SXCP_CAMERA_CONFIG,),
                "location_config": (SXCP_LOCATION_CONFIG,),
                "composition_config": (SXCP_COMPOSITION_CONFIG,),
                "character_profile": (SXCP_CHARACTER_PROFILE,),
                "character_cast": (SXCP_CHARACTER_CAST,),
                "hardcore_position_config": (SXCP_HARDCORE_POSITION_CONFIG,),
                "extra_positive": ("STRING", {"default": "", "multiline": True}),
                "extra_negative": ("STRING", {"default": "", "multiline": True}),
            },
        }

    RETURN_TYPES = ("STRING", "STRING", "STRING", "STRING", "STRING", "STRING")
    RETURN_NAMES = ("prompt", "negative_prompt", "caption", "metadata_json", "category", "subcategory")
    FUNCTION = "build"
    CATEGORY = "prompt_builder"

    def build(
        self,
        row_number,
        start_index,
        seed,
        category_config="",
        cast_config="",
        generation_profile="",
        filter_config="",
        ethnicity_list="",
        seed_config="",
        camera_config="",
        location_config="",
        composition_config="",
        character_profile="",
        character_cast="",
        hardcore_position_config="",
        extra_positive="",
        extra_negative="",
    ):
        row = build_prompt_from_configs(
            row_number=row_number,
            start_index=start_index,
            seed=seed,
            category_config=category_config or "",
            cast_config=cast_config or "",
            generation_profile=generation_profile or "",
            filter_config=ethnicity_list or filter_config or "",
            seed_config=seed_config or "",
            camera_config=camera_config or "",
            location_config=location_config or "",
            composition_config=composition_config or "",
            character_profile=character_profile or "",
            character_cast=character_cast or "",
            hardcore_position_config=hardcore_position_config or "",
            extra_positive=extra_positive or "",
            extra_negative=extra_negative or "",
        )
        return (
            row["prompt"],
            row["negative_prompt"],
            row["caption"],
            json.dumps(row, ensure_ascii=True, sort_keys=True),
            row.get("main_category", ""),
            row.get("subcategory", ""),
        )


NODE_CLASS_MAPPINGS = {
    "SxCPPromptBuilder": SxCPPromptBuilder,
    "SxCPPromptBuilderFromConfigs": SxCPPromptBuilderFromConfigs,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "SxCPPromptBuilder": "SxCP Prompt Builder",
    "SxCPPromptBuilderFromConfigs": "SxCP Prompt Builder From Configs",
}