# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

"""Construction and execution of a fixed LTX refinement-phase adapter.

Checkpoint key parsing remains in :mod:`ltx2_adapter_parser`; this module
resolves the declared artifact and makes the existing Transformer linears
adapter-capable while retaining their quantization and tensor-parallel paths.
"""

from __future__ import annotations

from collections import defaultdict
from collections.abc import Iterable
from dataclasses import dataclass
from typing import Any

import torch
import torch.nn as nn
import torch.nn.functional as F
import vllm.envs as envs
from vllm.distributed import (
    split_tensor_along_last_dim,
    tensor_model_parallel_all_reduce,
)
from vllm.logger import init_logger
from vllm.model_executor.layers.batch_invariant import linear_batch_invariant
from vllm.model_executor.layers.linear import (
    ColumnParallelLinear,
    QKVParallelLinear,
    RowParallelLinear,
    UnquantizedLinearMethod,
)
from vllm.model_executor.layers.utils import dispatch_unquantized_gemm
from vllm.platforms import current_platform

from .ltx2_adapter_parser import AdapterManifest, AdapterTarget, iter_adapter_tensors, parse_ltx_adapter
from .ltx2_components import resolve_ltx_artifact
from .ltx2_recipes import LTX_DISTILLED_ADAPTER_SLOT

logger = init_logger(__name__)


@dataclass(frozen=True)
class _AdapterLayout:
    """Rank-local storage and execution layout for one LoRA pair."""

    kind: str
    input_size: int
    output_size: int
    a_input_start: int
    a_input_size: int
    b_output_start: int
    b_output_size: int
    output_start: int
    reduce_results: bool = False


def _module_device(module: nn.Module) -> torch.device:
    for tensor in (*module.parameters(), *module.buffers()):
        if tensor.device.type != "meta":
            return tensor.device
    raise RuntimeError(f"Cannot materialize an LTX phase adapter for meta module {module._get_name()}.")


def _build_layout(layer: nn.Module, target: AdapterTarget) -> _AdapterLayout:
    if isinstance(layer, QKVParallelLinear):
        if target.slice not in {"q", "k", "v"}:
            raise ValueError(
                f"LTX packed QKV target {target.module!r} requires a q/k/v adapter slice, got {target.slice!r}."
            )
        q_local = layer.num_heads * layer.head_size
        k_local = layer.num_kv_heads * layer.head_size
        v_local = layer.num_kv_heads * layer.v_head_size
        if target.slice == "q":
            output_size = layer.total_num_heads * layer.head_size
            b_output_start = layer.tp_rank * q_local
            b_output_size = q_local
            output_start = 0
        elif target.slice == "k":
            output_size = layer.total_num_kv_heads * layer.head_size
            shard_rank = layer.tp_rank // layer.num_kv_head_replicas
            b_output_start = shard_rank * k_local
            b_output_size = k_local
            output_start = q_local
        else:
            output_size = layer.total_num_kv_heads * layer.v_head_size
            shard_rank = layer.tp_rank // layer.num_kv_head_replicas
            b_output_start = shard_rank * v_local
            b_output_size = v_local
            output_start = q_local + k_local
        return _AdapterLayout(
            kind="qkv",
            input_size=layer.input_size,
            output_size=output_size,
            a_input_start=0,
            a_input_size=layer.input_size,
            b_output_start=b_output_start,
            b_output_size=b_output_size,
            output_start=output_start,
        )

    if isinstance(layer, ColumnParallelLinear):
        if layer.gather_output:
            raise ValueError(
                f"LTX phase adapter does not support gathered ColumnParallelLinear target {target.module!r}."
            )
        return _AdapterLayout(
            kind="column",
            input_size=layer.input_size,
            output_size=layer.output_size,
            a_input_start=0,
            a_input_size=layer.input_size,
            b_output_start=layer.tp_rank * layer.output_size_per_partition,
            b_output_size=layer.output_size_per_partition,
            output_start=0,
        )

    if isinstance(layer, RowParallelLinear):
        return _AdapterLayout(
            kind="row",
            input_size=layer.input_size,
            output_size=layer.output_size,
            a_input_start=layer.tp_rank * layer.input_size_per_partition,
            a_input_size=layer.input_size_per_partition,
            b_output_start=0,
            b_output_size=layer.output_size,
            output_start=0,
            reduce_results=layer.reduce_results,
        )

    if isinstance(layer, nn.Linear):
        return _AdapterLayout(
            kind="replicated",
            input_size=layer.in_features,
            output_size=layer.out_features,
            a_input_start=0,
            a_input_size=layer.in_features,
            b_output_start=0,
            b_output_size=layer.out_features,
            output_start=0,
        )

    raise TypeError(f"LTX phase adapter target {target.module!r} has unsupported layer type {type(layer).__name__}.")


