# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Block discovery for layerwise offload.

Shared between LayerWiseOffloadBackend and DistributedLayerwiseOffloadBackend.
"""

from __future__ import annotations

from torch import nn
from vllm.logger import init_logger

logger = init_logger(__name__)


def get_blocks_attr_names(model: nn.Module) -> list[str]:
    """Get block attribute names from model class."""
    attrs: list[str] = getattr(model.__class__, "_layerwise_offload_blocks_attrs", [])

    if not attrs:
        old_attr = getattr(model.__class__, "_layerwise_offload_blocks_attr", None)
        if old_attr is not None:
            logger.warning(
                "'_layerwise_offload_blocks_attr' is deprecated, "
                "please use '_layerwise_offload_blocks_attrs' instead. "
                "Example: _layerwise_offload_blocks_attrs = ['blocks']"
            )
            attrs = [old_attr] if isinstance(old_attr, str) else list(old_attr)

    return attrs


def set_blocks_attr_names(model: nn.Module, names: list[str]) -> None:
    if not hasattr(model.__class__, "_layerwise_offload_blocks_attrs"):
        setattr(model.__class__, "_layerwise_offload_blocks_attrs", names)


def get_blocks_from_dit(
    model: nn.Module,
    block_attrs: tuple[str, ...] | None = None,
) -> tuple[list[str], list[nn.Module]]:
    """Retrieve blocks from an explicit plan or the DiT class metadata."""
    explicit_attrs = block_attrs is not None
    blocks_attr_names = list(block_attrs) if explicit_attrs else get_blocks_attr_names(model)
    if not blocks_attr_names:
        logger.warning(
            f"No _layerwise_offload_blocks_attrs defined for {model.__class__.__name__}, "
            "skipping distributed layerwise offloading"
        )
        return [], []

    blocks: list[nn.Module] = []
    for name in blocks_attr_names:
        attr = getattr(model, name, None)
        if attr is None:
            source = "OffloadPlan.block_attrs" if explicit_attrs else "_layerwise_offload_blocks_attrs"
            raise AttributeError(
                f"Attribute '{name}' declared in {source} does not exist on model {model.__class__.__name__}"
            )
        try:
            attr_iter = iter(attr)
        except TypeError:
            if isinstance(attr, nn.Module):
                logger.warning(
                    "Attribute '%s' on %s is not iterable; treating it as one block.",
                    name,
                    model.__class__.__name__,
                )
                blocks.append(attr)
                continue

            logger.warning(
                "Attribute '%s' on %s is not iterable (got %s); skipping it.",
                name,
                model.__class__.__name__,
                type(attr).__name__,
            )
        else:
            blocks.extend(attr_iter)

    if not blocks:
        logger.warning(
            "No blocks found in %s for %s, skipping distributed layerwise offloading",
            blocks_attr_names,
            model.__class__.__name__,
        )
        return [], []

    return blocks_attr_names, blocks