"""
Monkey-patch for multitalk_loop.py prompt selection.

Replaces the original prompt selection block (iteration-based) with a
frame-based prompt_schedule-aware version. Falls back to original behavior
when no prompt_schedule is present in text_embeds.
"""

import logging
import textwrap

log = logging.getLogger("ComfyUI-MultiTalkPromptSchedule")

_PATCHED = False


def _select_prompt(text_embeds, audio_start_idx, iteration_count, log_fn):
    """Frame-based prompt selection with fallback to iteration-based.

    Returns:
        positive: list of prompt embeddings for this window
    """
    prompt_schedule = text_embeds.get("prompt_schedule", None)
    nag_schedule = text_embeds.get("nag_schedule", None)

    if prompt_schedule is not None:
        # Restore original NAG embeds before swapping (prevents cache corruption)
        if "_orig_nag_prompt_embeds" not in text_embeds:
            import copy as _copy
            text_embeds["_orig_nag_prompt_embeds"] = _copy.deepcopy(text_embeds.get("nag_prompt_embeds"))
            text_embeds["_orig_negative_prompt_embeds"] = _copy.deepcopy(text_embeds.get("negative_prompt_embeds"))
        else:
            text_embeds["nag_prompt_embeds"] = text_embeds["_orig_nag_prompt_embeds"]
            text_embeds["negative_prompt_embeds"] = text_embeds["_orig_negative_prompt_embeds"]

        # Frame-based prompt selection (MultiTalkPromptSchedule node)
        prompt_index = len(prompt_schedule) - 1  # default to last
        for idx, (sched_start, sched_end) in enumerate(prompt_schedule):
            if audio_start_idx < sched_end:
                prompt_index = idx
                break
        positive = [text_embeds["prompt_embeds"][prompt_index]]

        # Per-scene NAG negative if available
        if nag_schedule is not None and prompt_index in nag_schedule:
            nag_index = nag_schedule[prompt_index]
            nag_embed = [text_embeds["all_nag_embeds"][nag_index]]
            text_embeds["nag_prompt_embeds"] = nag_embed
            text_embeds["negative_prompt_embeds"] = nag_embed
            log_fn(f"Using per-scene NAG negative for scene {prompt_index}")
        elif nag_schedule is not None:
            if "all_nag_embeds" in text_embeds:
                text_embeds["nag_prompt_embeds"] = [text_embeds["all_nag_embeds"][0]]
                text_embeds["negative_prompt_embeds"] = [text_embeds["all_nag_embeds"][0]]

        sched_start, sched_end = prompt_schedule[prompt_index]
        log_fn(
            f"Using scheduled prompt {prompt_index} "
            f"(frames {sched_start}-{sched_end}) "
            f"at audio_start_idx={audio_start_idx}"
        )
    elif len(text_embeds["prompt_embeds"]) > 1:
        # Fallback: iteration-based (original Kijai behavior, pipe separator)
        prompt_index = min(
            iteration_count, len(text_embeds["prompt_embeds"]) - 1
        )
        positive = [text_embeds["prompt_embeds"][prompt_index]]
        log_fn(f"Using prompt index: {prompt_index}")
    else:
        positive = text_embeds["prompt_embeds"]

    return positive


