from __future__ import annotations

try:
    from .caption_naturalizer import naturalize_caption
    from .caption_policy import caption_profile_choices
    from .krea_formatter import format_krea2_prompt
    from .sdxl_formatter import (
        format_sdxl_prompt,
        sdxl_formatter_profile_choices,
        sdxl_quality_preset_choices,
        sdxl_style_preset_choices,
    )
except ImportError:  # Allows local smoke tests from the repository root.
    from caption_naturalizer import naturalize_caption
    from caption_policy import caption_profile_choices
    from krea_formatter import format_krea2_prompt
    from sdxl_formatter import (
        format_sdxl_prompt,
        sdxl_formatter_profile_choices,
        sdxl_quality_preset_choices,
        sdxl_style_preset_choices,
    )


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"}),
                "caption_profile": (caption_profile_choices(), {"default": "manual_controls"}),
                "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": {
                "source_text_input": ("STRING", {"forceInput": True}),
                "metadata_json": ("STRING", {"forceInput": True}),
            },
        }

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

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


class SxCPKrea2Formatter:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "source_text": ("STRING", {"default": "", "multiline": True}),
                "input_hint": (["auto", "metadata_json", "prompt"], {"default": "auto"}),
                "target": (["auto", "single", "softcore", "hardcore"], {"default": "auto"}),
                "detail_level": (["balanced", "concise", "dense"], {"default": "balanced"}),
                "style_mode": (["preserve", "photographic", "minimal"], {"default": "preserve"}),
                "preserve_trigger": ("BOOLEAN", {"default": False}),
            },
            "optional": {
                "source_text_input": ("STRING", {"forceInput": True}),
                "metadata_json": ("STRING", {"forceInput": True}),
                "negative_prompt": ("STRING", {"forceInput": True}),
                "extra_positive": ("STRING", {"default": "", "multiline": True}),
                "extra_negative": ("STRING", {"default": "", "multiline": True}),
            },
        }

    RETURN_TYPES = ("STRING", "STRING", "STRING", "STRING", "STRING", "STRING", "STRING")
    RETURN_NAMES = (
        "krea_prompt",
        "negative_prompt",
        "krea_softcore_prompt",
        "krea_hardcore_prompt",
        "softcore_negative_prompt",
        "hardcore_negative_prompt",
        "method",
    )
    FUNCTION = "build"
    CATEGORY = "prompt_builder"

    def build(
        self,
        source_text,
        input_hint,
        target,
        detail_level,
        style_mode,
        preserve_trigger,
        source_text_input="",
        metadata_json="",
        negative_prompt="",
        extra_positive="",
        extra_negative="",
    ):
        active_source_text = source_text_input or source_text or ""
        row = format_krea2_prompt(
            source_text=active_source_text,
            metadata_json=metadata_json or "",
            negative_prompt=negative_prompt or "",
            input_hint=input_hint,
            target=target,
            detail_level=detail_level,
            style_mode=style_mode,
            preserve_trigger=preserve_trigger,
            extra_positive=extra_positive or "",
            extra_negative=extra_negative or "",
        )
        return (
            row["krea_prompt"],
            row["negative_prompt"],
            row["krea_softcore_prompt"],
            row["krea_hardcore_prompt"],
            row["softcore_negative_prompt"],
            row["hardcore_negative_prompt"],
            row["method"],
        )


class SxCPSDXLFormatter:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "source_text": ("STRING", {"default": "", "multiline": True}),
                "input_hint": (["auto", "metadata_json", "prompt"], {"default": "auto"}),
                "target": (["auto", "single", "softcore", "hardcore"], {"default": "auto"}),
                "formatter_profile": (sdxl_formatter_profile_choices(), {"default": "manual_controls"}),
                "style_preset": (sdxl_style_preset_choices(), {"default": "flat_vector_pony"}),
                "quality_preset": (sdxl_quality_preset_choices(), {"default": "pony_high"}),
                "trigger": ("STRING", {"default": "mythp0rt", "multiline": False}),
                "prepend_trigger_to_prompt": ("BOOLEAN", {"default": True}),
                "preserve_trigger": ("BOOLEAN", {"default": False}),
                "nude_weight": ("FLOAT", {"default": 1.29, "min": 0.1, "max": 3.0, "step": 0.01}),
            },
            "optional": {
                "source_text_input": ("STRING", {"forceInput": True}),
                "metadata_json": ("STRING", {"forceInput": True}),
                "negative_prompt": ("STRING", {"forceInput": True}),
                "custom_style": ("STRING", {"default": "", "multiline": True}),
                "custom_quality": ("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")
    RETURN_NAMES = (
        "sdxl_prompt",
        "negative_prompt",
        "sdxl_softcore_prompt",
        "sdxl_hardcore_prompt",
        "softcore_negative_prompt",
        "hardcore_negative_prompt",
        "method",
    )
    FUNCTION = "build"
    CATEGORY = "prompt_builder"

    def build(
        self,
        source_text,
        input_hint,
        target,
        formatter_profile,
        style_preset,
        quality_preset,
        trigger,
        prepend_trigger_to_prompt,
        preserve_trigger,
        nude_weight,
        source_text_input="",
        metadata_json="",
        negative_prompt="",
        custom_style="",
        custom_quality="",
        extra_positive="",
        extra_negative="",
    ):
        active_source_text = source_text_input or source_text or ""
        row = format_sdxl_prompt(
            source_text=active_source_text,
            metadata_json=metadata_json or "",
            negative_prompt=negative_prompt or "",
            input_hint=input_hint,
            target=target,
            formatter_profile=formatter_profile,
            style_preset=style_preset,
            quality_preset=quality_preset,
            trigger=trigger,
            prepend_trigger=prepend_trigger_to_prompt,
            preserve_trigger=preserve_trigger,
            nude_weight=nude_weight,
            custom_style=custom_style or "",
            custom_quality=custom_quality or "",
            extra_positive=extra_positive or "",
            extra_negative=extra_negative or "",
        )
        return (
            row["sdxl_prompt"],
            row["negative_prompt"],
            row["sdxl_softcore_prompt"],
            row["sdxl_hardcore_prompt"],
            row["softcore_negative_prompt"],
            row["hardcore_negative_prompt"],
            row["method"],
        )


NODE_CLASS_MAPPINGS = {
    "SxCPCaptionNaturalizer": SxCPCaptionNaturalizer,
    "SxCPKrea2Formatter": SxCPKrea2Formatter,
    "SxCPSDXLFormatter": SxCPSDXLFormatter,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "SxCPCaptionNaturalizer": "SxCP Caption Naturalizer",
    "SxCPKrea2Formatter": "SxCP Krea2 Formatter",
    "SxCPSDXLFormatter": "SxCP SDXL Formatter",
}