"""Instance-scoped dispatch for the long-context KDA training engine."""

from __future__ import annotations

from functools import partial, wraps
from threading import RLock

from .training_engine import TrainingEngineConfig


_KDA_DISPATCH_LOCK = RLock()


def configure_kda_mixer_runtime(mixer, engine: TrainingEngineConfig):
    """Attach the configured KDA call to one mixer without global leakage.

    FLA's ``KimiDeltaAttention.forward`` resolves ``chunk_kda`` from its module
    global at call time.  The research trainer replaced that global once per
    process, which works for a one-model executable but leaks when a process holds
    models with different engines.  This wrapper swaps it only for this bound
    forward and restores it in ``finally``; the custom autograd function owns its
    backward, so the global is not needed after forward returns.
    """

    if not isinstance(engine, TrainingEngineConfig):
        raise TypeError("KDA runtime requires TrainingEngineConfig")
    segment_size = int(engine.recurrent_segment_tokens)
    metadata = {
        "format": "yunmo-next-kda-training-engine-v1",
        "installed": False,
        "segment_size": segment_size,
        "allow_packed_documents": bool(engine.recurrent_packed_documents),
        "official_forward_serial_backward": bool(
            engine.recurrent_official_forward_serial_backward
        ),
    }
    if segment_size == 0:
        mixer._yunmo_kda_training_engine = metadata
        return mixer
    if not (
        engine.recurrent_packed_documents
        and engine.recurrent_official_forward_serial_backward
    ):
        raise RuntimeError("invalid serial recurrent execution configuration")

    import fla.layers.kda as fla_kda_layer

    from .kda_serial import serial_chunk_kda

    engine = partial(
        serial_chunk_kda,
        segment_size=segment_size,
        chunk_size=64,
        allow_packed_documents=True,
        official_forward_serial_backward=True,
    )
    original_forward = mixer.forward

    @wraps(original_forward)
    def forward(*args, **kwargs):
        with _KDA_DISPATCH_LOCK:
            previous = fla_kda_layer.chunk_kda
            fla_kda_layer.chunk_kda = engine
            try:
                return original_forward(*args, **kwargs)
            finally:
                fla_kda_layer.chunk_kda = previous

    mixer.forward = forward
    metadata = {**metadata, "installed": True, "dispatch": "instance_scoped"}
    mixer._yunmo_kda_training_engine = metadata
    return mixer