"""Backend-agnostic LoRA manager.

Owns the catalog, lifecycle (REGISTERED -> MATERIALIZING -> MATERIALIZED
-> ENABLED), background prewarm, and delta math (B @ A) for dynamic LoRA
application. Subclasses plug in *where the live weights live*: a TRT
engine via IRefitter, or a PyTorch decoder's parameters directly.

Library:
  ``register_library(directory)`` discovers ``*.safetensors`` in a flat
  directory and registers each as a REGISTERED entry whose id is the
  filename stem. The catalog backs an "infinite library" workflow where
  hundreds of LoRAs cost nothing until enabled.

Threading:
  - register / enable / disable / set_strength / refit run on the
    inference-owning thread; refit and inference are mutually exclusive.
  - prewarm runs on a background ThreadPoolExecutor; the worker only
    fills the entry's deltas dict and never touches the engine.
"""

from __future__ import annotations

import abc
import concurrent.futures
import time
from dataclasses import dataclass
from enum import Enum
from pathlib import Path
from typing import Dict, List, Optional, Set, Tuple

from loguru import logger
import torch
from safetensors.torch import load_file


class LoRAState(Enum):
    REGISTERED = "registered"
    MATERIALIZING = "materializing"
    MATERIALIZED = "materialized"
    ENABLED = "enabled"


@dataclass
class LoRADescriptor:
    """Read-only public view of a LoRA in the library."""
    id: str
    path: str
    name: str
    state: str
    strength: float
    materialized_bytes: int


@dataclass
class _LoRAEntry:
    """Internal mutable state for one library entry.

    ``strength`` is preserved across enable/disable cycles so flipping a
    UI toggle off and back on doesn't reset the slider.
    """
    lora_id: str
    path: str
    name: str = ""
    state: LoRAState = LoRAState.REGISTERED
    strength: float = 0.0
    deltas: Optional[Dict[str, torch.Tensor]] = None
    future: Optional[concurrent.futures.Future] = None
    materialized_bytes: int = 0