class _AdapterPiece(nn.Module):
    """One rank-local A/B pair stored as non-persistent buffers."""

    def __init__(self, target: AdapterTarget, layout: _AdapterLayout) -> None:
        super().__init__()
        self.target = target
        self.layout = layout
        # The fixed graph is installed before base checkpoint loading.  The
        # empty buffers are replaced only after the base layer's quant/TP layout
        # has settled and before offload/compile snapshots are taken.
        self.register_buffer("lora_a", torch.empty(0), persistent=False)
        self.register_buffer("lora_b", torch.empty(0), persistent=False)

    def load(self, lora_a: torch.Tensor, lora_b: torch.Tensor, *, device: torch.device, dtype: torch.dtype) -> None:
        if tuple(lora_a.shape) != self.target.a_shape or tuple(lora_b.shape) != self.target.b_shape:
            raise ValueError(
                f"LTX adapter source shape changed for {self.target.source_module!r}: "
                f"A={tuple(lora_a.shape)}, B={tuple(lora_b.shape)}."
            )
        if lora_a.ndim != 2 or lora_b.ndim != 2:
            raise ValueError(f"LTX adapter target {self.target.source_module!r} must use rank-2 A/B tensors.")
        rank = self.target.a_shape[0]
        if lora_a.shape[1] != self.layout.input_size or lora_b.shape[0] != self.layout.output_size:
            raise ValueError(
                f"LTX adapter shape does not match {self.target.module!r}: "
                f"expected A=({rank}, {self.layout.input_size}), "
                f"B=({self.layout.output_size}, {rank}); "
                f"got A={tuple(lora_a.shape)}, B={tuple(lora_b.shape)}."
            )

        local_a = lora_a.narrow(1, self.layout.a_input_start, self.layout.a_input_size)
        local_b = lora_b.narrow(0, self.layout.b_output_start, self.layout.b_output_size)
        self.lora_a = local_a.to(device=device, dtype=dtype)
        self.lora_b = local_b.to(device=device, dtype=dtype)

    def delta(self, input_: torch.Tensor) -> torch.Tensor:
        if self.lora_a.numel() == 0 or self.lora_b.numel() == 0:
            raise RuntimeError(f"LTX phase adapter target {self.target.module!r} was not materialized.")
        if self.layout.kind == "row" and input_.shape[-1] != self.layout.a_input_size:
            input_ = input_.narrow(-1, self.layout.a_input_start, self.layout.a_input_size)
        input_ = input_.to(dtype=self.lora_a.dtype)
        delta = F.linear(F.linear(input_, self.lora_a), self.lora_b)
        if self.layout.kind == "row" and self.layout.reduce_results:
            delta = tensor_model_parallel_all_reduce(delta)
        return delta

    def weight_delta(self) -> torch.Tensor:
        """Materialize the rank-local BF16 weight delta used by official fusion."""
        if self.lora_a.numel() == 0 or self.lora_b.numel() == 0:
            raise RuntimeError(f"LTX phase adapter target {self.target.module!r} was not materialized.")
        return torch.matmul(self.lora_b, self.lora_a)


def _apply_unquantized_gemm_with_weight(
    layer: nn.Module,
    input_: torch.Tensor,
    weight: torch.Tensor,
    bias: torch.Tensor | None,
) -> torch.Tensor:
    # GEMM dispatch is selected from the host platform, not the tensor device.
    # In CPU tests on a ROCm host that selects rocm_unquantized_gemm, which has
    # no CPU kernel. Keep that combination on PyTorch's portable implementation.
    if current_platform.is_rocm() and input_.device.type == "cpu":
        return F.linear(input_, weight, bias)
    if envs.VLLM_BATCH_INVARIANT and current_platform.is_cuda_alike():
        return linear_batch_invariant(input_, weight, bias)
    return dispatch_unquantized_gemm()(layer, input_, weight, bias)


