"""Distributed transformer weight tracker with LoRA hot-swap.
Shared infrastructure used by both TDP and SP builders.
"""

from __future__ import annotations

import torch
import torch.distributed as dist

from ltx_core.loader.fuse_loras import fuse_lora_weights
from ltx_core.loader.primitives import LoraPathStrengthAndSDOps, LoraStateDictWithStrength, StateDict
from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder as Builder
from ltx_core.model.model_protocol import ModelType
from ltx_core.multigpu.sharded_sd import ShardedSD


def _apply_loras_inplace(
    source: dict[str, torch.Tensor],
    target: dict[str, torch.Tensor],
    loras: tuple[LoraPathStrengthAndSDOps, ...],
    builder: Builder,  # type: ignore[type-arg]
    lora_keys: frozenset[str],
) -> None:
    """Reset *target* to clean weights from *source*, then fuse all LoRAs in one pass."""
    for key, clean_weight in source.items():
        target[key].copy_(clean_weight)

    lora_sds = [
        LoraStateDictWithStrength(
            builder.load_sd(
                [lora.path],
                sd_ops=lora.sd_ops.with_additional_allowed_keys(lora_keys),
                registry=builder.registry,
                device=builder.lora_load_device,
            ),
            lora.strength,
        )
        for lora in loras
        if lora.strength != 0
    ]
    target_sd = StateDict(
        sd=target, device=next(iter(target.values())).device, size=0, dtype={next(iter(target.values())).dtype}
    )
    for key, fused in fuse_lora_weights(target_sd, lora_sds, fuse_rule=builder.fuse_rule):
        target[key].copy_(fused)


class TransformerWeightTracker:
    """Tracks cached transformer weights with distributed LoRA hot-swap.
    Shared across stage builders that operate on the same checkpoint.
    Does **not** own the model weights — it references tensors stored in a
    :class:`Registry` and receives a builder at :meth:`build` time.
    Uses two :class:`ShardedSD` instances (created on first :meth:`build` call):
    - ``stored_sd`` — cloned backup of the original (pre-LoRA) weights.
      Used to restore registry tensors before applying a different LoRA set.
    - ``broadcast_sd`` — zero-copy view into the registry tensors.
      After in-place LoRA fusion on the owning rank, this broadcasts the
      fused results to all other ranks so every rank sees the same weights.
    Both are created together and are always either both ``None`` or both set.
    With ``no_lora_swap``, the LoRA set is assumed fixed (none, or one set):
    the backup clone is skipped and any swap or reset raises.
    """

    def __init__(self, group: dist.ProcessGroup, bucket_mb: int = 256, no_lora_swap: bool = False) -> None:
        if bucket_mb <= 0:
            raise ValueError("bucket_mb must be > 0")
        self._group = group
        self._bucket_mb = bucket_mb
        self._no_lora_swap = no_lora_swap
        self._staging: torch.Tensor | None = None
        self.stored_sd: ShardedSD | None = None
        self.broadcast_sd: ShardedSD | None = None
        self.loras: list[tuple[str, float, str]] = []

    @property
    def staging(self) -> torch.Tensor:
        """The single broadcast scratch buffer, shared by both SDs, allocated on first use."""
        if self._staging is None:
            device = torch.device(f"cuda:{torch.cuda.current_device()}")
            self._staging = torch.empty(self._bucket_mb * 1024 * 1024, dtype=torch.uint8, device=device)
        return self._staging

    def loras_match(self, lora_list: list[tuple[str, float, str]]) -> bool:
        if len(lora_list) != len(self.loras):
            return False
        return sorted(lora_list) == sorted(self.loras)

    def reset_loras(self, target_sd: dict[str, torch.Tensor]) -> None:
        """Restore *target_sd* to original (pre-LoRA) weights.
        No-op if no LoRAs are currently applied. This is a cooperative
        operation — all ranks must call it simultaneously.
        """
        if not self.loras:
            return
        if self._no_lora_swap:
            raise RuntimeError("no_lora_swap tracker has no backup to reset from")
        if self.stored_sd is None:
            raise RuntimeError("stored_sd must be initialised before reset_loras")
        self.loras = []
        self.stored_sd.broadcast_shards_into(target_sd, self.staging)

    def _local_lora_keys(self) -> frozenset[str]:
        """Derive LoRA key names from the locally owned model keys."""
        if self.stored_sd is None:
            return frozenset()

        keys: set[str] = set()
        for k in self.stored_sd.local_shard:
            if k.endswith(".weight"):
                prefix = k[: -len(".weight")]
                keys.add(f"{prefix}.lora_A.weight")
                keys.add(f"{prefix}.lora_B.weight")
        return frozenset(keys)

    def apply_loras_(
        self,
        target_sd: dict[str, torch.Tensor],
        loras: tuple[LoraPathStrengthAndSDOps, ...],
        builder: Builder,  # type: ignore[type-arg]
    ) -> None:
        """Fuse *loras* into *target_sd* in-place (trailing ``_`` denotes in-place).
        Skips work when the requested LoRAs already match. Restores stored
        weights before applying new LoRAs. This is a cooperative operation —
        all ranks must call it simultaneously.
        """
        new_loras = [(lora.path, lora.strength, lora.sd_ops.name) for lora in loras]

        if self.loras_match(new_loras):
            return

        if self._no_lora_swap and self.loras:
            raise RuntimeError(f"no_lora_swap tracker cannot change LoRAs: have {self.loras}, requested {new_loras}")

        if all(lora.strength == 0 for lora in loras):
            self.reset_loras(target_sd)
            return

        if self.stored_sd is None or self.broadcast_sd is None:
            raise RuntimeError("ShardedSDs must be initialised before apply_loras_ (call build first)")

        source = self.stored_sd.local_shard
        target = {k: v for k, v in target_sd.items() if k in source}
        lora_keys = self._local_lora_keys()
        _apply_loras_inplace(source, target, loras, builder, lora_keys)

        self.broadcast_sd.broadcast_shards_into(target_sd, self.staging)
        self.loras = new_loras

    def build(
        self,
        builder: Builder[ModelType],
        device: torch.device | None = None,
        dtype: torch.dtype | None = None,
        **kwargs: object,
    ) -> ModelType:
        """Build the transformer model with distributed LoRA hot-swap.
        Populates the registry with clean weights on first call, then applies
        LoRAs in-place and broadcasts to all ranks. Assumes the builder carries
        a non-dummy :class:`Registry` so that weights can be cached and reused
        across calls.
        """
        loras = builder.loras
        clean_builder = builder.with_loras(())

        model_paths = list(builder.model_path) if isinstance(builder.model_path, tuple) else [builder.model_path]

        # First call: populate the registry with clean weights.
        if clean_builder.registry.get(model_paths, clean_builder.model_sd_ops) is None:
            clean_builder.build(device=device, dtype=dtype, **kwargs)

        cached_sd = clean_builder.registry.get(model_paths, clean_builder.model_sd_ops)
        if cached_sd is None:
            raise RuntimeError("Expected model state dict in registry but found None")

        if self.stored_sd is None:
            self.stored_sd = ShardedSD.from_state_dict(cached_sd.sd, self._group, clone=not self._no_lora_swap)
            self.broadcast_sd = ShardedSD.from_state_dict(cached_sd.sd, self._group, clone=False)

        if loras:
            self.apply_loras_(cached_sd.sd, loras, builder)
        else:
            self.reset_loras(cached_sd.sd)

        return clean_builder.build(device=device, dtype=dtype, **kwargs)