class LoRACatalogBase(abc.ABC):
    """The family-agnostic slice of a LoRA manager: catalog dict,
    REGISTERED → MATERIALIZING → MATERIALIZED (→ ENABLED) state
    machine, background prewarm executor, and read-only inspection.

    Two manager families share this layer without sharing weight
    semantics: the delta-merge managers (:class:`LoRAManagerBase` — ACE
    eager + TRT refit, ``base + Σ strength·delta``) and the SA3
    parametrization manager (:mod:`acestep.engine.sa3_lora`), whose
    materialize/enable transaction is a different shape (staged state
    dicts + torch parametrize registration). Enable/disable/strength
    live on the subclasses — their transaction and rollback semantics
    differ by design.

    Subclasses implement :meth:`_materialize_worker` (worker-thread
    body: fill the entry's staged payload without touching the live
    weights; on failure restore REGISTERED state and re-raise).
    """

    def __init__(self) -> None:
        # Library + lifecycle state.  Insertion order is preserved so
        # ``remove_lora(-1)`` can pop the most-recently-registered entry,
        # matching the legacy stack-style API.
        self._loras: Dict[str, _LoRAEntry] = {}
        self._executor: Optional[concurrent.futures.ThreadPoolExecutor] = None

    # ------------------------------------------------------------------
    # Library: catalog without RAM cost
    # ------------------------------------------------------------------

    @staticmethod
    def _make_id(path: str) -> str:
        """Filename-stem id. Two files with the same stem collide; the
        registrar treats this as identity, so that's fine for a flat
        library directory but means a caller can't register two distinct
        LoRAs that happen to share a stem."""
        return Path(path).stem

    def register_lora(
        self, path: str, name: Optional[str] = None,
        lora_id: Optional[str] = None,
    ) -> str:
        """Add a LoRA to the catalog without materializing deltas.

        Idempotent: re-registering the same id (filename stem) returns
        the existing id and leaves any in-flight prewarm / enabled
        state alone.  The existing entry's name is NOT overwritten on
        re-register; pass an explicit ``name`` only on first registration.

        ``lora_id`` overrides the default filename-stem id; the library
        scan passes pre-disambiguated ids (see
        :func:`acestep.paths.assign_lora_ids`) so same-stem files across
        roots don't silently shadow each other.
        """
        if lora_id is None:
            lora_id = self._make_id(path)
        if lora_id in self._loras:
            existing = self._loras[lora_id]
            if existing.path != str(path):
                logger.warning(
                    "LoRA id {!r} already registered to {}; ignoring re-register from {}",
                    lora_id, existing.path, path,
                )
            return lora_id
        self._loras[lora_id] = _LoRAEntry(
            lora_id=lora_id, path=str(path),
            name=name if name is not None else lora_id,
        )
        logger.info("Registered LoRA: {} (path={})", lora_id, path)
        return lora_id

    def register_library(
        self, directory: Optional[Path] = None,
    ) -> List[str]:
        """Discover and register every ``*.safetensors`` reachable from
        the configured LoRA roots.

        When ``directory`` is supplied, scans only that one root
        (recursive). When omitted, scans :func:`acestep.paths.loras_dir`
        plus everything listed in ``ACESTEP_EXTRA_LORA_DIRS`` via
        :func:`acestep.paths.discover_all_loras`, so operators can
        point at training output directories without moving files into
        the primary library.

        Returns the list of registered ids in scan order (primary root
        first, then each extra dir; each root sorted by filename).
        Missing directories are silently skipped.
        """
        from acestep.paths import assign_lora_ids, discover_all_loras, discover_loras
        if directory is None:
            files = discover_all_loras()
        else:
            files = discover_loras(directory)
        ids: List[str] = []
        for p, lora_id in assign_lora_ids(files):
            try:
                ids.append(self.register_lora(str(p), lora_id=lora_id))
            except Exception as e:
                logger.warning("Failed to register {}: {}", p, e)
        if files:
            logger.info(
                "Registered library: {} LoRAs across {} root(s)",
                len(ids),
                1 if directory is not None else "all configured",
            )
        return ids

    # ------------------------------------------------------------------
    # Lifecycle: REGISTERED <-> MATERIALIZING <-> MATERIALIZED <-> ENABLED
    # ------------------------------------------------------------------

    def _get_executor(self) -> concurrent.futures.ThreadPoolExecutor:
        """Lazy-init single-worker pool for prewarm.  Single worker
        because the materialization runs (B @ A) on the same CUDA device
        as inference; oversubscribing would just block on the GPU."""
        if self._executor is None:
            self._executor = concurrent.futures.ThreadPoolExecutor(
                max_workers=1, thread_name_prefix="lora_prewarm",
            )
        return self._executor

    def prewarm_lora(self, lora_id: str) -> concurrent.futures.Future:
        """Kick off background materialization of ``lora_id``.

        Returns a Future that resolves to ``None`` once deltas are in
        CPU RAM.  Subsequent ``enable_lora`` will skip the materialization
        step.  Calling on a MATERIALIZED / ENABLED entry returns an
        already-completed future; calling on a MATERIALIZING entry
        returns the in-flight future.
        """
        entry = self._require_entry(lora_id)
        if entry.state in (LoRAState.MATERIALIZED, LoRAState.ENABLED):
            f: concurrent.futures.Future = concurrent.futures.Future()
            f.set_result(None)
            return f
        if entry.state == LoRAState.MATERIALIZING:
            assert entry.future is not None
            return entry.future
        entry.state = LoRAState.MATERIALIZING
        entry.future = self._get_executor().submit(
            self._materialize_worker, entry,
        )
        logger.info("Prewarming LoRA: {}", lora_id)
        return entry.future

    @abc.abstractmethod
    def _materialize_worker(self, entry: _LoRAEntry) -> None:
        """Worker-thread body: load + stage the entry's payload without
        touching the live weights, then flip it to MATERIALIZED. On
        failure restore REGISTERED (when still MATERIALIZING) and
        re-raise. If the entry was concurrently disabled (state changed
        away from MATERIALIZING), the result must be dropped rather
        than resurrected."""

    # ------------------------------------------------------------------
    # Inspection
    # ------------------------------------------------------------------

    def list_loras(self) -> List[LoRADescriptor]:
        return [
            LoRADescriptor(
                id=e.lora_id, path=e.path, name=e.name,
                state=e.state.value, strength=e.strength,
                materialized_bytes=e.materialized_bytes,
            )
            for e in self._loras.values()
        ]

    def get_lora(self, lora_id: str) -> LoRADescriptor:
        e = self._require_entry(lora_id)
        return LoRADescriptor(
            id=e.lora_id, path=e.path, name=e.name,
            state=e.state.value, strength=e.strength,
            materialized_bytes=e.materialized_bytes,
        )

    def _require_entry(self, lora_id: str) -> _LoRAEntry:
        if lora_id not in self._loras:
            raise ValueError(f"LoRA {lora_id!r} not registered")
        return self._loras[lora_id]

    @property
    def has_active_loras(self) -> bool:
        return any(e.state == LoRAState.ENABLED for e in self._loras.values())

    @property
    def active_lora_count(self) -> int:
        return sum(
            1 for e in self._loras.values() if e.state == LoRAState.ENABLED
        )

    @property
    def active_lora_ids(self) -> List[str]:
        return [
            e.lora_id for e in self._loras.values()
            if e.state == LoRAState.ENABLED
        ]

    @property
    def total_materialized_bytes(self) -> int:
        return sum(e.materialized_bytes for e in self._loras.values())