def _forward_with_weight(
    layer: nn.Module,
    input_: torch.Tensor,
    weight: torch.Tensor,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor | None]:
    """Run an unquantized base layer with an ephemeral fused weight."""
    if isinstance(layer, nn.Linear):
        return F.linear(input_, weight, layer.bias)

    quant_method = getattr(layer, "quant_method", None)
    if not isinstance(quant_method, UnquantizedLinearMethod):
        raise ValueError("LTX layer-fused phase LoRA requires unquantized Transformer linears.")

    if isinstance(layer, RowParallelLinear):
        if layer.input_is_parallel:
            input_parallel = input_
        else:
            input_parallel = split_tensor_along_last_dim(input_, num_partitions=layer.tp_size)[
                layer.tp_rank
            ].contiguous()
        bias = None if layer.tp_rank > 0 or layer.skip_bias_add else layer.bias
        output_parallel = _apply_unquantized_gemm_with_weight(layer, input_parallel, weight, bias)
        if layer.reduce_results and layer.tp_size > 1:
            output = tensor_model_parallel_all_reduce(output_parallel)
        else:
            output = output_parallel
    elif isinstance(layer, ColumnParallelLinear):
        bias = layer.bias if not layer.skip_bias_add else None
        output = _apply_unquantized_gemm_with_weight(layer, input_, weight, bias)
    else:
        raise TypeError(f"Unsupported LTX layer-fused linear type {type(layer).__name__}.")

    if not layer.return_bias:
        return output
    return output, layer.bias if layer.skip_bias_add else None


class _PhaseAdapterLinear(nn.Module):
    """Delegate the base linear unchanged and add an enabled adapter slot."""

    def __init__(
        self,
        base_layer: nn.Module,
        targets: Iterable[AdapterTarget],
        *,
        adapter_dtype: torch.dtype,
        layer_fused: bool = False,
    ) -> None:
        super().__init__()
        self.base_layer = base_layer
        self.adapter_dtype = adapter_dtype
        if layer_fused and not (
            isinstance(base_layer, nn.Linear)
            or isinstance(getattr(base_layer, "quant_method", None), UnquantizedLinearMethod)
        ):
            raise ValueError("LTX layer-fused phase LoRA requires unquantized Transformer linears.")
        self.layer_fused = layer_fused
        pieces = [_AdapterPiece(target, _build_layout(base_layer, target)) for target in targets]
        self.pieces = nn.ModuleList(pieces)
        self._pieces_by_target = {piece.target: piece for piece in pieces}
        self._enabled = False

    def load_target(self, target: AdapterTarget, lora_a: torch.Tensor, lora_b: torch.Tensor) -> None:
        try:
            piece = self._pieces_by_target[target]
        except KeyError as exc:
            raise ValueError(f"Adapter target {target.module!r} is not installed on this layer.") from exc
        piece.load(lora_a, lora_b, device=_module_device(self.base_layer), dtype=self.adapter_dtype)

    def set_enabled(self, enabled: bool) -> None:
        self._enabled = enabled

    def _add_adapter_delta(self, input_: torch.Tensor, output: torch.Tensor) -> torch.Tensor:
        if not self._enabled:
            return output
        if output.requires_grad:
            output = output.clone()
        for piece in self.pieces:
            delta = piece.delta(input_).to(dtype=output.dtype)
            start = piece.layout.output_start
            output[..., start : start + piece.layout.b_output_size].add_(delta)
        return output

    @torch.no_grad()
    def _fused_weight(self) -> torch.Tensor:
        base_weight = self.base_layer.weight
        fused_weight = base_weight.clone()
        for piece in self.pieces:
            start = piece.layout.output_start
            base_slice = base_weight.narrow(0, start, piece.layout.b_output_size)
            delta = piece.weight_delta()
            if delta.shape != base_slice.shape:
                raise ValueError(
                    f"LTX layer-fused delta shape {tuple(delta.shape)} does not match base weight slice "
                    f"{tuple(base_slice.shape)} for {piece.target.module!r}."
                )
            # Match official fusion exactly: materialize B@A in the adapter
            # dtype, then add the base tensor in-place before the main GEMM.
            delta.add_(base_slice)
            fused_weight.narrow(0, start, piece.layout.b_output_size).copy_(delta)
        return fused_weight

    def forward(self, input_: torch.Tensor):
        if not self._enabled:
            return self.base_layer(input_)
        if self.layer_fused:
            return _forward_with_weight(self.base_layer, input_, self._fused_weight())
        output = self.base_layer(input_)
        if isinstance(output, tuple):
            return (self._add_adapter_delta(input_, output[0]), *output[1:])
        return self._add_adapter_delta(input_, output)


