"""
Muse Collective — MuseGuide
============================
Applies reference/keyframe guidance to a video latent before sampling.

Handles:
  - Image keyframe guidance (reference image locked to latent start)
  - Motion video segments via IC-LoRA
  - Licon MSR (Prefix) identity injection, via ComfyUI-LTXVideo's
    LTXAddVideoICLoRAGuide/LTXVAddGuide nodes
  - Retake mode (preserve base video, regenerate a region)
  - Conditioning crop after sampling
"""

import asyncio
import json
import logging
import math
import os
import platform
import subprocess
import wave

import av
import numpy as np
import torch
import torch.nn.functional as F

import comfy
import comfy.sd
import comfy.utils
import folder_paths
import node_helpers
from aiohttp import web
from server import PromptServer
from comfy_extras import nodes_lt

log = logging.getLogger(__name__)


# ── Character-analysis backend route ────────────────────────────────────────
# Backs the timeline character-card panel's "Analyze" button (JS: muse_director_v2.js).
# Sends the reference image(s) to a locally-configured vision-language model and
# returns a short physical-description string for the character/object.

_MUSE_PROVIDER_DEFAULTS = {
    "ollama": {"url": "http://127.0.0.1:11434", "model": "huihui_ai/qwen3.5-abliterated:2b"},
    "lmstudio": {"url": "http://127.0.0.1:1234", "model": ""},
    "custom": {"url": "", "model": ""},
    # Gemini's OpenAI-compatible endpoint. The API key is read server-side from
    # GEMINI_API_KEY so it never touches the browser or gets saved into a workflow
    # JSON. No trailing /v1 — Gemini's compat layer serves /chat/completions
    # directly off this base, unlike the LM Studio/custom /v1/... convention.
    "gemini": {"url": "https://generativelanguage.googleapis.com/v1beta/openai", "model": "gemini-2.5-flash"},
}

_MUSE_ANALYZE_PROMPT = (
    "Identify whether the main subject of this image is a person, or an object/"
    "prop/creature/vehicle. Then write a short, plain description: for a person, "
    "cover hair (color and style), face, build, and what they're wearing, in two "
    "or three sentences; for an object, cover what it is plus its shape, color, "
    "material, and any standout details, in two sentences. Skip the background "
    "entirely, don't say which category you picked, and keep it brief."
)


def _muse_resolve_provider(request_data):
    provider = (request_data.get("provider") or "ollama").lower()
    defaults = _MUSE_PROVIDER_DEFAULTS.get(provider, _MUSE_PROVIDER_DEFAULTS["ollama"])
    base_url = (request_data.get("base_url") or defaults["url"]).rstrip("/")
    model = request_data.get("model") or defaults["model"]
    return provider, base_url, model


def _error(message, status=200):
    return web.json_response({"status": "error", "message": message}, status=status)


def _strip_data_url_prefix(b64_images):
    items = b64_images if isinstance(b64_images, list) else [b64_images]
    cleaned = []
    for item in items:
        cleaned.append(item.split(",", 1)[1] if "," in item else item)
    return cleaned


def _strip_thinking_block(text):
    if "<think>" in text:
        return text.split("</think>")[-1].strip()
    return text


async def _call_ollama_vision(session, base_url, model_name, images_b64):
    payload = {"model": model_name, "prompt": _MUSE_ANALYZE_PROMPT,
               "images": images_b64, "stream": False, "keep_alive": 0}
    async with session.post(f"{base_url}/api/generate", json=payload, timeout=120) as resp:
        if resp.status != 200:
            raise RuntimeError(f"Ollama HTTP {resp.status}: {await resp.text()}")
        body = await resp.json()
        return (body.get("response") or "").strip()


def _build_vision_chat_content(images_b64):
    content = [{"type": "text", "text": _MUSE_ANALYZE_PROMPT}]
    for b64 in images_b64:
        content.append({"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{b64}"}})
    return content


async def _call_gemini_vision(session, base_url, model_name, images_b64):
    api_key = os.environ.get("GEMINI_API_KEY")
    if not api_key:
        raise RuntimeError("GEMINI_API_KEY environment variable is not set. Set it and restart ComfyUI.")
    payload = {
        "model": model_name,
        "messages": [{"role": "user", "content": _build_vision_chat_content(images_b64)}],
        "max_tokens": 2048, "stream": False,
    }
    headers = {"Authorization": f"Bearer {api_key}"}
    async with session.post(f"{base_url}/chat/completions", json=payload, headers=headers, timeout=120) as resp:
        if resp.status != 200:
            raise RuntimeError(f"Gemini HTTP {resp.status}: {await resp.text()}")
        body = await resp.json()
        try:
            return (body["choices"][0]["message"].get("content") or "").strip()
        except (KeyError, IndexError, TypeError):
            raise RuntimeError("Unexpected response shape from Gemini.")


async def _call_openai_compatible_vision(session, base_url, provider, model_name, images_b64):
    payload = {
        "model": model_name,
        "messages": [{"role": "user", "content": _build_vision_chat_content(images_b64)}],
        "max_tokens": 2048, "stream": False,
    }
    async with session.post(f"{base_url}/v1/chat/completions", json=payload, timeout=120) as resp:
        if resp.status != 200:
            raise RuntimeError(f"{provider} HTTP {resp.status}: {await resp.text()}")
        body = await resp.json()
        try:
            msg = body["choices"][0]["message"]
            return (msg.get("content") or msg.get("reasoning_content") or "").strip()
        except (KeyError, IndexError, TypeError):
            raise RuntimeError(f"Unexpected response shape from {provider}.")


@PromptServer.instance.routes.post("/muse_director/analyze_character")
async def muse_analyze_character_endpoint(request):
    try:
        import aiohttp
        request_data = await request.json()
        char_index = int(request_data.get("char_index", 0))
        provider, base_url, model_name = _muse_resolve_provider(request_data)

        if provider == "off":
            return _error("Analyze is set to Off / Manual.")

        raw_images = request_data.get("image_b64", "")
        if not raw_images:
            return _error("No image provided for analysis.")
        images_b64 = _strip_data_url_prefix(raw_images)
        if not images_b64:
            return _error("No valid base64 images decoded.")

        if provider in ("lmstudio", "custom") and not model_name:
            return _error(f"No model name set for {provider}. Enter your loaded model's name.")

        log.info("[MuseDirector] Analyzing Character %d via %s (%s, model '%s')...",
                 char_index + 1, provider, base_url, model_name)

        try:
            async with aiohttp.ClientSession() as session:
                if provider == "ollama":
                    generated_text = await _call_ollama_vision(session, base_url, model_name, images_b64)
                elif provider == "gemini":
                    generated_text = await _call_gemini_vision(session, base_url, model_name, images_b64)
                else:
                    generated_text = await _call_openai_compatible_vision(session, base_url, provider, model_name, images_b64)
        except aiohttp.ClientConnectorError:
            return _error(f"Could not connect to {provider} at {base_url}. Make sure the server is running and reachable.")
        except RuntimeError as e:
            return _error(str(e))

        generated_text = _strip_thinking_block(generated_text)
        log.info("[MuseDirector] Character analysis complete: %s", generated_text)
        return web.json_response({"status": "success", "description": generated_text})

    except Exception as e:
        log.error("[MuseDirector] Failed to analyze character: %s", e)
        return _error(str(e), status=500)