def apply_prompt_schedule_patch():
    """Monkey-patch multitalk_loop to support prompt_schedule in text_embeds."""
    global _PATCHED
    if _PATCHED:
        return

    import importlib

    loop_mod = None
    # Try both naming conventions (hyphen vs underscore)
    for module_name in [
        "custom_nodes.ComfyUI-WanVideoWrapper.multitalk.multitalk_loop",
        "custom_nodes.ComfyUI_WanVideoWrapper.multitalk.multitalk_loop",
    ]:
        try:
            loop_mod = importlib.import_module(module_name)
            break
        except ImportError:
            continue

    if loop_mod is None:
        log.warning(
            "[MultiTalkPromptSchedule] ComfyUI-WanVideoWrapper not found. "
            "Patch not applied — the node will still load but scheduling "
            "won't work until WanVideoWrapper is installed."
        )
        return

    import inspect
    import re

    original_fn = loop_mod.multitalk_loop

    # Get the source to check if the patch is already present
    # (e.g. user already manually patched)
    try:
        src = inspect.getsource(original_fn)
        if "prompt_schedule" in src:
            log.info(
                "[MultiTalkPromptSchedule] multitalk_loop already contains "
                "prompt_schedule logic. Skipping monkey-patch."
            )
            _PATCHED = True
            return
    except (OSError, TypeError):
        # Can't inspect source (compiled, etc.) — proceed with patch
        pass

    # ── Build the patched function ──────────────────────────────────────
    # We wrap the original function and intercept the prompt selection.
    # The trick: multitalk_loop is a method (self, **kwargs).  We wrap it
    # so that before it runs, we inject a _select_prompt callback into
    # text_embeds that the original code path will pick up.
    #
    # Strategy: We pre-process text_embeds so that when the original code
    # hits `if len(text_embeds["prompt_embeds"]) > 1:`, the embeds are
    # already set to the correct prompt for the FIRST window.  For
    # subsequent windows we can't intercept since the loop is internal.
    #
    # Better strategy: We actually patch the source code of the function
    # at the AST/text level to inject our prompt_schedule block.

    src = inspect.getsource(original_fn)

    # The original prompt selection block we need to replace:
    original_block = textwrap.dedent("""\
        # Use the appropriate prompt for this section
        if len(text_embeds["prompt_embeds"]) > 1:
            prompt_index = min(iteration_count, len(text_embeds["prompt_embeds"]) - 1)
            positive = [text_embeds["prompt_embeds"][prompt_index]]
            log.info(f"Using prompt index: {prompt_index}")
        else:
            positive = text_embeds["prompt_embeds"]""")

    patched_block = textwrap.dedent("""\
        # Use the appropriate prompt for this section
        # [Patched by ComfyUI-MultiTalkPromptSchedule]
        log.info(f"[MultiTalkPromptSchedule PATCH] text_embeds keys: {list(text_embeds.keys())}")
        prompt_schedule = text_embeds.get("prompt_schedule", None)
        nag_schedule = text_embeds.get("nag_schedule", None)
        if prompt_schedule is not None:
            # Restore original NAG embeds before swapping (prevents cache corruption between runs)
            if "_orig_nag_prompt_embeds" not in text_embeds:
                import copy as _copy
                text_embeds["_orig_nag_prompt_embeds"] = _copy.deepcopy(text_embeds.get("nag_prompt_embeds"))
                text_embeds["_orig_negative_prompt_embeds"] = _copy.deepcopy(text_embeds.get("negative_prompt_embeds"))
            else:
                text_embeds["nag_prompt_embeds"] = text_embeds["_orig_nag_prompt_embeds"]
                text_embeds["negative_prompt_embeds"] = text_embeds["_orig_negative_prompt_embeds"]
            prompt_index = len(prompt_schedule) - 1
            for idx, (sched_start, sched_end) in enumerate(prompt_schedule):
                if audio_start_idx < sched_end:
                    prompt_index = idx
                    break
            positive = [text_embeds["prompt_embeds"][prompt_index]]
            if nag_schedule is not None and prompt_index in nag_schedule:
                nag_index = nag_schedule[prompt_index]
                nag_embed = [text_embeds["all_nag_embeds"][nag_index]]
                text_embeds["nag_prompt_embeds"] = nag_embed
                text_embeds["negative_prompt_embeds"] = nag_embed
                log.info(f"Using per-scene NAG negative for scene {prompt_index}")
            elif nag_schedule is not None:
                if "all_nag_embeds" in text_embeds:
                    text_embeds["nag_prompt_embeds"] = [text_embeds["all_nag_embeds"][0]]
                    text_embeds["negative_prompt_embeds"] = [text_embeds["all_nag_embeds"][0]]
            sched_start, sched_end = prompt_schedule[prompt_index]
            log.info(f"Using scheduled prompt {prompt_index} (frames {sched_start}-{sched_end}) at audio_start_idx={audio_start_idx}")
        elif len(text_embeds["prompt_embeds"]) > 1:
            prompt_index = min(iteration_count, len(text_embeds["prompt_embeds"]) - 1)
            positive = [text_embeds["prompt_embeds"][prompt_index]]
            log.info(f"Using prompt index: {prompt_index}")
        else:
            positive = text_embeds["prompt_embeds"]""")

    # Normalize whitespace for matching: the source may have different
    # indentation (8 spaces inside the while loop).
    def _normalize_indent(block):
        """Strip common leading whitespace so we can match flexibly."""
        lines = block.splitlines()
        return [l.strip() for l in lines if l.strip()]

    original_lines = _normalize_indent(original_block)

    # Find the block in the source by matching stripped lines
    src_lines = src.splitlines()
    match_start = None
    for i in range(len(src_lines)):
        if src_lines[i].strip() == original_lines[0]:
            # Check if subsequent lines match
            matched = True
            for j, orig_line in enumerate(original_lines):
                if i + j >= len(src_lines):
                    matched = False
                    break
                if src_lines[i + j].strip() != orig_line:
                    matched = False
                    break
            if matched:
                match_start = i
                break

    if match_start is None:
        log.warning(
            "[MultiTalkPromptSchedule] Could not find the original prompt "
            "selection block in multitalk_loop source. The upstream code may "
            "have changed. Patch not applied."
        )
        return

    # Determine the indentation of the matched block
    indent = ""
    first_line = src_lines[match_start]
    indent = first_line[: len(first_line) - len(first_line.lstrip())]

    # Build the replacement with proper indentation
    patched_lines = []
    for line in patched_block.splitlines():
        if line.strip():
            patched_lines.append(indent + line)
        else:
            patched_lines.append("")

    # Replace in source
    match_end = match_start + len(original_lines)
    new_src_lines = (
        src_lines[:match_start] + patched_lines + src_lines[match_end:]
    )

    # Patch 2: deepcopy audio_embedding to prevent cache corruption.
    # The original `audio_embedding = multitalk_audio_embeds` is just a
    # reference. When audio is short, the loop pads audio_embedding[i]
    # in-place (torch.cat) which corrupts the cached node output.
    # On the next run, embeddings are already padded → wrong sizes → frozen video.
    audio_orig = "audio_embedding = multitalk_audio_embeds"
    audio_patched = (
        "# [Patched by ComfyUI-MultiTalkPromptSchedule] deepcopy to prevent cache corruption\n"
        "audio_embedding = [t.clone() for t in multitalk_audio_embeds]"
    )
    patched_count = 0
    for i, line in enumerate(new_src_lines):
        if audio_orig in line:
            indent2 = line[: len(line) - len(line.lstrip())]
            new_src_lines[i] = "\n".join(indent2 + l for l in audio_patched.splitlines())
            patched_count += 1
    if patched_count > 0:
        log.info(
            "[MultiTalkPromptSchedule] Also patched audio_embedding "
            "to deepcopy (prevents cache corruption on re-run)."
        )
    else:
        log.warning(
            "[MultiTalkPromptSchedule] Could not find "
            "'audio_embedding = multitalk_audio_embeds' to patch."
        )

    new_src = "\n".join(new_src_lines)

    # Remove the `def multitalk_loop(self, **kwargs):` indentation
    # since we'll exec it as a module-level function.
    # Actually, we need to dedent the function definition first.
    new_src = textwrap.dedent(new_src)

    # Compile and exec in the module's namespace
    try:
        code = compile(new_src, loop_mod.__file__, "exec")
        exec(code, loop_mod.__dict__)
        log.info(
            "[MultiTalkPromptSchedule] Successfully patched multitalk_loop "
            "with prompt_schedule support."
        )
        _PATCHED = True

        # CRITICAL: nodes_sampler.py imports multitalk_loop with
        # `from .multitalk.multitalk_loop import multitalk_loop`
        # which creates a separate reference. We must update it too,
        # otherwise the sampler keeps calling the original function.
        #
        # We scan ALL loaded modules because ComfyUI may register them
        # under unexpected names in sys.modules (hyphen vs underscore, etc.)
        import sys
        patched_fn = loop_mod.multitalk_loop
        patched_count = 0
        for mod_key, mod in list(sys.modules.items()):
            if mod is None or mod is loop_mod:
                continue
            try:
                if hasattr(mod, 'multitalk_loop') and callable(getattr(mod, 'multitalk_loop')):
                    old_fn = getattr(mod, 'multitalk_loop')
                    if old_fn is not patched_fn:
                        setattr(mod, 'multitalk_loop', patched_fn)
                        log.info(
                            f"[MultiTalkPromptSchedule] Patched multitalk_loop "
                            f"reference in module: {mod_key}"
                        )
                        patched_count += 1
            except Exception:
                continue
        if patched_count == 0:
            log.warning(
                "[MultiTalkPromptSchedule] No other modules with multitalk_loop "
                "found. The sampler may not be loaded yet — registering an "
                "import hook to patch it when it loads."
            )
            # Deferred patch: if nodes_sampler loads AFTER us, patch it then
            _deferred_patch_fn = patched_fn
            _deferred_loop_mod = loop_mod

            class _MultitalkImportHook:
                def find_module(self, name, path=None):
                    if 'nodes_sampler' in name:
                        return self
                    return None

                def load_module(self, name):
                    # Let the normal import proceed
                    sys.meta_path.remove(self)
                    mod = importlib.import_module(name)
                    # Now patch it
                    if hasattr(mod, 'multitalk_loop'):
                        mod.multitalk_loop = _deferred_patch_fn
                        log.info(
                            f"[MultiTalkPromptSchedule] Deferred patch applied "
                            f"to {name}"
                        )
                    sys.meta_path.append(self)
                    return mod

            sys.meta_path.append(_MultitalkImportHook())
        else:
            log.info(
                f"[MultiTalkPromptSchedule] Patched {patched_count} module(s) "
                f"with updated multitalk_loop reference."
            )

    except Exception as e:
        log.error(
            f"[MultiTalkPromptSchedule] Failed to compile patched "
            f"multitalk_loop: {e}. Patch not applied."
        )