# prompt_comment.py
#
# ComfyUI のプロンプトでコメントアウト機能を提供するノード
# 指定した記号以降のテキストを削除し、クリーンなプロンプトを出力します
#
# 対応コメント記法:
#   - # (Python/Shell style)
#   - // (C/JavaScript style)
#   - -- (SQL style)
#   - Custom (任意の文字列)

import re


class PromptComment:
    """
    プロンプト内のコメントを除去するノード。
    各行で指定したコメント記号以降のテキストを削除します。
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "text": ("STRING", {"multiline": True, "default": ""}),
                "comment_style": (["#", "//", "--", "Custom"],),
                "custom_marker": ("STRING", {"default": "#"}),
                "remove_empty_lines": ("BOOLEAN", {"default": True}),
                "trim_whitespace": ("BOOLEAN", {"default": True}),
            }
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("text",)
    FUNCTION = "strip_comments"
    CATEGORY = "text/utils"

    def strip_comments(
        self,
        text,
        comment_style,
        custom_marker,
        remove_empty_lines,
        trim_whitespace,
    ):
        # コメントマーカーを決定
        if comment_style == "Custom":
            marker = custom_marker
        else:
            marker = comment_style

        if not marker:
            # マーカーが空の場合は何もしない
            return (text,)

        # 各行を処理
        lines = text.split("\n")
        processed_lines = []

        for line in lines:
            # コメントマーカーの位置を検索
            # 文字列リテラル内のマーカーは無視しない（シンプルな実装）
            marker_pos = line.find(marker)

            if marker_pos != -1:
                # コメントマーカー以前のテキストのみ保持
                line = line[:marker_pos]

            if trim_whitespace:
                line = line.strip()

            # 空行の処理
            if remove_empty_lines:
                if line:  # 空でない行のみ追加
                    processed_lines.append(line)
            else:
                processed_lines.append(line)

        # 行を結合
        result = "\n".join(processed_lines)

        # 最終的なトリム
        if trim_whitespace:
            result = result.strip()

        return (result,)


class PromptCommentMultiMarker:
    """
    複数のコメントマーカーに対応したプロンプトコメント除去ノード。
    #, //, -- を同時に認識します。
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "text": ("STRING", {"multiline": True, "default": ""}),
                "use_hash": ("BOOLEAN", {"default": True}),      # #
                "use_double_slash": ("BOOLEAN", {"default": False}),  # //
                "use_double_dash": ("BOOLEAN", {"default": False}),   # --
                "remove_empty_lines": ("BOOLEAN", {"default": True}),
                "trim_whitespace": ("BOOLEAN", {"default": True}),
            }
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("text",)
    FUNCTION = "strip_comments"
    CATEGORY = "text/utils"

    def strip_comments(
        self,
        text,
        use_hash,
        use_double_slash,
        use_double_dash,
        remove_empty_lines,
        trim_whitespace,
    ):
        # 有効なマーカーのリストを構築
        markers = []
        if use_hash:
            markers.append("#")
        if use_double_slash:
            markers.append("//")
        if use_double_dash:
            markers.append("--")

        if not markers:
            # マーカーが指定されていない場合は何もしない
            return (text,)

        # 各行を処理
        lines = text.split("\n")
        processed_lines = []

        for line in lines:
            # 最も早い位置にあるコメントマーカーを検索
            earliest_pos = len(line)
            for marker in markers:
                pos = line.find(marker)
                if pos != -1 and pos < earliest_pos:
                    earliest_pos = pos

            if earliest_pos < len(line):
                # コメントマーカー以前のテキストのみ保持
                line = line[:earliest_pos]

            if trim_whitespace:
                line = line.strip()

            # 空行の処理
            if remove_empty_lines:
                if line:  # 空でない行のみ追加
                    processed_lines.append(line)
            else:
                processed_lines.append(line)

        # 行を結合
        result = "\n".join(processed_lines)

        # 最終的なトリム
        if trim_whitespace:
            result = result.strip()

        return (result,)


NODE_CLASS_MAPPINGS = {
    "PromptComment": PromptComment,
    "PromptCommentMultiMarker": PromptCommentMultiMarker,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "PromptComment": "Prompt Comment (Strip)",
    "PromptCommentMultiMarker": "Prompt Comment (Multi-Marker)",
}