class LoRAManagerBase(LoRACatalogBase):
    """Catalog + lifecycle + delta math; subclasses do the engine writeback.

    The delta-merge manager family (ACE lineage): materialization
    computes full-rank ``B @ A`` deltas and enable/disable/strength
    recompose ``base + Σ strength·delta`` into the live weights.

    Subclasses must:

    - Populate ``self._base_weights`` (dict: param_name -> CPU/GPU tensor in
      the live weight's native dtype) and ``self._param_dtype`` (dict:
      param_name -> torch.dtype) during __init__, then call
      ``super().__init__()`` to set up the lifecycle bookkeeping.
    - Implement :meth:`_apply_to_engine`, which writes
      ``base + sum_over_enabled(strength * delta)`` to the live weights
      for the supplied ``param_names``.

    Param naming is decoder-relative (matches ``decoder.named_parameters()``).
    The TRT subclass adds the ``decoder.`` prefix only at the refitter
    boundary.
    """

    def __init__(self) -> None:
        # Subclasses set _base_weights, _param_dtype, and _param_numel
        # before super().__init__. _param_numel is the element count of
        # the base weight (orientation-independent) so _compute_deltas
        # can sanity-check that a LoRA was trained for THIS base model
        # rather than silently producing wrong-shape deltas (e.g. a 2B
        # LoRA materialized against an XL engine).
        if not hasattr(self, "_base_weights"):
            self._base_weights: Dict[str, torch.Tensor] = {}
        if not hasattr(self, "_param_dtype"):
            self._param_dtype: Dict[str, torch.dtype] = {}
        if not hasattr(self, "_param_numel"):
            self._param_numel: Dict[str, int] = {}

        super().__init__()
        self._ever_dirty: Set[str] = set()

    # ------------------------------------------------------------------
    # Subclass hook
    # ------------------------------------------------------------------

    @abc.abstractmethod
    def _apply_to_engine(self, param_names: Set[str]) -> None:
        """Write ``base + Σ strength·delta`` to the live weights.

        ``param_names`` is the set of decoder-relative param names that
        need updating. Subclasses iterate, accumulate enabled-LoRA
        contributions in a buffer, and push to the live weights.
        """

    # ------------------------------------------------------------------
    # Lifecycle transition hooks (default: no-ops)
    #
    # The base lifecycle (enable_lora / disable_lora) calls these around
    # state transitions so subclasses can attach side effects without
    # overriding the lifecycle proper. The motivating case is the eager
    # backend's "promote materialized deltas to GPU on enable, drop the
    # GPU mirror on disable" pattern: deltas live cheaply in CPU RAM
    # while merely MATERIALIZED, but the active set lives in VRAM so
    # slider-driven refits stay zero-copy on the device.
    # ------------------------------------------------------------------

    def _on_enabled(self, entry: "_LoRAEntry") -> None:
        """Hook: ``entry`` just transitioned to ENABLED. Called BEFORE
        the contributing-strength refit fires so subclasses can stage
        runtime data the refit will read."""

    def _on_disabled(self, entry: "_LoRAEntry") -> None:
        """Hook: ``entry`` just transitioned away from ENABLED. Called
        AFTER ``entry.deltas`` is cleared but BEFORE the rollback refit
        fires (so subclasses can drop their mirrors first)."""

    def _materialize_worker(self, entry: _LoRAEntry) -> None:
        """Worker-thread body.  Loads safetensors, computes deltas,
        writes them to the entry.  Engine state untouched.

        If the entry was concurrently disabled (state changed away from
        MATERIALIZING), the result is dropped to avoid resurrecting the
        deltas after disable cleared them.
        """
        t0 = time.perf_counter()
        try:
            deltas, bytes_ = self._compute_deltas(entry.path)
        except Exception:
            if entry.state == LoRAState.MATERIALIZING:
                entry.state = LoRAState.REGISTERED
                entry.future = None
            raise
        if entry.state == LoRAState.MATERIALIZING:
            entry.deltas = deltas
            entry.materialized_bytes = bytes_
            entry.state = LoRAState.MATERIALIZED
            elapsed = (time.perf_counter() - t0) * 1000
            logger.info(
                "Materialized LoRA {} ({} params, {:.1f} MB) in {:.1f}ms",
                entry.lora_id, len(deltas), bytes_ / 1e6, elapsed,
            )

    def _compute_deltas(
        self, lora_path: str,
    ) -> Tuple[Dict[str, torch.Tensor], int]:
        """Load LoRA from disk and compute full-rank deltas (B @ A).

        Pure compute; safe to call from a worker thread.  Returns
        (deltas_dict, total_bytes_in_cpu_ram). Each delta is stored in
        the live weight's native dtype on CPU. Subclasses pull deltas
        through to whatever device they need.
        """
        raw = load_file(lora_path)
        pairs: Dict[str, Dict[str, torch.Tensor]] = {}
        for key, tensor in raw.items():
            parts = key.replace("base_model.model.", "")
            if ".lora_A.weight" in parts:
                param_name = parts.replace(".lora_A.weight", ".weight")
                pairs.setdefault(param_name, {})["A"] = tensor
            elif ".lora_B.weight" in parts:
                param_name = parts.replace(".lora_B.weight", ".weight")
                pairs.setdefault(param_name, {})["B"] = tensor

        # Zero PEFT-pattern matches means this is not an ACE-Step LoRA
        # at all. Without this check the pair map is empty, the
        # shape-mismatch guard below never trips (it only fires on
        # matched pairs), and enable_lora "succeeds" having applied
        # nothing — a silent no-op. Wrong-family files must fail loudly.
        if not pairs:
            sa3_like = any(
                ".parametrizations.weight." in k for k in raw
            )
            hint = (
                " The key layout matches the SA3 (stable-audio-3) "
                "parametrization format; SA3-family LoRAs cannot load "
                "on an ACE-Step engine."
                if sa3_like else ""
            )
            raise RuntimeError(
                f"LoRA {Path(lora_path).name} is not an ACE-Step LoRA: "
                f"none of its {len(raw)} tensors match the PEFT "
                f"lora_A.weight/lora_B.weight key pattern, so there is "
                f"nothing to apply.{hint}"
            )

        device = self._delta_compute_device()
        deltas: Dict[str, torch.Tensor] = {}
        total_bytes = 0
        skipped = 0
        shape_mismatch = 0
        first_mismatch: Optional[tuple] = None
        for param_name, ab in pairs.items():
            if "A" not in ab or "B" not in ab:
                continue
            if param_name not in self._param_dtype:
                skipped += 1
                continue
            target_dt = self._param_dtype[param_name]
            try:
                A = ab["A"].to(device=device, dtype=torch.float32)
                B = ab["B"].to(device=device, dtype=torch.float32)
                d = (B @ A).to(dtype=target_dt)
            except torch.OutOfMemoryError:
                # The GPU matmul is a speed nicety; on a VRAM-pressured
                # pod (TRT workspaces + stems + engine swaps) it can OOM
                # even though the deltas themselves live on CPU. Fall
                # back to CPU compute for the rest of this LoRA instead
                # of failing the enable.
                if device.type == "cpu":
                    raise
                logger.warning(
                    "_compute_deltas({}): GPU OOM during delta matmul; "
                    "falling back to CPU compute",
                    Path(lora_path).name,
                )
                device = torch.device("cpu")
                torch.cuda.empty_cache()
                A = ab["A"].to(device=device, dtype=torch.float32)
                B = ab["B"].to(device=device, dtype=torch.float32)
                d = (B @ A).to(dtype=target_dt)
            d = d.to(device=self._delta_storage_device()).contiguous()
            # Catch base-model mismatches (e.g. 2B LoRA on XL engine):
            # element count is invariant under transpose, so a numel
            # mismatch here is unambiguous. We don't compare full
            # shape because the engine stores some weights transposed
            # vs the LoRA's torch [out, in] convention.
            expected_numel = self._param_numel.get(param_name)
            if expected_numel is not None and d.numel() != expected_numel:
                shape_mismatch += 1
                if first_mismatch is None:
                    first_mismatch = (
                        param_name, tuple(d.shape), d.numel(), expected_numel,
                    )
                continue
            deltas[param_name] = d
            total_bytes += d.numel() * d.element_size()
        # When deltas are stored on CPU (TRT path) the GPU was only a
        # matmul scratchpad; hand its pooled temporaries back to the
        # driver so they don't sit reserved against the next TRT
        # workspace or stem-extraction allocation.
        if torch.cuda.is_available() and self._delta_storage_device().type == "cpu":
            torch.cuda.empty_cache()
        if skipped:
            logger.debug(
                "_compute_deltas({}): {} params skipped (not in engine)",
                Path(lora_path).name, skipped,
            )
        # If every applicable LoRA tensor mismatches the base, this is a
        # base-model mismatch — surface a clear error rather than silently
        # returning empty deltas (which would let enable_lora succeed but
        # do nothing) or crashing later inside _apply_to_engine's add_.
        if shape_mismatch and not deltas:
            assert first_mismatch is not None
            bad_param, lora_shape, lora_numel, base_numel = first_mismatch
            raise RuntimeError(
                f"LoRA {Path(lora_path).name} is incompatible with the "
                f"loaded base model: {shape_mismatch} tensor(s) have "
                f"shapes that don't fit any engine weight slot. "
                f"E.g. {bad_param!r}: LoRA produced shape {lora_shape} "
                f"(numel={lora_numel}) but the base weight has numel="
                f"{base_numel}. Common cause: a 2B LoRA was loaded "
                f"against an XL engine (hidden=2048 vs 4096) or "
                f"vice-versa. Verify the LoRA's training base checkpoint."
            )
        if shape_mismatch:
            logger.warning(
                "_compute_deltas({}): {} tensors had shape mismatches "
                "vs base; {} usable deltas remain. Partial overlap likely "
                "means the LoRA targets a subset that happens to fit.",
                Path(lora_path).name, shape_mismatch, len(deltas),
            )
        return deltas, total_bytes

    def _delta_compute_device(self) -> torch.device:
        """Where the (B @ A) matmul runs. GPU when available."""
        return torch.device(
            "cuda" if torch.cuda.is_available() else "cpu"
        )

    def _delta_storage_device(self) -> torch.device:
        """Where deltas live in RAM after materialization.

        TRT keeps them on CPU because refit buffers and numpy views are
        CPU. Eager keeps them on the decoder's device so refits don't
        pay an H2D transfer per slider tick.
        """
        return torch.device("cpu")

    def enable_lora(
        self, lora_id: str, strength: Optional[float] = None,
    ) -> None:
        """Promote a LoRA to ENABLED.  Synchronous; materializes if needed.

        ``strength``, when provided, overrides the entry's stored strength
        BEFORE the refit fires.  Pass it whenever you know the target
        strength up-front: it lets a one-shot caller atomically transition
        REGISTERED -> ENABLED-at-S without an intermediate refit at 0
        followed by a second refit at S, which the streaming pipeline can
        observe as a one-tick "missing LoRA" glitch in the first decode
        window.

        Refits the engine weights iff the resulting strength is non-zero
        (a strength-0 enable is a placeholder for a slider that hasn't
        ramped yet, and the refit would be a no-op).
        """
        entry = self._require_entry(lora_id)
        if strength is not None:
            entry.strength = float(strength)
        if entry.state == LoRAState.ENABLED:
            return
        if entry.state == LoRAState.MATERIALIZING:
            assert entry.future is not None
            entry.future.result()
        if entry.state == LoRAState.REGISTERED:
            t0 = time.perf_counter()
            deltas, bytes_ = self._compute_deltas(entry.path)
            entry.deltas = deltas
            entry.materialized_bytes = bytes_
            entry.state = LoRAState.MATERIALIZED
            elapsed = (time.perf_counter() - t0) * 1000
            logger.info(
                "Materialized LoRA {} inline ({} params, {:.1f} MB) in {:.1f}ms",
                lora_id, len(deltas), bytes_ / 1e6, elapsed,
            )
        entry.state = LoRAState.ENABLED
        entry.future = None
        # Subclass hook: stage backend-specific runtime state (e.g. the
        # eager backend's CPU->GPU delta promotion) BEFORE the refit so
        # the refit can read it.
        self._on_enabled(entry)
        if entry.strength != 0.0 and entry.deltas:
            self._refit_weights(set(entry.deltas.keys()))
        logger.info(
            "Enabled LoRA {} ({} params, {:.1f} MB, strength={:.2f})",
            lora_id, len(entry.deltas or {}),
            entry.materialized_bytes / 1e6, entry.strength,
        )

    def disable_lora(self, lora_id: str) -> None:
        """Drop deltas from CPU RAM and refit if the LoRA was contributing.

        Strength is preserved on the entry so re-enable returns to the
        same slider position.
        """
        entry = self._require_entry(lora_id)
        if entry.state == LoRAState.REGISTERED:
            return

        if entry.state == LoRAState.MATERIALIZING and entry.future is not None:
            try:
                entry.future.result()
            except Exception:
                pass

        was_contributing = (
            entry.state == LoRAState.ENABLED and entry.strength != 0.0
        )
        affected_params: Set[str] = (
            set(entry.deltas.keys()) if entry.deltas else set()
        )

        entry.state = LoRAState.REGISTERED
        entry.deltas = None
        entry.materialized_bytes = 0
        entry.future = None

        # Subclass hook: drop backend-specific runtime state (e.g. the
        # eager backend's GPU delta mirror). Runs before the rollback
        # refit so the composing code never sees a stale mirror.
        self._on_disabled(entry)

        if was_contributing:
            self._refit_weights(affected_params)

        # Return the just-freed GPU delta mirror to the CUDA driver.
        # _on_disabled drops the only Python references to the mirror
        # tensors, but PyTorch's caching allocator retains the
        # underlying VRAM segments in its pool to amortize future
        # alloc cost. For LoRA disables that's the wrong trade:
        # nvidia-smi still reports the memory as "used" (e.g. a
        # 1.2 GB delta footprint sits resident), and worse, the pool
        # fragments across a session of mixed-size enable/disable
        # cycles. The fragmentation surface manifests on a fleet pod
        # as a mid-session OOM after a few LoRA rotations:
        #
        #   [Server] WS dispatch error: CUDA out of memory.
        #     Tried to allocate 32 MiB. GPU 0 has 13.25 MiB free.
        #     process using 31.31 GiB / 31.36 GiB.
        #     PyTorch allocated 12.98 GiB, 435.36 MiB reserved but
        #     unallocated.  ← the cached pool we want to reclaim
        #
        # The PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True env
        # also avoids fragmentation but is global-process and
        # opt-in. Calling empty_cache() here is precise (only fires
        # on the actual free event), cheap (microseconds-to-low-ms
        # since refs are already dropped), and doesn't require an
        # operator action. Cap PR (#140 in demon-public-demo)
        # bounds concurrent enables; this commit handles the
        # cumulative-rotation pattern that the cap alone can't
        # prevent.
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

        logger.info(
            "Disabled LoRA {} (was_contributing={})",
            lora_id, was_contributing,
        )

    def set_lora_strength(self, lora_id: str, strength: float) -> None:
        """Adjust the strength of an ENABLED LoRA.

        Raises ValueError if the LoRA is not enabled.  Auto-enable was
        considered and rejected: it would hide materialization cost
        (hundreds of ms) behind a slider event.  The UI is responsible
        for calling ``enable_lora`` before a slider becomes interactive.
        """
        entry = self._require_entry(lora_id)
        if entry.state != LoRAState.ENABLED:
            raise ValueError(
                f"LoRA {lora_id!r} is not enabled (state={entry.state.value}). "
                "Call enable_lora() first."
            )
        if entry.strength == strength:
            return
        old = entry.strength
        entry.strength = strength
        if entry.deltas:
            self._refit_weights(set(entry.deltas.keys()))
        logger.info(
            "LoRA {} strength: {:.3f} -> {:.3f} ({} params)",
            lora_id, old, strength, len(entry.deltas or {}),
        )

    def remove_lora(self, lora_id=-1) -> bool:
        """Drop a LoRA from the catalog entirely.

        Disables first if enabled.  Default ``-1`` removes the most
        recently registered entry, preserving the legacy stack-pop API.
        """
        if lora_id == -1:
            if not self._loras:
                return False
            lora_id = next(reversed(self._loras))
        if lora_id not in self._loras:
            return False
        self.disable_lora(lora_id)
        del self._loras[lora_id]
        logger.info("Removed LoRA {} from catalog", lora_id)
        return True

    def remove_all(self) -> None:
        """Remove every LoRA from the catalog and restore engine to base."""
        for lid in list(self._loras.keys()):
            self.remove_lora(lid)

    def close(self) -> None:
        """Drop catalog state and shut down the prewarm executor.

        Called by :meth:`DiffusionEngine.close` on session teardown so the
        materialized deltas (CPU RAM) and any backend-specific runtime
        mirrors (GPU, populated by ``_on_enabled``) don't outlive the
        session that allocated them. The engine refit itself is *not*
        rolled back here — the engine is being destroyed seconds later
        and the refit's only effect would be to dirty pages we're about
        to free.

        Idempotent: subsequent calls are no-ops.
        """
        # Drop CPU deltas + per-id mirrors. Subclasses override
        # ``_on_disabled`` to release backend-specific GPU buffers
        # (EagerLoRAManager._gpu_deltas), but we mark every entry as
        # REGISTERED first so that hook fires once per ever-enabled LoRA.
        for entry in self._loras.values():
            if entry.state == LoRAState.ENABLED:
                entry.state = LoRAState.REGISTERED
                entry.deltas = None
                entry.materialized_bytes = 0
                try:
                    self._on_disabled(entry)
                except Exception as e:
                    logger.warning(
                        "_on_disabled raised for {} during close: {}",
                        entry.lora_id, e,
                    )
            else:
                entry.deltas = None
                entry.materialized_bytes = 0
        self._loras.clear()
        # Drop base-weight snapshots (TRT: CPU, Eager: same dtype/device
        # as live params — GPU on a normal session). Refit buffers
        # (TRT only) live alongside.
        self._base_weights.clear()
        for attr in ("_refit_bufs", "_param_to_trt", "_np_dtype",
                     "_param_dtype", "_decoder_params", "_gpu_deltas"):
            d = getattr(self, attr, None)
            if isinstance(d, dict):
                d.clear()
        # Shut down the prewarm thread pool. A worker thread holding a
        # reference to a partially-materialized entry's deltas would
        # otherwise keep CPU buffers alive past close().
        if self._executor is not None:
            self._executor.shutdown(wait=False, cancel_futures=True)
            self._executor = None
        # Drop TRT-specific bookkeeping last so the IRefitter's reference
        # to the engine clears before the engine is destroyed.
        for attr in ("_refitter", "_engine", "_trt", "_trt_logger"):
            if hasattr(self, attr):
                setattr(self, attr, None)

    # ------------------------------------------------------------------
    # Backward-compat one-shot API
    # ------------------------------------------------------------------

    def apply_lora(self, lora_path: str, strength: float = 1.0) -> str:
        """Register, set strength, and enable a LoRA in one call.

        Idempotent on path: calling twice is the same as registering
        once and setting the second-call's strength.
        """
        lora_id = self.register_lora(lora_path)
        self._loras[lora_id].strength = float(strength)
        self.enable_lora(lora_id)
        return lora_id

    # ------------------------------------------------------------------
    # Inspection (delta-specific; catalog inspection lives on the base)
    # ------------------------------------------------------------------

    @property
    def refittable_param_count(self) -> int:
        return len(self._param_dtype)

    # ------------------------------------------------------------------
    # Refit driver (subclass writeback)
    # ------------------------------------------------------------------

    def _refit_weights(self, param_names: Set[str]) -> None:
        """Time + log the engine writeback. Subclass does the actual work.

        Splitting time/log out of ``_apply_to_engine`` keeps the subclass
        method narrow and lets the TRT path reuse pre-allocated refit
        buffers without bookkeeping noise.
        """
        if not param_names:
            return
        t0 = time.perf_counter()
        self._apply_to_engine(param_names)
        for name in param_names:
            self._ever_dirty.add(name)
        elapsed = (time.perf_counter() - t0) * 1000
        logger.info(
            "Refitted {} weights in {:.1f}ms", len(param_names), elapsed,
        )


