import re
import logging

log = logging.getLogger(__name__)


class MusePromptSplitter:
    """
    Splits one block of LLM-generated text into up to 4 separate segment prompt
    strings, parsed by 'PROMPT_N:' markers. Pairs with an LLM node (e.g. Ollama,
    QwenVL) whose system prompt instructs it to expand a simple scene description
    into fully-formatted LTX 2.3 segment prompts, labeled PROMPT_1, PROMPT_2, etc.
    Wire each output into MuseDirectorSamplerV2_5's segment_override_1..4 inputs.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "llm_output": ("STRING", {"multiline": True, "default": "", "forceInput": True}),
            },
        }

    RETURN_TYPES = ("STRING",) * 4
    RETURN_NAMES = tuple(f"prompt_{i}" for i in range(1, 5))
    FUNCTION = "execute"
    CATEGORY = "Muse Collective"
    DESCRIPTION = (
        "Splits one LLM-generated text block into up to 4 separate segment prompts, "
        "parsed by 'PROMPT_N:' markers. Wire each output into "
        "MuseDirectorSamplerV2_5's segment_override_1..4 inputs."
    )

    def execute(self, llm_output):
        results = [""] * 4
        if not llm_output or not llm_output.strip():
            log.warning("[MusePromptSplitter] Empty input — all outputs will be empty.")
            return tuple(results)

        matches = re.findall(
            r'PROMPT_(\d+)\s*:\s*(.*?)(?=PROMPT_\d+\s*:|\Z)',
            llm_output,
            re.IGNORECASE | re.DOTALL,
        )

        if not matches:
            log.warning(
                "[MusePromptSplitter] No 'PROMPT_N:' markers found in LLM output — check the "
                "LLM's system prompt is instructing it to label outputs correctly. "
                "Raw output (first 200 chars): %s", llm_output[:200],
            )
            return tuple(results)

        for idx_str, text in matches:
            try:
                idx = int(idx_str)
            except ValueError:
                continue
            if 1 <= idx <= 4:
                results[idx - 1] = text.strip()

        log.info("[MusePromptSplitter] Parsed %d labeled prompt(s) from LLM output.",
                  sum(1 for r in results if r))
        return tuple(results)


NODE_CLASS_MAPPINGS = {
    "MusePromptSplitter": MusePromptSplitter,
}
NODE_DISPLAY_NAME_MAPPINGS = {
    "MusePromptSplitter": "Muse Prompt Splitter",
}