"""Backend family registry.

Maps a ``SessionConfig.backend`` family name to a factory that builds
the session's :class:`~acestep.streaming.generator_backend.GeneratorBackend`.
Adding a model family is an entry here plus a backend module — the
session, runner, wire protocol, and UI do not change (see
``round_3_BACKEND_PLAN_FINAL.md``).

Factory contract: ``factory(streaming_session) -> GeneratorBackend``.
The factory pulls whatever it needs off the (fully constructed)
StreamingSession; that keeps this registry free of per-family argument
plumbing.

A registration dict is deliberate (vs entry points / import scanning):
with a handful of in-tree families, an explicit dict is greppable and
import-cheap. Revisit only if out-of-tree families become real.
"""

from acestep.engine.obs import logger


def _make_acestep(ss):
    from acestep.paths import checkpoint_scale
    from acestep.steering import SteeringController, ensure_steering_vectors
    from acestep.streaming.ace_backend import ACEStepBackend

    # SteeringController is the source of truth for slot_count and the
    # vector catalog; ensure_steering_vectors fetches/caches the
    # checkpoint's probe bundle (None for checkpoints without one — XL,
    # fetch failures — which degrades the controller to is_loaded=False
    # and drops the steering capability/knobs for the session).
    steering = SteeringController(ensure_steering_vectors(ss.checkpoint))

    return ACEStepBackend(
        ss.session, ss.stream,
        state=ss.state,
        use_midi=True,  # always "MIDI" mode; KnobState provides values
        use_sde=ss.use_sde, use_lora=ss.use_lora,
        midi_knobs=ss.virtual_knobs,
        engine_obj=ss.engine_obj,
        vae_window=ss.vae_window, crop_seconds=ss.crop_seconds,
        k1_name=ss.k1_name, seed=1528, skip_threshold=5e-4,
        walk_window=ss.walk_window,
        walk_window_s=ss.walk_window_s,
        neg_conditioning=ss.cond_negative,
        steering=steering,
        # Scale label for the lora_compatible predicate; None for
        # checkpoints outside the scale map = "don't filter".
        checkpoint_scale=checkpoint_scale(ss.checkpoint),
    )


def _make_sa3(ss):
    # Assembles SA3Backend.from_context from the construction payload
    # the per-family create path (acestep.streaming.sa3_session) stashed
    # on the session: the process-cached SA3Context, the precomputed
    # conditioning capture, and the SAME-encoded source anchor. Plain
    # session attributes (knob state, steps/depth, vae_window) come off
    # the session itself, per the factory contract above.
    from acestep.streaming.sa3_backend import SA3Backend

    init = getattr(ss, "backend_init", None)
    if not init or "context" not in init:
        raise ValueError(
            "backend 'sa3' requires the per-family create path "
            "(acestep.streaming.sa3_session.create_sa3_session) to stash "
            "its construction payload; an ACE-shaped session cannot "
            "assemble an SA3 backend"
        )
    return SA3Backend.from_context(
        init["context"],
        prompt=ss.state.prompt_text,
        duration_s=float(init["duration_s"]),
        knob_state=ss.virtual_knobs,
        state=ss.state,
        cond=init["cond"],
        # Prompt-B capture for the A/B crossfade; .get so an in-process
        # payload predating the blend surface stays blend-neutral.
        cond_b=init.get("cond_b"),
        source_latent_bct=init["source_latent_bct"],
        # Resolved accel values (compile already normalized to eager by
        # the create path); .get so an in-process payload without them
        # stays on the eager default.
        dit_backend=init.get("dit_backend", "eager"),
        codec_backend=init.get("codec_backend", "eager"),
        steps=int(ss.config.steps),
        depth=int(ss.state.current_depth),
        vae_window_s=float(ss.vae_window),
        # SA3 LoRA (plan Phase 1): the create path constructs the
        # family manager against the process-cached model and stashes
        # it here; .get so an in-process payload predating the LoRA
        # surface stays LoRA-less rather than failing assembly.
        lora_manager=init.get("lora_manager"),
        use_lora=bool(ss.use_lora),
    )


FAMILIES = {
    "acestep": _make_acestep,
    "sa3": _make_sa3,
}

# ---------------------------------------------------------------------------
# Checkpoint aliases (plan §3.5): a server --checkpoint name resolves to
# (backend family, model id). Names absent from the map are ACE
# checkpoint directory names, exactly as before — so plain "xl" and the
# canonical directory names keep working.
# ---------------------------------------------------------------------------

CHECKPOINT_ALIASES = {
    "xl": ("acestep", "acestep-v15-xl-turbo"),
    "sa3-small": ("sa3", "small-music"),
    "sa3-medium": ("sa3", "medium"),
}