class EagerLoRAManager(LoRAManagerBase):
    """LoRA writeback into a PyTorch decoder's parameters in place.

    Storage strategy ("hybrid"):

    - ``MATERIALIZED`` deltas live on **CPU**, matching the TRT backend.
      A library workflow that prewarms dozens of LoRAs costs system
      RAM, not VRAM.
    - On the ENABLED transition, the entry's deltas are mirrored into a
      per-id **GPU** buffer (``self._gpu_deltas[id]``). Refits read from
      that mirror, so slider-driven strength updates run zero-copy on
      the device. The mirror is freed on disable.
    - The base-weight snapshot is also lazy: nothing is cloned at
      construction. The first refit clones the live params it's about
      to overwrite into ``self._base_weights``, so a session that never
      enables a LoRA pays no extra VRAM.

    All writes use ``param.data.copy_`` / ``.add_`` so the live
    ``nn.Parameter`` identity stays the same — required for safety
    under ``torch.compile``: the compiled graph reads params as tensor
    inputs (not constants), and storage-level mutations are visible to
    the next forward without retracing.
    """

    def __init__(
        self,
        decoder: torch.nn.Module,
        device: Optional[torch.device] = None,
    ):
        if decoder is None:
            raise ValueError("EagerLoRAManager requires a decoder module")

        # If the decoder has been wrapped by ``torch.compile``, the live
        # parameters live on the original module; named_parameters()
        # routes through the OptimizedModule wrapper transparently, so
        # this works for both compiled and eager decoders.
        named = list(decoder.named_parameters())
        if not named:
            raise RuntimeError(
                "Decoder has no parameters; cannot init EagerLoRAManager. "
                "(If running with skip_decoder=True, use the TRT path.)"
            )

        # Resolve device: caller-provided wins, else infer from the first
        # parameter. We tolerate a meta-device decoder by deferring to CPU.
        first_param = named[0][1]
        inferred = first_param.device
        self._device = device if device is not None else (
            inferred if inferred.type != "meta" else torch.device("cpu")
        )

        self._decoder_params: Dict[str, torch.nn.Parameter] = {}
        # Populated lazily on first refit (see _ensure_base_snapshot).
        self._base_weights: Dict[str, torch.Tensor] = {}
        self._param_dtype: Dict[str, torch.dtype] = {}
        # Per-id GPU mirror of MATERIALIZED CPU deltas, populated on
        # _on_enabled and dropped on _on_disabled.
        self._gpu_deltas: Dict[str, Dict[str, torch.Tensor]] = {}

        self._param_numel: Dict[str, int] = {}
        for name, param in named:
            self._decoder_params[name] = param
            self._param_dtype[name] = param.dtype
            self._param_numel[name] = param.numel()

        logger.info(
            "Eager LoRA manager ready: {} decoder params indexed on {} "
            "(base snapshot deferred until first enable)",
            len(self._param_dtype), self._device,
        )

        super().__init__()

    def _delta_compute_device(self) -> torch.device:
        return self._device

    # _delta_storage_device falls back to base default (CPU): MATERIALIZED
    # deltas live in system RAM. The hot, GPU-resident mirror is only
    # populated for ENABLED entries via _on_enabled.

    def _on_enabled(self, entry: "_LoRAEntry") -> None:
        """Promote CPU deltas to a GPU mirror so refits stay zero-copy.

        Pays one H2D transfer at enable time; subsequent refits (slider
        ticks, disable rollbacks) read straight from the device mirror.
        """
        if not entry.deltas:
            return
        self._gpu_deltas[entry.lora_id] = {
            name: d.to(device=self._device, non_blocking=True).contiguous()
            for name, d in entry.deltas.items()
        }

    def _on_disabled(self, entry: "_LoRAEntry") -> None:
        """Drop the GPU mirror; the CPU MATERIALIZED copy was already
        cleared by the lifecycle. Frees VRAM equal to the LoRA's delta
        footprint."""
        self._gpu_deltas.pop(entry.lora_id, None)

    def _ensure_base_snapshot(self, param_names: Set[str]) -> None:
        """Lazily clone live params into ``_base_weights`` for the
        requested set.

        Called on the refit hot path. The clone happens at most once
        per param: after first capture, subsequent refits reuse the
        snapshot to recompose ``base + Σ s·d``. A session that only
        ever sits at strength 0 (placeholder pattern) never triggers a
        snapshot because no refit fires.
        """
        for name in param_names:
            if name in self._base_weights:
                continue
            param = self._decoder_params.get(name)
            if param is None:
                continue
            self._base_weights[name] = param.data.detach().clone()

    def _apply_to_engine(self, param_names: Set[str]) -> None:
        """Compose ``base + Σ strength·delta`` directly into param.data.

        Reads deltas from ``self._gpu_deltas`` (the on-device mirror),
        not from ``entry.deltas`` (which live on CPU). The mirror is
        guaranteed populated for every ENABLED entry by ``_on_enabled``.
        """
        self._ensure_base_snapshot(param_names)

        for param_name in param_names:
            param = self._decoder_params.get(param_name)
            if param is None:
                continue
            param.data.copy_(self._base_weights[param_name])

            for entry in self._loras.values():
                if entry.state != LoRAState.ENABLED:
                    continue
                if entry.strength == 0.0:
                    continue
                gpu = self._gpu_deltas.get(entry.lora_id)
                if gpu and param_name in gpu:
                    param.data.add_(gpu[param_name], alpha=entry.strength)