@PromptServer.instance.routes.post("/muse_director/unload_ollama")
async def muse_unload_ollama_endpoint(request):
    """Evict the analysis model from VRAM right before a generation run, so the VLM
    doesn't compete with LTX for VRAM."""
    try:
        import aiohttp
        data = await request.json()
        provider, base_url, model_name = _muse_resolve_provider(data)
        if provider != "ollama":
            return web.json_response({"status": "skipped", "message": f"Provider '{provider}' has no unload step."})
        async with aiohttp.ClientSession() as session:
            payload = {"model": model_name, "keep_alive": 0}
            async with session.post(f"{base_url}/api/generate", json=payload, timeout=15) as response:
                await response.text()
        return web.json_response({"status": "success"})
    except Exception as e:
        log.warning("[MuseDirector] unload_ollama failed (non-fatal): %s", e)
        return web.json_response({"status": "error", "message": str(e)})


# ── Timeline media upload + audio waveform backend ──────────────────────────
# Backs the timeline's drag-and-drop video/audio track: a chunked upload route
# (ComfyUI's default endpoint rejects large files with a 413), an upload
# dedup check, an audio-waveform/peaks route so the timeline can draw a
# waveform for an already-uploaded file, and an "open upload folder" button.

_MUSE_UPLOAD_SUBDIR = "muse"
_MUSE_AUDIO_EXTENSIONS = {".wav", ".mp3", ".ogg", ".flac", ".m4a"}
_MUSE_WAVEFORM_POINTS = 200


def _muse_upload_dir():
    upload_dir = os.path.join(folder_paths.get_input_directory(), _MUSE_UPLOAD_SUBDIR)
    os.makedirs(upload_dir, exist_ok=True)
    return upload_dir