def resolve_checkpoint(name: str) -> tuple:
    """``--checkpoint`` name -> ``(backend_family, model_id)``."""
    return CHECKPOINT_ALIASES.get(name, ("acestep", name))


# ---------------------------------------------------------------------------
# Startup-warmup policy (plan §3.5: warmup is backend policy, replacing
# checkpoint-name sniffing in server.py). "ace_trt" drives the synthetic
# ACE warmup session (acestep.streaming.warmup: TRT decoder-engine load,
# LoRA-refit manager, first-tick pipeline build — ~30s of one-time
# engine-resident state). "none" skips it: SA3's one-time cost is the
# SA3Context load, which the per-family create path process-caches, so
# the first real session pays it once and the rest are warm.
# ---------------------------------------------------------------------------

WARMUP_POLICIES = {
    "acestep": "ace_trt",
    "sa3": "none",
}


def warmup_policy(family: str) -> str:
    """Startup-warmup policy for ``family`` ("ace_trt" | "none")."""
    return WARMUP_POLICIES.get(family, "none")


def _acestep_knob_universe():
    from acestep.steering.policy import (
        AUTO_AXES,
        MANUAL_MAX_LAYER,
        MANUAL_MAX_STEP,
        PROBE_N,
    )
    from acestep.streaming.knobs import (
        knob_specs,
        manual_slot_specs,
        steering_axis_spec,
    )

    # Every spec the family can ever expose: both SDE-mode variants plus
    # a representative LoRA-strength knob (the per-id specs all come from
    # lora_strength_spec, so one placeholder id covers the pattern), plus
    # the steering surface — the four auto axes and one representative
    # manual slot (per-slot specs all come from manual_slot_specs).
    # Catalog geometry uses the canonical v15-turbo bundle's 144 cells;
    # no network fetch happens here (policy tables only).
    steering = [
        steering_axis_spec(
            ax.name,
            axis=ax.axis,
            inject_layer=max(
                0, min(MANUAL_MAX_LAYER, ax.probe_layer + ax.layer_offset),
            ),
            probe_step=ax.probe_step,
            probe_n=PROBE_N,
            blurb=ax.blurb,
        )
        for ax in AUTO_AXES
    ] + manual_slot_specs(
        1,
        src_max=143,
        catalog_len=144,
        layer_max=MANUAL_MAX_LAYER,
        step_max=MANUAL_MAX_STEP,
    )
    return (
        knob_specs(False, loras=["<lora_id>"])
        + knob_specs(True, loras=["<lora_id>"])
        + steering
    )


# Per-family knob universes for the cross-backend homonym rule (plan
# §3.3): the full set of KnobSpecs a family can ever expose, obtainable
# WITHOUT constructing the (GPU-heavy) backend. The homonym drift guard
# (tests/unit/test_knob_homonyms.py) runs over these manifests — a knob
# name shared across families must mean exactly the same thing, or it
# must be renamed (prefix / group), so the first lazily-reused name
# can't become a silent semantic fork. Keyed identically to FAMILIES;
# the guard enforces the keys stay in sync.
def _sa3_knob_universe():
    from acestep.streaming.sa3_backend import sa3_knob_specs

    # One representative LoRA-strength knob (the per-id specs all come
    # from the shared lora_strength_spec factory, so one placeholder id
    # covers the pattern — same convention as the ACE universe). The
    # name is shared with ACE's universe deliberately: the homonym
    # guard proves the spec shapes are identical across families.
    return sa3_knob_specs(loras=["<lora_id>"])


FAMILY_KNOB_UNIVERSES = {
    "acestep": _acestep_knob_universe,
    "sa3": _sa3_knob_universe,
}


def _create_sa3_session(cls, **kwargs):
    from acestep.streaming.sa3_session import create_sa3_session

    return create_sa3_session(cls, **kwargs)


# Families whose per-connect setup doesn't fit the ACE create path
# (TRT profiles, model load, demucs, conditioning encode). Keyed like
# FAMILIES; absent = the family rides StreamingSession.create's default
# body. Contract: ``creator(cls, *, audio, config, checkpoint,
# session_id, **rest) -> StreamingSession``.
SESSION_CREATORS = {
    "sa3": _create_sa3_session,
}


def make_backend(name: str, streaming_session):
    """Build the GeneratorBackend for ``name``.

    Unknown families fail loudly at session create (config-time error,
    never a silent fallback): the client asked for a generator this
    server build does not ship.
    """
    try:
        factory = FAMILIES[name]
    except KeyError:
        known = ", ".join(sorted(FAMILIES))
        logger.error("unknown_backend_family name={} known={}", name, known)
        raise ValueError(
            f"unknown backend family {name!r} (registered: {known})"
        ) from None
    return factory(streaming_session)
