from __future__ import annotations

import json

try:
    from .prompt_builder import build_prompt, build_seed_config_json, category_choices, subcategory_choices
    from .caption_naturalizer import naturalize_caption
except ImportError:
    from prompt_builder import build_prompt, build_seed_config_json, category_choices, subcategory_choices
    from caption_naturalizer import naturalize_caption


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": (["full", "minimal"], {"default": "full"}),
                "ethnicity": (["any", "asian", "white_asian"], {"default": "any"}),
                "poses": (["standard", "evocative"], {"default": "standard"}),
                "backside_bias": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}),
                "figure": (["curvy", "balanced", "bombshell"], {"default": "curvy"}),
                "no_plus_women": ("BOOLEAN", {"default": False}),
                "no_black": ("BOOLEAN", {"default": False}),
                "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": {
                "seed_config": ("STRING", {"default": "", "multiline": True}),
                "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,
        backside_bias,
        figure,
        no_plus_women,
        no_black,
        women_count,
        men_count,
        minimal_clothing_ratio,
        standard_pose_ratio,
        trigger,
        prepend_trigger_to_prompt,
        seed_config="",
        extra_positive="",
        extra_negative="",
    ):
        row = build_prompt(
            category=category,
            subcategory=subcategory,
            row_number=row_number,
            start_index=start_index,
            seed=seed,
            clothing=clothing,
            ethnicity=ethnicity,
            poses=poses,
            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 "",
        )
        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 SxCPSeedControl:
    @classmethod
    def INPUT_TYPES(cls):
        seed_spec = {"default": -1, "min": -1, "max": 0xFFFFFFFF, "step": 1}
        return {
            "required": {
                "category_seed": ("INT", seed_spec),
                "subcategory_seed": ("INT", seed_spec),
                "content_seed": ("INT", seed_spec),
                "person_seed": ("INT", seed_spec),
                "scene_seed": ("INT", seed_spec),
                "pose_seed": ("INT", seed_spec),
                "role_seed": ("INT", seed_spec),
                "expression_seed": ("INT", seed_spec),
                "composition_seed": ("INT", seed_spec),
            }
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("seed_config",)
    FUNCTION = "build"
    CATEGORY = "prompt_builder"

    def build(
        self,
        category_seed,
        subcategory_seed,
        content_seed,
        person_seed,
        scene_seed,
        pose_seed,
        role_seed,
        expression_seed,
        composition_seed,
    ):
        return (
            build_seed_config_json(
                category_seed=category_seed,
                subcategory_seed=subcategory_seed,
                content_seed=content_seed,
                person_seed=person_seed,
                scene_seed=scene_seed,
                pose_seed=pose_seed,
                role_seed=role_seed,
                expression_seed=expression_seed,
                composition_seed=composition_seed,
            ),
        )


class SxCPCaptionNaturalizer:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "source_text": ("STRING", {"default": "", "multiline": True}),
                "input_hint": (["auto", "metadata_json", "caption_or_prompt"], {"default": "auto"}),
                "detail_level": (["balanced", "concise", "dense"], {"default": "balanced"}),
                "style_policy": (["drop_style_tail", "keep_style_terms"], {"default": "drop_style_tail"}),
                "trigger": ("STRING", {"default": "sxcppnl7"}),
                "include_trigger": ("BOOLEAN", {"default": True}),
            },
            "optional": {
                "metadata_json": ("STRING", {"default": "", "multiline": True}),
            },
        }

    RETURN_TYPES = ("STRING", "STRING")
    RETURN_NAMES = ("natural_caption", "method")
    FUNCTION = "build"
    CATEGORY = "prompt_builder"

    def build(
        self,
        source_text,
        input_hint,
        detail_level,
        style_policy,
        trigger,
        include_trigger,
        metadata_json="",
    ):
        return naturalize_caption(
            source_text=source_text or "",
            metadata_json=metadata_json or "",
            input_hint=input_hint,
            trigger=trigger,
            include_trigger=include_trigger,
            detail_level=detail_level,
            style_policy=style_policy,
        )


NODE_CLASS_MAPPINGS = {
    "SxCPPromptBuilder": SxCPPromptBuilder,
    "SxCPSeedControl": SxCPSeedControl,
    "SxCPCaptionNaturalizer": SxCPCaptionNaturalizer,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "SxCPPromptBuilder": "SxCP Prompt Builder",
    "SxCPSeedControl": "SxCP Seed Control",
    "SxCPCaptionNaturalizer": "SxCP Caption Naturalizer",
}