"""Context windows: scenes longer than the model can hold at once.

Nothing is ported: ComfyUI's own `comfy.context_windows` owns the whole mechanism. This builds its
handler and hangs it on the model, the way the stock LTXVContextWindows node does. Core picks it up
inside `calc_cond_batch`, which every sampler already goes through.

The capability gate is the important part. The handler has long existed for plain video, but the
packed audio+video latent needs `BaseModel.map_context_window_to_modalities` to unpack the stream,
map each video window onto its audio window and re-slice the guides. Without it, windowing an LTX
latent would cut the packed tensor blindly and quietly wreck the sound and the guides, so that is
refused, in words, rather than tried.

Core's schedule names put the words the other way round from what this once offered
(`standard_uniform`, not `uniform_standard`); the old spellings are still read, from projects that
saved them.
"""

import inspect

from ..._core import control, patching

ID = "context_windows"
TITLE = "Context windows"
MOUNT = "generation.sampling"
STAGE = "sampling"
CATEGORY = "sampling"
STATUS = "experimental"
REQUIRES = ["ltx_av"]

SCHEDULES = ["standard_uniform", "standard_static", "looped_uniform", "batched"]
FUSES = ["pyramid", "relative", "flat", "overlap-linear"]
_ALIASES = {"uniform_standard": "standard_uniform", "static_standard": "standard_static",
            "uniform_looped": "looped_uniform"}
_TEMPORAL = 8                      # LTX: real frames per latent frame

SETTINGS = {
    "enabled": {"type": "bool", "default": False, "label": "Context windows (long scenes)",
                "hint": "Sample a long scene in overlapping windows so it fits the model's memory."},
    "length": {"type": "int", "default": 145, "min": 9, "max": 2049, "step": 8, "label": "Window length (frames)",
               "hint": "Frames the model sees at once. Shorter uses less memory; longer keeps motion steadier.",
               "when": {"enabled": True}},
    "overlap": {"type": "int", "default": 40, "min": 0, "max": 512, "step": 8, "label": "Window overlap (frames)",
                "hint": "Frames neighbouring windows share. More overlap joins smoother but runs slower; it must be shorter than the window.",
                "when": {"enabled": True, "schedule": [v for v in SCHEDULES if v != "batched"]}},
    "schedule": {"type": "enum", "default": "standard_uniform", "label": "Window schedule",
                 "hint": "standard_uniform: windows shift a little each step. standard_static: the same windows every step. looped_uniform: the end wraps round to the start. batched: back-to-back windows with no overlap.",
                 "options": [{"value": v, "label": v} for v in SCHEDULES], "when": {"enabled": True}},
    "fuse": {"type": "enum", "default": "pyramid", "label": "Window blend",
             "hint": "How windows are mixed: pyramid (relative is the same in this ComfyUI) trusts the middle of each window most; overlap-linear fades across the overlap only; flat averages equally.",
             "options": [{"value": v, "label": v} for v in FUSES], "when": {"enabled": True}},
    "freenoise": {"type": "bool", "default": True, "label": "FreeNoise blending", "when": {"enabled": True},
                  "hint": "Shuffle the starting noise so neighbouring windows share it."},
    "retain_first": {"type": "bool", "default": False, "label": "Pin anchor in every window",
                     "hint": "Keeps the first frame in every window so the whole clip stays tied to your starting image.",
                     "when": {"enabled": True}},
}


def _cleanup(patcher, key):
    """A handler left on this model by an earlier install of ours (a clone carries it forward)."""
    options = patcher.model_options
    if patching._ours(options.get("context_handler"), key):
        options.pop("context_handler", None)


def install(patcher, values, key):
    _cleanup(patcher, key)
    if not values.get("enabled"):
        return None
    try:
        import comfy.context_windows as cw
        import comfy.patcher_extension as pe
    except ImportError:
        raise control.Unavailable("this ComfyUI has no context-window support (needs ComfyUI >= v0.29.0)")
    if not hasattr(getattr(patcher, "model", None), "map_context_window_to_modalities"):
        raise control.Unavailable("this ComfyUI cannot window an audio+video latent (it lacks "
                           "map_context_window_to_modalities); update ComfyUI")

    schedule = _ALIASES.get(str(values.get("schedule")), str(values.get("schedule")))
    try:
        sched = cw.get_matching_context_schedule(schedule)
    except ValueError:
        raise control.Unavailable(f"schedule {schedule!r} is not one this ComfyUI knows")
    try:
        fuse = cw.get_matching_fuse_method(str(values.get("fuse")))
    except ValueError:
        raise control.Unavailable(f"blend {values.get('fuse')!r} is not one this ComfyUI knows")

    latent_len = max(((int(values["length"]) - 1) // _TEMPORAL) + 1, 1)
    overlap = max(int(values["overlap"]) // _TEMPORAL, 0)
    clamped = overlap >= latent_len                       # core's schedules loop forever / return no windows
    overlap = min(overlap, latent_len - 1)
    retain = "0" if values.get("retain_first") else ""
    kwargs = dict(context_schedule=sched, fuse_method=fuse, context_length=latent_len,
                  context_overlap=overlap, context_stride=1,
                  closed_loop=False, dim=2, freenoise=bool(values.get("freenoise")),
                  cond_retain_index_list=retain, latent_retain_index_list=retain,
                  split_conds_to_windows=False)
    accepted = set(inspect.signature(cw.IndexListContextHandler).parameters)
    dropped = [k for k in kwargs if k not in accepted]       # this core's handler lacks some keywords
    for k in dropped:
        kwargs.pop(k)
    handler = patching.tag(cw.IndexListContextHandler(**kwargs), key)
    patcher.model_options["context_handler"] = handler
    patcher.add_wrapper_with_key(pe.WrappersMP.PREPARE_SAMPLING, f"{key}.prepare", cw._prepare_sampling_wrapper)
    if values.get("freenoise"):
        patcher.add_wrapper_with_key(pe.WrappersMP.SAMPLER_SAMPLE, f"{key}.freenoise", cw._sampler_sample_wrapper)
    note = f"windows of {values['length']} frames, {values['overlap']} overlap ({schedule}, {values['fuse']})"
    if clamped:
        note = (f"windows of {values['length']} frames, overlap cut to {overlap * _TEMPORAL} "
                f"(it must be shorter than the window) ({schedule}, {values['fuse']})")
    if retain and "latent_retain_index_list" in dropped:
        note += "; this ComfyUI cannot pin the anchor in the latent, only in the conditioning"
    elif retain and "cond_retain_index_list" in dropped:
        note += "; this ComfyUI cannot pin the anchor at all"
    if values.get("freenoise") and "freenoise" in dropped:
        note += "; this ComfyUI has no FreeNoise blending"
    return note


PROVIDES = {"modifier": install}