class LTXPhaseAdapterRuntime:
    """Install and operate one fixed phase adapter on an existing transformer."""

    def __init__(
        self,
        transformer: nn.Module,
        manifest: AdapterManifest,
        *,
        dtype: torch.dtype,
        layer_fused: bool = False,
    ) -> None:
        self.transformer = transformer
        self.manifest = manifest
        self._wrappers: dict[str, _PhaseAdapterLinear] = {}
        self._target_names: tuple[str, ...] = ()
        self._materialized = False
        grouped: dict[str, list[AdapterTarget]] = defaultdict(list)
        for target in manifest.targets:
            grouped[target.module].append(target)

        original_layers = {name: self.transformer.get_submodule(name) for name in grouped}
        for name, targets in grouped.items():
            wrapper = _PhaseAdapterLinear(
                original_layers[name],
                targets,
                adapter_dtype=dtype,
                layer_fused=layer_fused,
            )
            parent_name, _, child_name = name.rpartition(".")
            parent = self.transformer.get_submodule(parent_name) if parent_name else self.transformer
            if child_name not in parent._modules:
                raise ValueError(f"LTX phase adapter target {name!r} is no longer a direct module child.")
            parent._modules[child_name] = wrapper
            self._wrappers[name] = wrapper

        self._target_names = tuple(sorted(self._wrappers, key=len, reverse=True))
        # LTX's custom loader receives the original checkpoint names.  Keep the
        # adaptation local to that loader rather than changing AutoWeightsLoader.
        self.transformer._phase_adapter_parameter_name = self.base_parameter_name
        logger.info(
            "Installed %d LTX %s phase-adapter linears",
            len(self._wrappers),
            "layer-fused" if layer_fused else "dynamic",
        )

    def base_parameter_name(self, name: str) -> str:
        for target_name in self._target_names:
            prefix = f"{target_name}."
            if name.startswith(prefix):
                return f"{target_name}.base_layer.{name[len(prefix) :]}"
        return name

    def finalize(self) -> None:
        """Stream rank-local A/B buffers after base weight processing completes."""
        if self._materialized:
            return
        for target, lora_a, lora_b in iter_adapter_tensors(self.manifest):
            self._wrappers[target.module].load_target(target, lora_a, lora_b)
        self._materialized = True
        logger.info("Materialized %d LTX phase-adapter tensor pairs", len(self.manifest.targets))

    def activate(self, adapter_slot: str | None) -> None:
        if adapter_slot is not None:
            if adapter_slot != LTX_DISTILLED_ADAPTER_SLOT:
                raise ValueError(f"Unknown LTX phase adapter slot {adapter_slot!r}.")
            if not self._materialized:
                raise RuntimeError("LTX phase adapter data must be materialized before activation.")
        for wrapper in self._wrappers.values():
            wrapper.set_enabled(adapter_slot is not None)


def build_ltx_phase_adapter(pipeline: Any) -> LTXPhaseAdapterRuntime | None:
    """Build the fixed adapter required by a multi-phase LTX recipe."""
    adapter_slots = {phase.adapter_slot for phase in pipeline.pipeline_recipe.phases if phase.adapter_slot is not None}
    if not adapter_slots:
        return None
    if adapter_slots != {LTX_DISTILLED_ADAPTER_SLOT}:
        raise ValueError(f"Unsupported LTX phase adapter slots: {sorted(adapter_slots)}.")

    profile = pipeline.component_profile
    if profile.distilled_lora_filename is None or profile.artifact_repo_id is None:
        raise ValueError(f"{profile.name} does not declare the required distilled adapter artifact.")
    if getattr(pipeline.od_config, "lora_path", None) is not None:
        raise ValueError(
            f"{pipeline.__class__.__name__} reserves LoRA execution for its phase adapter; "
            "request or static LoRA composition is not supported yet."
        )

    quantization_config = getattr(pipeline.od_config, "quantization_config", None)
    layer_fused = quantization_config is None
    adapter_path = resolve_ltx_artifact(
        pipeline.od_config.model,
        profile.artifact_repo_id,
        profile.distilled_lora_filename,
    )

    if layer_fused and pipeline.od_config.dtype is not torch.bfloat16:
        raise ValueError("LTX layer-fused phase LoRA requires a bfloat16 pipeline dtype.")

    manifest = parse_ltx_adapter(pipeline.transformer, adapter_path)
    runtime = LTXPhaseAdapterRuntime(
        pipeline.transformer,
        manifest,
        dtype=pipeline.od_config.dtype,
        layer_fused=layer_fused,
    )
    return runtime