from __future__ import annotations

import json

try:
    from .prompt_builder import (
        build_camera_config_json,
        build_insta_of_options_json,
        build_insta_of_pair,
        build_prompt,
        build_seed_config_json,
        camera_angle_choices,
        camera_distance_choices,
        camera_lens_choices,
        camera_mode_choices,
        camera_orientation_choices,
        camera_phone_choices,
        camera_priority_choices,
        camera_shot_choices,
        category_choices,
        subcategory_choices,
    )
    from .caption_naturalizer import naturalize_caption
except ImportError:
    from prompt_builder import (
        build_camera_config_json,
        build_insta_of_options_json,
        build_insta_of_pair,
        build_prompt,
        build_seed_config_json,
        camera_angle_choices,
        camera_distance_choices,
        camera_lens_choices,
        camera_mode_choices,
        camera_orientation_choices,
        camera_phone_choices,
        camera_priority_choices,
        camera_shot_choices,
        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}),
                "camera_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="",
        camera_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 "",
            camera_config=camera_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 SxCPCameraControl:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "camera_mode": (camera_mode_choices(), {"default": "handheld_selfie"}),
                "shot_size": (camera_shot_choices(), {"default": "auto"}),
                "angle": (camera_angle_choices(), {"default": "auto"}),
                "lens": (camera_lens_choices(), {"default": "smartphone_wide"}),
                "distance": (camera_distance_choices(), {"default": "arm_length"}),
                "orientation": (camera_orientation_choices(), {"default": "vertical_story"}),
                "phone_visibility": (camera_phone_choices(), {"default": "phone_visible"}),
                "priority": (camera_priority_choices(), {"default": "locked"}),
            }
        }

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

    def build(
        self,
        camera_mode,
        shot_size,
        angle,
        lens,
        distance,
        orientation,
        phone_visibility,
        priority,
    ):
        return (
            build_camera_config_json(
                camera_mode=camera_mode,
                shot_size=shot_size,
                angle=angle,
                lens=lens,
                distance=distance,
                orientation=orientation,
                phone_visibility=phone_visibility,
                priority=priority,
            ),
        )


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,
        )


class SxCPInstaOFOptions:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "softcore_cast": (["solo", "same_as_hardcore"], {"default": "solo"}),
                "hardcore_cast": (["use_counts", "couple", "threesome", "group"], {"default": "use_counts"}),
                "hardcore_women_count": ("INT", {"default": 1, "min": 0, "max": 12, "step": 1}),
                "hardcore_men_count": ("INT", {"default": 1, "min": 0, "max": 12, "step": 1}),
                "softcore_level": (["social_tease", "lingerie_tease", "implied_nude", "explicit_tease"], {"default": "lingerie_tease"}),
                "hardcore_level": (["explicit", "hardcore"], {"default": "hardcore"}),
                "platform_style": (["hybrid", "instagram", "onlyfans"], {"default": "hybrid"}),
                "continuity": (["same_creator_same_room", "same_creator_new_scene"], {"default": "same_creator_same_room"}),
                "softcore_camera_mode": (camera_mode_choices(), {"default": "handheld_selfie"}),
                "hardcore_camera_mode": (["same_as_softcore"] + camera_mode_choices(), {"default": "same_as_softcore"}),
            }
        }

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

    def build(
        self,
        softcore_cast,
        hardcore_cast,
        hardcore_women_count,
        hardcore_men_count,
        softcore_level,
        hardcore_level,
        platform_style,
        continuity,
        softcore_camera_mode,
        hardcore_camera_mode,
    ):
        return (
            build_insta_of_options_json(
                softcore_cast=softcore_cast,
                hardcore_cast=hardcore_cast,
                hardcore_women_count=hardcore_women_count,
                hardcore_men_count=hardcore_men_count,
                softcore_level=softcore_level,
                hardcore_level=hardcore_level,
                platform_style=platform_style,
                continuity=continuity,
                softcore_camera_mode=softcore_camera_mode,
                hardcore_camera_mode=hardcore_camera_mode,
            ),
        )


class SxCPInstaOFPromptPair:
    @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}),
                "ethnicity": (["any", "asian", "white_asian"], {"default": "any"}),
                "figure": (["curvy", "balanced", "bombshell"], {"default": "curvy"}),
                "no_plus_women": ("BOOLEAN", {"default": False}),
                "no_black": ("BOOLEAN", {"default": False}),
                "trigger": ("STRING", {"default": "sxcpinup_coloredpencil"}),
                "prepend_trigger_to_prompt": ("BOOLEAN", {"default": True}),
            },
            "optional": {
                "seed_config": ("STRING", {"default": "", "multiline": True}),
                "options_json": ("STRING", {"default": "", "multiline": True}),
                "camera_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", "STRING", "STRING")
    RETURN_NAMES = (
        "softcore_prompt",
        "hardcore_prompt",
        "softcore_negative_prompt",
        "hardcore_negative_prompt",
        "softcore_caption",
        "hardcore_caption",
        "shared_descriptor",
        "metadata_json",
    )
    FUNCTION = "build"
    CATEGORY = "prompt_builder"

    def build(
        self,
        row_number,
        start_index,
        seed,
        ethnicity,
        figure,
        no_plus_women,
        no_black,
        trigger,
        prepend_trigger_to_prompt,
        seed_config="",
        options_json="",
        camera_config="",
        extra_positive="",
        extra_negative="",
    ):
        row = build_insta_of_pair(
            row_number=row_number,
            start_index=start_index,
            seed=seed,
            ethnicity=ethnicity,
            figure=figure,
            no_plus_women=no_plus_women,
            no_black=no_black,
            trigger=trigger,
            prepend_trigger_to_prompt=prepend_trigger_to_prompt,
            seed_config=seed_config or "",
            options_json=options_json or "",
            camera_config=camera_config or "",
            extra_positive=extra_positive or "",
            extra_negative=extra_negative or "",
        )
        return (
            row["softcore_prompt"],
            row["hardcore_prompt"],
            row["softcore_negative_prompt"],
            row["hardcore_negative_prompt"],
            row["softcore_caption"],
            row["hardcore_caption"],
            row["shared_descriptor"],
            json.dumps(row, ensure_ascii=True, sort_keys=True),
        )


NODE_CLASS_MAPPINGS = {
    "SxCPPromptBuilder": SxCPPromptBuilder,
    "SxCPSeedControl": SxCPSeedControl,
    "SxCPCameraControl": SxCPCameraControl,
    "SxCPCaptionNaturalizer": SxCPCaptionNaturalizer,
    "SxCPInstaOFOptions": SxCPInstaOFOptions,
    "SxCPInstaOFPromptPair": SxCPInstaOFPromptPair,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "SxCPPromptBuilder": "SxCP Prompt Builder",
    "SxCPSeedControl": "SxCP Seed Control",
    "SxCPCameraControl": "SxCP Camera Control",
    "SxCPCaptionNaturalizer": "SxCP Caption Naturalizer",
    "SxCPInstaOFOptions": "SxCP Insta/OF Options",
    "SxCPInstaOFPromptPair": "SxCP Insta/OF Prompt Pair",
}