def _pcm16_peaks(raw_bytes, num_points=_MUSE_WAVEFORM_POINTS):
    """Downsample raw mono 16-bit PCM audio into `num_points` normalized
    (0-1) peak-amplitude values, for drawing a compact waveform."""
    samples = np.frombuffer(raw_bytes, dtype=np.int16)
    if samples.size == 0:
        return [0.0] * num_points
    stride = max(1, samples.size // num_points)
    peaks = []
    for i in range(num_points):
        window = samples[i * stride:(i + 1) * stride]
        peaks.append(float(np.max(np.abs(window))) / 32767.0 if window.size else 0.0)
    return peaks


def _decode_track_to_pcm16_mono(path, sample_rate):
    """Decode the first audio stream of a media file via PyAV, resampled to
    mono 16-bit PCM at `sample_rate`. Returns raw bytes, or None if the file
    has no audio track."""
    with av.open(path) as container:
        if not container.streams.audio:
            return None
        stream = container.streams.audio[0]
        resampler = av.AudioResampler(format="s16", layout="mono", rate=sample_rate)
        chunks = bytearray()
        for frame in container.decode(stream):
            for resampled in resampler.resample(frame):
                chunks.extend(resampled.to_ndarray().tobytes())
        for resampled in resampler.resample(None):  # flush
            chunks.extend(resampled.to_ndarray().tobytes())
        return bytes(chunks) if chunks else None


def _read_wav_peaks(wav_path):
    with wave.open(wav_path, "rb") as wav_file:
        frame_count = wav_file.getnframes()
        if frame_count == 0:
            return [0.0] * _MUSE_WAVEFORM_POINTS
        return _pcm16_peaks(wav_file.readframes(frame_count))


def _get_audio_peaks(audio_path):
    """Peaks for a file that IS audio (as opposed to a video to extract audio
    from) — used by both the get-audio route and the video extractor below."""
    if os.path.splitext(audio_path)[1].lower() == ".wav":
        try:
            return _read_wav_peaks(audio_path)
        except Exception:
            pass  # fall through to the generic PyAV path
    try:
        pcm = _decode_track_to_pcm16_mono(audio_path, sample_rate=8000)
        return _pcm16_peaks(pcm) if pcm else None
    except Exception as e:
        log.warning("[MuseDirector] Failed to read audio peaks for %s: %s", audio_path, e)
        return None


def _extract_audio_from_video(video_path):
    """Pull the audio track out of a video file into a cached sidecar .wav
    (44.1kHz mono), returning (relative_wav_path, peaks) — or (None, None) if
    the video has no audio track. Reuses an existing extraction if present
    and already at the right sample rate."""
    cached_wav = os.path.splitext(video_path)[0] + "_extracted_audio.wav"
    if os.path.exists(cached_wav) and os.path.getsize(cached_wav) > 44:
        try:
            with wave.open(cached_wav, "rb") as check:
                if check.getframerate() == 44100:
                    input_dir = folder_paths.get_input_directory()
                    rel = os.path.relpath(cached_wav, input_dir).replace("\\", "/")
                    return rel, _read_wav_peaks(cached_wav)
        except Exception:
            pass

    try:
        pcm = _decode_track_to_pcm16_mono(video_path, sample_rate=44100)
    except Exception as e:
        log.warning("[MuseDirector] Error extracting audio from %s: %s", video_path, e)
        return None, None
    if not pcm:
        return None, None

    with wave.open(cached_wav, "wb") as wav_file:
        wav_file.setnchannels(1)
        wav_file.setsampwidth(2)
        wav_file.setframerate(44100)
        wav_file.writeframes(pcm)

    input_dir = folder_paths.get_input_directory()
    rel_path = os.path.relpath(cached_wav, input_dir).replace("\\", "/")
    return rel_path, _pcm16_peaks(pcm)


def _find_matching_upload(filename, size_str):
    """Dedup check for the chunked-upload flow: does a file matching this
    name (and, if given, exact byte size) already exist in the muse upload
    folder or the root input dir? Falls back to a `*_<filename>` suffix
    search, since the frontend prefixes retried uploads to avoid collisions."""
    input_dir = folder_paths.get_input_directory()
    upload_dir = _muse_upload_dir()

    def size_matches(path):
        if not size_str:
            return True
        try:
            return os.path.getsize(path) == int(size_str)
        except (ValueError, OSError):
            return True

    for candidate in (os.path.join(upload_dir, filename), os.path.join(input_dir, filename)):
        if os.path.isfile(candidate) and size_matches(candidate):
            return os.path.relpath(candidate, input_dir).replace("\\", "/")

    base_name = os.path.basename(filename)
    suffix = f"_{base_name}"
    for search_dir in (upload_dir, input_dir):
        if not os.path.isdir(search_dir):
            continue
        for entry in os.listdir(search_dir):
            if entry == base_name or entry.endswith(suffix):
                candidate = os.path.join(search_dir, entry)
                if os.path.isfile(candidate) and size_matches(candidate):
                    return os.path.relpath(candidate, input_dir).replace("\\", "/")
    return None


@PromptServer.instance.routes.get("/muse_director/check_upload")
async def muse_check_upload_endpoint(request):
    filename = request.query.get("filename", "")
    if not filename:
        return web.json_response({"exists": False})
    match = _find_matching_upload(filename, request.query.get("size", ""))
    if match:
        return web.json_response({"exists": True, "name": match})
    return web.json_response({"exists": False})


@PromptServer.instance.routes.get("/muse_director/get_audio")
async def muse_get_audio_endpoint(request):
    filename = request.query.get("filename")
    if not filename:
        return web.json_response({"error": "Missing filename"}, status=400)

    input_dir = folder_paths.get_input_directory()
    clean_name = filename.replace("\\", "/")
    file_path = os.path.join(input_dir, clean_name)
    if not os.path.exists(file_path):
        basename = os.path.basename(clean_name)
        in_muse_dir = os.path.join(input_dir, _MUSE_UPLOAD_SUBDIR, basename)
        file_path = in_muse_dir if os.path.exists(in_muse_dir) else os.path.join(input_dir, basename)

    if not os.path.isfile(file_path):
        return web.json_response({"error": "File not found"}, status=404)

    if os.path.splitext(file_path)[1].lower() in _MUSE_AUDIO_EXTENSIONS:
        peaks = None
        try:
            peaks = _get_audio_peaks(file_path)
        except Exception as e:
            log.warning("[MuseDirector] Failed to read audio peaks for %s: %s", file_path, e)
        rel_path = os.path.relpath(file_path, input_dir).replace("\\", "/")
        return web.json_response({"audio_file": rel_path, "peaks": peaks})

    audio_file, peaks = None, None
    try:
        loop = asyncio.get_event_loop()
        audio_file, peaks = await loop.run_in_executor(None, _extract_audio_from_video, file_path)
    except Exception as e:
        log.warning("[MuseDirector] Error extracting audio from %s: %s", file_path, e)
    return web.json_response({"audio_file": audio_file, "peaks": peaks})


@PromptServer.instance.routes.get("/muse_director/open_upload_folder")
async def muse_open_upload_folder_endpoint(request):
    upload_dir = _muse_upload_dir()
    try:
        system_name = platform.system()
        if system_name == "Windows":
            subprocess.Popen(["explorer", os.path.normpath(upload_dir)])
        elif system_name == "Darwin":
            subprocess.Popen(["open", upload_dir])
        else:
            subprocess.Popen(["xdg-open", upload_dir])
        return web.json_response({"success": True})
    except Exception as e:
        log.warning("[MuseDirector] Failed to open upload folder: %s", e)
        return web.json_response({"success": False, "error": str(e)}, status=500)


def _write_upload_chunk_to_disk(file_field, file_path, append):
    with open(file_path, "ab" if append else "wb") as f:
        f.write(file_field.file.read())


@PromptServer.instance.routes.post("/muse_director/upload_chunk")
async def muse_upload_chunk_endpoint(request):
    """Chunked upload for large video/audio files — ComfyUI's built-in upload
    route enforces a size limit (413 Payload Too Large) that this bypasses by
    streaming the file to disk a piece at a time."""
    form = await request.post()
    file_field = form.get("file")
    raw_filename = form.get("filename")
    chunk_index = int(form.get("chunk_index"))
    total_chunks = int(form.get("total_chunks"))

    upload_dir = _muse_upload_dir()
    filename = os.path.basename(raw_filename)  # strip any path components — no traversal
    file_path = os.path.join(upload_dir, filename)
    if not os.path.realpath(file_path).startswith(os.path.realpath(upload_dir)):
        return web.json_response({"error": "Invalid filename"}, status=400)

    loop = asyncio.get_event_loop()
    await loop.run_in_executor(None, _write_upload_chunk_to_disk, file_field, file_path, chunk_index > 0)

    if chunk_index != total_chunks - 1:
        return web.json_response({"status": "ok"})

    audio_file, peaks = None, None
    try:
        audio_file, peaks = await loop.run_in_executor(None, _extract_audio_from_video, file_path)
    except Exception as e:
        log.warning("[MuseDirector] Error extracting audio from final upload chunk: %s", e)
    return web.json_response({
        "name": f"{_MUSE_UPLOAD_SUBDIR}/{filename}",
        "audio_file": audio_file,
        "peaks": peaks,
    })


# ── Licon MSR (Prefix) — identity injection via ComfyUI-LTXVideo ───────────────
# Needs LTXAddVideoICLoRAGuide, a ComfyUI-LTXVideo custom node (not ComfyUI core) —
# it builds the IC-LoRA reference-slideshow conditioning the identity signal rides
# on. Deliberate exception to this file otherwise sticking to core: no core
# primitive does what that node does.

def _call_node(node_class, **kwargs):
    """ComfyUI ships two node calling conventions: newer comfy_api io.ComfyNode
    subclasses expose a classmethod `execute`, older-style nodes are instantiated
    and called through whatever method name `FUNCTION` points at. Try the
    classmethod form first, fall back to instantiating."""
    if hasattr(node_class, "execute"):
        return node_class.execute(**kwargs)
    entrypoint_name = getattr(node_class, "FUNCTION", None)
    if not entrypoint_name:
        raise RuntimeError(f"{node_class!r} has no 'execute' classmethod and no FUNCTION attribute")
    instance = node_class()
    if not hasattr(instance, entrypoint_name):
        raise RuntimeError(f"{node_class!r} has no method '{entrypoint_name}'")
    return getattr(instance, entrypoint_name)(**kwargs)


def _flatten_node_return(raw_result):
    """A node call can come back as a plain tuple/list, a dict with a 'result'
    key, or a comfy_api io.NodeOutput-style object wrapping its payload behind
    one of a few possible attribute names — collapse all of those into a plain
    tuple of outputs."""
    if raw_result is None:
        return ()
    if isinstance(raw_result, (tuple, list)):
        return tuple(raw_result)
    if isinstance(raw_result, dict):
        payload = raw_result.get("result")
        if isinstance(payload, (tuple, list)):
            return tuple(payload)
    for attr_name in ("result", "outputs", "values", "args"):
        payload = getattr(raw_result, attr_name, None)
        if callable(payload):
            try:
                payload = payload()
            except Exception:
                continue
        if isinstance(payload, (tuple, list)):
            return tuple(payload)
    return (raw_result,)


# ── Image resize helper ───────────────────────────────────────────────────────

def _resize_image(tensor: torch.Tensor, target_w: int, target_h: int,
                  method: str, divisible_by: int = 1) -> torch.Tensor:
    """Resize [N, H, W, C] float32 tensor to target dimensions."""
    def snap(val, div):
        div = max(1, int(div))
        return max(div, (val // div) * div)

    tw = snap(target_w, divisible_by)
    th = snap(target_h, divisible_by)

    N, H, W, C = tensor.shape
    if H == th and W == tw:
        return tensor

    t = tensor.permute(0, 3, 1, 2)  # NHWC → NCHW

    if method == "stretch to fit":
        r = F.interpolate(t, size=(th, tw), mode="bilinear", align_corners=False)

    elif method in ("maintain aspect ratio",):
        ratio = min(tw / W, th / H)
        nw = snap(int(W * ratio), divisible_by)
        nh = snap(int(H * ratio), divisible_by)
        r = F.interpolate(t, size=(nh, nw), mode="bilinear", align_corners=False)

    elif method in ("pad", "pad green"):
        ratio = min(tw / W, th / H)
        nw = snap(int(W * ratio), divisible_by)
        nh = snap(int(H * ratio), divisible_by)
        inner = F.interpolate(t, size=(nh, nw), mode="bilinear", align_corners=False)
        pl = (tw - nw) // 2
        pt = (th - nh) // 2
        if method == "pad green":
            r = torch.zeros((N, C, th, tw), dtype=t.dtype, device=t.device)
            r[:, 0] = 102 / 255.0
            r[:, 1] = 1.0
            r[:, :, pt:pt + nh, pl:pl + nw] = inner
        else:
            r = F.pad(inner, (pl, tw - nw - pl, pt, th - nh - pt), value=0)

    elif method == "crop":
        ratio = max(tw / W, th / H)
        nw, nh = int(W * ratio), int(H * ratio)
        inner = F.interpolate(t, size=(nh, nw), mode="bilinear", align_corners=False)
        l = (nw - tw) // 2
        top = (nh - th) // 2
        r = inner[:, :, top:top + th, l:l + tw]

    else:
        r = F.interpolate(t, size=(th, tw), mode="bilinear", align_corners=False)

    return r.permute(0, 2, 3, 1)  # NCHW → NHWC


# ── Latent helpers ────────────────────────────────────────────────────────────

def _clone_noise_mask(latent, latent_image):
    if "noise_mask" in latent and latent["noise_mask"] is not None:
        return latent["noise_mask"].clone()
    B, _, F, _, _ = latent_image.shape
    return torch.ones((B, 1, F, 1, 1), dtype=torch.float32, device=latent_image.device)


def _resize_latent_spatial(latent_image, noise_mask, w, h, method):
    b, c, f, lh, lw = latent_image.shape
    if lw == w and lh == h:
        return latent_image, noise_mask
    lat4 = latent_image.permute(0, 2, 1, 3, 4).reshape(b * f, c, lh, lw)
    lat4 = comfy.utils.common_upscale(lat4, w, h, method, "disabled")
    latent_image = lat4.reshape(b, f, c, h, w).permute(0, 2, 1, 3, 4)
    if noise_mask is not None and (noise_mask.shape[-1] > 1 or noise_mask.shape[-2] > 1):
        mb, mc, mf, mh, mw = noise_mask.shape
        m4 = noise_mask.permute(0, 2, 1, 3, 4).reshape(mb * mf, mc, mh, mw)
        m4 = comfy.utils.common_upscale(m4, w, h, method, "disabled")
        noise_mask = m4.reshape(mb, mf, mc, h, w).permute(0, 2, 1, 3, 4)
    return latent_image, noise_mask


def _ceil_to_multiple(value, multiple):
    multiple = max(1, int(multiple))
    return int(math.ceil(value / multiple) * multiple)


def _snap_latent_to_downscale(latent_image, noise_mask, downscale_factor, method):
    factor = int(max(1, round(float(downscale_factor))))
    if factor <= 1:
        return latent_image, noise_mask
    _, _, _, h, w = latent_image.shape
    nw = _ceil_to_multiple(w, factor)
    nh = _ceil_to_multiple(h, factor)
    if nw == w and nh == h:
        return latent_image, noise_mask
    log.warning("[MuseGuide] Auto-snapping latent %sx%s → %sx%s for IC-LoRA downscale %s", w, h, nw, nh, factor)
    return _resize_latent_spatial(latent_image, noise_mask, nw, nh, method)


def _dilate_latent(latent: dict, horizontal_scale: int, vertical_scale: int) -> dict:
    if horizontal_scale == 1 and vertical_scale == 1:
        return latent
    samples = latent["samples"]
    mask = latent.get("noise_mask", None)
    ds = samples.shape[:3] + (samples.shape[3] * vertical_scale, samples.shape[4] * horizontal_scale)
    dilated = torch.zeros(ds, device=samples.device, dtype=samples.dtype)
    dilated[..., ::vertical_scale, ::horizontal_scale] = samples
    dms = (dilated.shape[0], 1, dilated.shape[2], dilated.shape[3], dilated.shape[4])
    dmask = torch.full(dms, -1.0, device=samples.device, dtype=samples.dtype)
    dmask[..., ::vertical_scale, ::horizontal_scale] = (mask if mask is not None else 1.0)
    return {"samples": dilated, "noise_mask": dmask}


# ── IC-LoRA helpers ───────────────────────────────────────────────────────────

def _load_lora_model_only(model, ic_lora_name, strength):
    lora_path = folder_paths.get_full_path_or_raise("loras", ic_lora_name)
    lora, metadata = comfy.utils.load_torch_file(lora_path, safe_load=True, return_metadata=True)
    try:
        downscale = float(metadata["reference_downscale_factor"])
    except Exception:
        downscale = 1.0
        log.warning("[MuseGuide] Could not read reference_downscale_factor from %s, using 1.0", ic_lora_name)
    if strength == 0:
        return model, downscale
    model_lora, _ = comfy.sd.load_lora_for_models(model, None, lora, strength, 0)
    return model_lora, downscale


def _encode_video_iclora(vae, latent_w, latent_h, images, scale_factors,
                          latent_downscale_factor, crop, use_tiled, tile_size, tile_overlap,
                          resize_method="crop"):
    tsf, wsf, hsf = scale_factors
    keep = ((images.shape[0] - 1) // tsf) * tsf + 1
    images = images[:keep]
    tw = max(8, int(latent_w * wsf / latent_downscale_factor))
    th = max(8, int(latent_h * hsf / latent_downscale_factor))
    if resize_method == "maintain aspect ratio":
        resize_method = "pad"
    pixels = _resize_image(images, tw, th, resize_method, divisible_by=1)
    enc = pixels[:, :, :, :3]
    if use_tiled:
        guide_lat = vae.encode_tiled(enc, tile_x=tile_size, tile_y=tile_size, overlap=tile_overlap)
    else:
        guide_lat = vae.encode(enc)
    return pixels, guide_lat


# ── Guide attention helpers ───────────────────────────────────────────────────

def _get_guide_attention_entries(conditioning):
    for item in conditioning:
        entries = item[1].get("guide_attention_entries")
        if entries is not None:
            return entries
    return []


def _set_guide_attention_entries(conditioning, entries):
    return node_helpers.conditioning_set_values(conditioning, {"guide_attention_entries": entries})


def _append_guide_attention_entry(conditioning, pre_filter_count, latent_shape, attention_strength=1.0):
    entries = [*_get_guide_attention_entries(conditioning)]
    entries.append({
        "pre_filter_count": int(pre_filter_count),
        "strength": float(attention_strength),
        "pixel_mask": None,
        "latent_shape": list(latent_shape),
    })
    return _set_guide_attention_entries(conditioning, entries)


# ── Video loading ─────────────────────────────────────────────────────────────

class _ResampleGuideFrames:
    @staticmethod
    def execute(images, source_fps, target_fps, target_num_frames, mode):
        if images is None:
            return images
        n = int(images.shape[0])
        target_num_frames = int(target_num_frames)
        if n <= 1:
            return images.repeat(target_num_frames, 1, 1, 1) if target_num_frames > 1 else images
        source_fps = float(max(0.001, source_fps))
        target_fps = float(max(0.001, target_fps))
        if target_num_frames <= 0:
            duration = (n - 1) / source_fps
            target_num_frames = max(1, int(round(duration * target_fps)) + 1)
        if target_num_frames == n and abs(target_fps - source_fps) < 1e-6:
            return images
        positions = torch.linspace(0, n - 1, target_num_frames, device=images.device, dtype=torch.float32)
        if mode == "nearest":
            idx = torch.round(positions).long().clamp(0, n - 1)
            return images.index_select(0, idx)
        idx0 = torch.floor(positions).long().clamp(0, n - 1)
        idx1 = torch.ceil(positions).long().clamp(0, n - 1)
        alpha = (positions - idx0.to(positions.dtype)).view(-1, 1, 1, 1)
        f0 = images.index_select(0, idx0).float()
        f1 = images.index_select(0, idx1).float()
        return (f0 * (1 - alpha) + f1 * alpha).to(images.dtype)


def _resolve_video_path(video_file):
    if os.path.isabs(str(video_file)) and os.path.exists(str(video_file)):
        return str(video_file)
    input_dir = folder_paths.get_input_directory()
    candidate = os.path.join(input_dir, str(video_file))
    if os.path.exists(candidate):
        return candidate
    try:
        annotated = folder_paths.get_annotated_filepath(str(video_file))
        if annotated and os.path.exists(annotated):
            return annotated
    except Exception:
        pass
    raise FileNotFoundError(f"[MuseGuide] Could not find video: {video_file}")


def _load_motion_video_frames(video_file, trim_start_frames, length_frames, director_fps, resample_mode="nearest"):
    path = _resolve_video_path(video_file)
    target_fps = max(1.0, float(director_fps))
    start_s = max(0.0, float(trim_start_frames) / target_fps)
    dur_s = max(0.0, float(length_frames) / target_fps)
    end_s = start_s + dur_s if dur_s > 0 else None
    container = av.open(path)
    stream = container.streams.video[0]
    stream.thread_type = "AUTO"
    try:
        source_fps = float(stream.average_rate) if stream.average_rate else float(stream.base_rate)
    except Exception:
        source_fps = target_fps
    if source_fps <= 0:
        source_fps = target_fps
    if start_s > 0:
        try:
            seek_pts = int(max(0, start_s - 0.5) / float(stream.time_base)) if stream.time_base else int(max(0, start_s - 0.5) * av.time_base)
            container.seek(seek_pts, stream=stream, backward=True)
        except Exception as e:
            log.warning("[MuseGuide] Seek failed: %s", e)
    frames = []
    decoded = 0
    for frame in container.decode(stream):
        t = frame.time if frame.time is not None else (
            float(frame.pts * stream.time_base) if frame.pts is not None and stream.time_base else float(decoded / source_fps)
        )
        decoded += 1
        if t < start_s - 0.01:
            continue
        if end_s is not None and t >= end_s:
            break
        frames.append(frame.to_ndarray(format="rgb24"))
    container.close()
    if not frames:
        raise ValueError(f"[MuseGuide] No frames decoded from: {video_file}")
    frames_np = np.array(frames, dtype=np.float32) / 255.0
    images = torch.from_numpy(frames_np)
    target_count = max(1, int(round(float(length_frames))))
    return _ResampleGuideFrames.execute(images, source_fps, target_fps, target_count, resample_mode)


# ── Conditioning crop helpers ─────────────────────────────────────────────────

def _cond_get_value(conditioning, key, default=None):
    for item in conditioning:
        val = item[1].get(key)
        if val is not None:
            return val
    return default


def _get_exact_crop_count(conditioning):
    val = _cond_get_value(conditioning, "nghtdrp_guide_crop_latent_frames")
    if val is not None:
        try:
            return max(0, int(val))
        except Exception:
            return 0
    kf = _cond_get_value(conditioning, "keyframe_idxs")
    if kf is None:
        return 0
    try:
        return int(torch.unique(kf[:, 0, :, 0]).shape[0])
    except Exception:
        return 0


def _get_noise_mask_for_crop(latent):
    img = latent["samples"]
    mask = latent.get("noise_mask")
    if mask is None:
        B, _, F, _, _ = img.shape
        return torch.ones((B, 1, F, 1, 1), dtype=torch.float32, device=img.device)
    return mask.clone()


# ── Main node ─────────────────────────────────────────────────────────────────

class MuseGuide:
    """
    Encodes reference images and motion video segments as keyframes into the latent.
    Supports IC-LoRA motion guidance, Licon MSR (Prefix) identity injection, and retake mode.
    """

    @classmethod
    def INPUT_TYPES(cls):
        loras = folder_paths.get_filename_list("loras") or ["put_ic_lora_in_ComfyUI_models_loras"]
        return {
            "required": {
                "positive":   ("CONDITIONING",),
                "negative":   ("CONDITIONING",),
                "vae":        ("VAE",),
                "latent":     ("LATENT",),
                "guide_data": ("GUIDE_DATA",),
            },
            "optional": {
                "motion_guide_data":      ("MOTION_GUIDE_DATA",),
                "model":                  ("MODEL",),
                "ic_lora_name":           (["None"] + loras, {"default": "None"}),
                "ic_lora_strength":       ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01}),
                "scale_by":               ("FLOAT", {"default": 1.0, "min": 0.01, "max": 8.0, "step": 0.01}),
                "upscale_method":         (["nearest-exact", "bilinear", "area", "bicubic", "bislerp"], {"default": "bicubic"}),
                "image_attention_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
                "crop":                   (["disabled", "center"], {"default": "center"}),
                "auto_snap_ic_grid":      ("BOOLEAN", {"default": True}),
                "use_tiled_encode":       ("BOOLEAN", {"default": False}),
                "tile_size":              ("INT", {"default": 256, "min": 64, "max": 512, "step": 32}),
                "tile_overlap":           ("INT", {"default": 64, "min": 16, "max": 256, "step": 16}),
                "retake_mode":            ("BOOLEAN", {"default": False}),
                "msr_strength":           ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01,
                                            "tooltip": "Licon MSR (Prefix) only. Per-stage override for the reference "
                                                       "strength baked into guide_data['msr'] by the Director — 0 means "
                                                       "'use the Director's reference_strength', >0 wins (set it lower "
                                                       "on a refinement stage so references hold detail without "
                                                       "repainting the opening)."}),
            },
        }

    RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT", "MODEL", "FLOAT")
    RETURN_NAMES = ("positive", "negative", "latent", "model", "latent_downscale_factor")
    FUNCTION = "execute"
    CATEGORY = "Muse Collective"

    @classmethod
    def execute(cls, positive, negative, vae, latent, guide_data,
                motion_guide_data=None, model=None, ic_lora_name="None",
                ic_lora_strength=1.0, scale_by=1.0, upscale_method="bicubic",
                image_attention_strength=1.0, crop="center", auto_snap_ic_grid=True,
                use_tiled_encode=False, tile_size=256, tile_overlap=64, retake_mode=False,
                msr_strength=0.0):

        motion_segments = (motion_guide_data or {}).get("segments", []) if motion_guide_data else []
        image_guides_count = len(guide_data.get("images", [])) if guide_data else 0
        log.info("[MuseGuide] execute — image_guides: %d, motion_segments: %d, ic_lora: %s",
                 image_guides_count, len(motion_segments), ic_lora_name)

        # Licon MSR (Prefix): the Director handed us raw references via guide_data["msr"]
        # instead of the normal images/insert_frames/strengths lists. Inject them at THIS
        # node's resolution (so a half-res Stage 1 and full-res Stage 2 each stay
        # self-consistent) and return early — none of the normal guide paths below apply.
        _msr = (guide_data or {}).get("msr")
        if _msr is not None:
            latent_downscale_factor = 1.0
            if model is not None and ic_lora_name != "None":
                model, latent_downscale_factor = _load_lora_model_only(model, ic_lora_name, ic_lora_strength)
            return cls._inject_msr(positive, negative, vae, latent, _msr, model,
                                    latent_downscale_factor, msr_strength)

        # Resolve resize method
        active_resize = (guide_data or {}).get("resize_method") or \
                        (motion_guide_data or {}).get("resize_method") or \
                        ("crop" if crop == "center" else "stretch to fit")

        # IC-LoRA
        latent_downscale_factor = 1.0
        if model is not None and ic_lora_name != "None":
            model, latent_downscale_factor = _load_lora_model_only(model, ic_lora_name, ic_lora_strength)

        scale_factors = vae.downscale_index_formula
        latent_image = latent["samples"].clone()
        noise_mask = _clone_noise_mask(latent, latent_image)

        # scale_by is for IC-LoRA reference video encoding only — never resize the generation latent here
        if auto_snap_ic_grid and model is not None and ic_lora_name != "None":
            latent_image, noise_mask = _snap_latent_to_downscale(latent_image, noise_mask, latent_downscale_factor, upscale_method)

        _, _, latent_length, latent_height, latent_width = latent_image.shape
        initial_latent_length = int(latent_length)
        time_scale_factor = scale_factors[0]
        ltxv_length = (latent_length - 1) * time_scale_factor + 1
        is_lora_active = model is not None and ic_lora_name != "None"

        # Parse timeline data for retake detection
        timeline_data_str = (guide_data or {}).get("timeline_data", "{}")
        try:
            tdata = json.loads(timeline_data_str)
        except Exception:
            tdata = {}

        is_retake_active = bool(retake_mode) or tdata.get("retakeMode", False)
        is_empty_latent = latent_image.abs().max().item() < 1e-5

        # ── Retake Mode ───────────────────────────────────────────────────────
        if is_retake_active:
            log.info("[MuseGuide] Retake mode active. is_empty_latent: %s", is_empty_latent)
            target_pix_w = latent_width * 32
            target_pix_h = latent_height * 32

            retake_start   = int(tdata.get("retakeStart", 0))
            retake_len     = int(tdata.get("retakeLength", 0))
            retake_strength = float(tdata.get("retakeStrength", 1.0))
            start_frame    = int((guide_data or {}).get("start_frame", 0))
            relative_start = max(0, retake_start - start_frame)

            l_start = relative_start // time_scale_factor
            l_end   = int(math.ceil((relative_start + retake_len) / time_scale_factor))
            l_start = min(l_start, latent_length)
            l_end   = min(l_end, latent_length)

            need_base_video = not (not is_empty_latent and l_start == 0 and l_end >= latent_length)

            retake_vid_info = tdata.get("retakeVideo") or {}
            video_file = retake_vid_info.get("imageFile", "") if isinstance(retake_vid_info, dict) else ""
            director_fps = float((motion_guide_data or {}).get("frame_rate", (guide_data or {}).get("frame_rate", 24)))

            if need_base_video:
                if not video_file:
                    raise ValueError(
                        "[MuseGuide] Retake mode is active but no base video has been selected on the timeline."
                    )
                try:
                    video_frames = _load_motion_video_frames(video_file, start_frame, ltxv_length, director_fps)
                    retake_resize = active_resize if active_resize != "maintain aspect ratio" else "pad"
                    pixels = _resize_image(video_frames, target_pix_w, target_pix_h, retake_resize, divisible_by=1)
                    keep = ((pixels.shape[0] - 1) // time_scale_factor) * time_scale_factor + 1
                    enc_src = pixels[:keep, :, :, :3]
                    if use_tiled_encode:
                        base_latent = vae.encode_tiled(enc_src, tile_x=tile_size, tile_y=tile_size, overlap=tile_overlap)
                    else:
                        base_latent = vae.encode(enc_src)
                    base_latent = base_latent.to(device=latent_image.device, dtype=latent_image.dtype)
                    paste_len = min(base_latent.shape[2], latent_length)
                    if is_empty_latent:
                        latent_image[:, :, :paste_len] = base_latent[:, :, :paste_len]
                    else:
                        if l_start > 0:
                            latent_image[:, :, :l_start] = base_latent[:, :, :l_start]
                        if l_end < paste_len:
                            latent_image[:, :, l_end:paste_len] = base_latent[:, :, l_end:paste_len]
                except Exception as exc:
                    log.warning("[MuseGuide] Base video load/encode failed: %s", exc)

            noise_mask = torch.zeros_like(noise_mask)
            if l_end > l_start:
                noise_mask[:, :, l_start:l_end] = retake_strength

            crop_frames = max(0, int(latent_image.shape[2]) - initial_latent_length)
            positive = node_helpers.conditioning_set_values(positive, {"nghtdrp_guide_crop_latent_frames": crop_frames})
            negative = node_helpers.conditioning_set_values(negative, {"nghtdrp_guide_crop_latent_frames": crop_frames})
            return (positive, negative, {"samples": latent_image, "noise_mask": noise_mask}, model, float(latent_downscale_factor))

        # ── Ghost Mask pre-extend ────────────────────────────────────────────
        # The incoming latent is left at its clean, visible length upstream —
        # extending it to fit the hidden reference tail is this node's job, so
        # video/audio lengths stay matched before this point. The extra region
        # is zeros/ones (fully denoisable); the reference keyframes appended
        # below get written into it, and it's cropped back off after sampling.
        ghost_pre_extend = int((guide_data or {}).get("ghost_pre_extend", 0)) if guide_data else 0
        if ghost_pre_extend > 0:
            b, c, _, h, w = latent_image.shape
            mb, mc, _, mh, mw = noise_mask.shape
            latent_image = torch.cat(
                [latent_image, torch.zeros((b, c, ghost_pre_extend, h, w), dtype=latent_image.dtype, device=latent_image.device)],
                dim=2,
            )
            noise_mask = torch.cat(
                [noise_mask, torch.ones((mb, mc, ghost_pre_extend, mh, mw), dtype=noise_mask.dtype, device=noise_mask.device)],
                dim=2,
            )
            latent_length = latent_image.shape[2]

        # ── Standard Keyframe Guidance ────────────────────────────────────────
        images       = (guide_data or {}).get("images", [])
        insert_frames = (guide_data or {}).get("insert_frames", [])
        strengths    = (guide_data or {}).get("strengths", [])
        director_fps = float((motion_guide_data or {}).get("frame_rate", (guide_data or {}).get("frame_rate", 24)))
        segments     = (motion_guide_data or {}).get("segments", [])

        if images or segments:
            log.info("[MuseGuide] Appended Keyframe Guidance. is_lora_active: %s", is_lora_active)

            # A. Image guides
            for idx, img_tensor in enumerate(images):
                f_idx    = insert_frames[idx] if idx < len(insert_frames) else 0
                strength = float(strengths[idx] if idx < len(strengths) else 1.0)
                log.info("[MuseGuide] image_guide[%d]: insert_frame=%s strength=%.3f shape=%s mean=%.4f std=%.4f",
                          idx, f_idx, strength, list(img_tensor.shape),
                          float(img_tensor.mean()), float(img_tensor.std()))
                if strength <= 0.0:
                    log.warning("[MuseGuide] image_guide[%d]: SKIPPED — strength <= 0", idx)
                    continue

                B_img, H_img, W_img, C_img = img_tensor.shape
                tpw = int(latent_width * 32)
                tph = int(latent_height * 32)
                if tpw != W_img or tph != H_img:
                    nchw = img_tensor.permute(0, 3, 1, 2)
                    nchw = comfy.utils.common_upscale(nchw, tpw, tph, upscale_method, "disabled")
                    img_tensor = nchw.permute(0, 2, 3, 1)

                image_pixels, guide_latent = nodes_lt.LTXVAddGuide.encode(vae, latent_width, latent_height, img_tensor, scale_factors)
                frame_idx, latent_idx = nodes_lt.LTXVAddGuide.get_latent_index(positive, latent_length, len(image_pixels), int(f_idx), scale_factors)

                if latent_idx >= latent_length:
                    log.warning("[MuseGuide] image_guide[%d]: SKIPPED — latent_idx %d >= latent_length %d "
                                "(guide fell outside the latent — check ghost_pre_extend/padding)",
                                idx, latent_idx, latent_length)
                    continue

                max_f = latent_length - latent_idx
                if guide_latent.shape[2] > max_f:
                    guide_latent = guide_latent[:, :, :max_f]

                tokens_added   = guide_latent.shape[2] * guide_latent.shape[3] * guide_latent.shape[4]
                guide_orig_shape = list(guide_latent.shape[2:])

                positive, negative, latent_image, noise_mask = nodes_lt.LTXVAddGuide.append_keyframe(
                    positive, negative, frame_idx, latent_image, noise_mask, guide_latent, strength, scale_factors
                )
                log.info("[MuseGuide] image_guide[%d]: APPLIED at frame_idx=%s latent_idx=%s guide_latent_shape=%s",
                          idx, frame_idx, latent_idx, list(guide_latent.shape))
                if is_lora_active:
                    positive = _append_guide_attention_entry(positive, tokens_added, guide_orig_shape, image_attention_strength)
                    negative = _append_guide_attention_entry(negative, tokens_added, guide_orig_shape, image_attention_strength)

            # B. Motion video segments
            for seg in segments:
                try:
                    video_file = seg.get("videoFile")
                    if not video_file:
                        continue

                    seg_start      = int(seg.get("start", 0))
                    seg_length     = int(seg.get("length", 1))
                    trim_start     = int(seg.get("trimStart", 0))
                    vid_strength   = float(seg.get("videoStrength", 1.0))
                    vid_attn_str   = float(seg.get("videoAttentionStrength", 0.65))

                    if seg_length <= 0 or vid_strength <= 0.0:
                        continue

                    video_frames = _load_motion_video_frames(video_file, trim_start, seg_length, director_fps, seg.get("resampleMode", "nearest"))
                    keep = ((video_frames.shape[0] - 1) // time_scale_factor) * time_scale_factor + 1
                    video_frames = video_frames[:keep]
                    causal_fix = int(seg_start) == 0 or keep == 1
                    encode_frames = video_frames if causal_fix else torch.cat([video_frames[:1], video_frames], dim=0)

                    _, guide_latent = _encode_video_iclora(
                        vae, latent_width, latent_height, encode_frames, scale_factors,
                        latent_downscale_factor, crop, use_tiled_encode, tile_size, tile_overlap,
                        resize_method=active_resize,
                    )
                    if not causal_fix:
                        guide_latent = guide_latent[:, :, 1:]

                    latent_idx = (seg_start + time_scale_factor - 1) // time_scale_factor if seg_start > 0 else 0
                    if latent_idx >= latent_length:
                        continue

                    if seg_start > 0 and guide_latent.shape[2] > 1:
                        guide_latent = guide_latent[:, :, 1:]
                        seg_start  += time_scale_factor
                        latent_idx += 1
                        if latent_idx >= latent_length:
                            continue

                    max_f = latent_length - latent_idx
                    if guide_latent.shape[2] > max_f:
                        guide_latent = guide_latent[:, :, :max_f]

                    guide_orig_shape = list(guide_latent.shape[2:])
                    Bg, Cg, Fg, Hg, Wg = guide_latent.shape
                    guide_mask = torch.ones((Bg, 1, Fg, Hg, Wg), device=guide_latent.device, dtype=guide_latent.dtype)

                    if seg_start > 0:
                        for i, s in enumerate([0.25, 0.65]):
                            if i < Fg:
                                guide_mask[:, :, i] = 1.0 + vid_strength * (1.0 - s)

                    ldf = int(max(1, round(float(latent_downscale_factor))))
                    if ldf > 1:
                        dilated = _dilate_latent({"samples": guide_latent, "noise_mask": guide_mask}, ldf, ldf)
                        guide_mask   = dilated["noise_mask"]
                        guide_latent = dilated["samples"]

                    tokens_added = guide_latent.shape[2] * guide_latent.shape[3] * guide_latent.shape[4]
                    positive, negative, latent_image, noise_mask = nodes_lt.LTXVAddGuide.append_keyframe(
                        positive, negative, seg_start, latent_image, noise_mask,
                        guide_latent, vid_strength, scale_factors,
                        guide_mask=guide_mask,
                        latent_downscale_factor=float(latent_downscale_factor),
                        causal_fix=causal_fix,
                    )
                    if is_lora_active:
                        positive = _append_guide_attention_entry(positive, tokens_added, guide_orig_shape, vid_attn_str)
                        negative = _append_guide_attention_entry(negative, tokens_added, guide_orig_shape, vid_attn_str)

                except Exception as exc:
                    raise RuntimeError(f"[MuseGuide] Motion segment failed for {seg}: {exc}") from exc

        else:
            log.info("[MuseGuide] No guides present — passing through.")

        crop_frames = max(0, int(latent_image.shape[2]) - initial_latent_length)
        positive = node_helpers.conditioning_set_values(positive, {"nghtdrp_guide_crop_latent_frames": crop_frames})
        negative = node_helpers.conditioning_set_values(negative, {"nghtdrp_guide_crop_latent_frames": crop_frames})
        return (positive, negative, {"samples": latent_image, "noise_mask": noise_mask}, model, float(latent_downscale_factor))

    @staticmethod
    def _lookup_msr_node_classes():
        from nodes import NODE_CLASS_MAPPINGS
        ic_lora_guide_cls = NODE_CLASS_MAPPINGS.get("LTXAddVideoICLoRAGuide")
        add_guide_cls = NODE_CLASS_MAPPINGS.get("LTXVAddGuide")
        if ic_lora_guide_cls is None or add_guide_cls is None:
            raise ValueError(
                "Licon MSR (Prefix) mode needs the 'LTXAddVideoICLoRAGuide' node from the "
                "ComfyUI-LTXVideo custom_nodes package and the core 'LTXVAddGuide' node — "
                "one of them doesn't appear to be installed/loaded."
            )
        return ic_lora_guide_cls, add_guide_cls

    @staticmethod
    def _pad_latent_for_msr_prefix(latent, keyframe_count, prefix_latents):
        """MSR reserves `prefix_latents` latent frames up front for identity
        keyframes; if there are fewer keyframes than that, pad the generated
        region out with denoisable zero/one frames so the total still lines up
        — the downstream crop (stock LTXVCropGuides, which is position-aware
        via conditioning metadata rather than a fixed tail length) needs the
        true clean region to start exactly at `prefix_latents`."""
        pad_count = max(0, prefix_latents - keyframe_count)
        if pad_count <= 0:
            return latent
        samples = latent["samples"]
        mask = _clone_noise_mask(latent, samples)
        b, c, _, h, w = samples.shape
        mb, mc, _, mh, mw = mask.shape
        padded_samples = torch.cat(
            [samples, torch.zeros((b, c, pad_count, h, w), dtype=samples.dtype, device=samples.device)], dim=2,
        )
        padded_mask = torch.cat(
            [mask, torch.ones((mb, mc, pad_count, mh, mw), dtype=mask.dtype, device=mask.device)], dim=2,
        )
        return {"samples": padded_samples, "noise_mask": padded_mask}

    @classmethod
    def _inject_msr(cls, positive, negative, vae, latent, msr, model, latent_downscale_factor, msr_strength=0.0):
        """Licon MSR (Prefix): inject the reference slideshow + per-character
        identity keyframes built into guide_data["msr"], at this node's current
        resolution. Two passes over the fixed reference material:

          1. The full reference slideshow goes through LTXAddVideoICLoRAGuide,
             which is where the actual identity conditioning gets attached —
             only its conditioning output is kept, its own latent is discarded.
          2. Each individual reference image is then dropped in as its own
             frozen keyframe via LTXVAddGuide, one call per reference, each
             call only contributing its updated latent — the conditioning
             stays pinned to whatever step 1 produced, it is not touched again
             inside this loop.

        Returns the node's usual 5-tuple: (positive, negative, latent, model,
        latent_downscale_factor).
        """
        ic_lora_guide_cls, add_guide_cls = cls._lookup_msr_node_classes()

        def clamp01(value):
            return max(0.0, min(1.0, float(value)))

        slideshow = msr["slideshow"]
        keyframes = msr["keyframes"]
        prefix_latents = int(msr["prefix_latents"])
        # A caller-supplied msr_strength > 0 overrides the Director's own value —
        # useful for a refinement stage where the references should hold detail
        # without repainting the opening frames from scratch.
        strength = clamp01(msr_strength) if float(msr_strength) > 0.0 else clamp01(msr["strength"])
        downscale = float(msr["downscale"])

        working_latent = cls._pad_latent_for_msr_prefix(latent, len(keyframes), prefix_latents)

        conditioning_result = _flatten_node_return(_call_node(
            ic_lora_guide_cls,
            positive=positive, negative=negative, vae=vae, latent=working_latent, image=slideshow,
            frame_idx=0, strength=strength, latent_downscale_factor=downscale,
            crop="center", use_tiled_encode=False, tile_size=256, tile_overlap=64,
        ))
        positive, negative = conditioning_result[0], conditioning_result[1]

        for reference_image in keyframes:
            keyframe_result = _flatten_node_return(_call_node(
                add_guide_cls,
                positive=positive, negative=negative, vae=vae, latent=working_latent,
                image=reference_image, frame_idx=0, strength=strength,
            ))
            working_latent = keyframe_result[2]

        log.info("[MuseGuide] Licon MSR: injected %d keyframe(s), prefix_latents=%d, strength=%.2f",
                 len(keyframes), prefix_latents, strength)

        # MSR is a prefix, not a tail — no nghtdrp_guide_crop_latent_frames value is
        # set here (that's a Ghost-Mask-only, fixed-length-tail concept). The
        # IC-LoRA/AddGuide calls above already wrote position-aware keyframe
        # metadata into the conditioning, so whatever calls this must crop with the
        # stock, core LTXVCropGuides node rather than this file's own MuseCropGuides.
        return (positive, negative, working_latent, model, float(latent_downscale_factor))


# ── Crop node ─────────────────────────────────────────────────────────────────

class MuseCropGuides:
    """
    Trims the appended guide keyframes from the latent after sampling.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "positive": ("CONDITIONING",),
                "negative": ("CONDITIONING",),
                "latent":   ("LATENT",),
            }
        }

    RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT")
    RETURN_NAMES = ("positive", "negative", "latent")
    FUNCTION = "execute"
    CATEGORY = "Muse Collective"

    def execute(self, positive, negative, latent):
        latent_image = latent["samples"].clone()
        noise_mask   = _get_noise_mask_for_crop(latent)
        crop_frames  = _get_exact_crop_count(positive)

        if crop_frames > 0:
            crop_frames = min(crop_frames, max(0, latent_image.shape[2] - 1))
            latent_image = latent_image[:, :, :-crop_frames]
            noise_mask   = noise_mask[:, :, :-crop_frames]

        clear = {
            "keyframe_idxs": None,
            "guide_attention_entries": None,
            "nghtdrp_guide_crop_latent_frames": None,
        }
        positive = node_helpers.conditioning_set_values(positive, clear)
        negative = node_helpers.conditioning_set_values(negative, clear)
        return (positive, negative, {"samples": latent_image, "noise_mask": noise_mask})


# ── Registration ──────────────────────────────────────────────────────────────

NODE_CLASS_MAPPINGS = {
    "MuseGuide":      MuseGuide,
    "MuseCropGuides": MuseCropGuides,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "MuseGuide":      "Muse Guide",
    "MuseCropGuides": "Muse Crop Guides",
}