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

from __future__ import annotations

import json
import os
from collections.abc import Sequence
from typing import TYPE_CHECKING, Literal

import torch
import torch.nn as nn
from vllm.model_executor.layers.linear import (
    ColumnParallelLinear,
    LinearBase,
    ReplicatedLinear,
    RowParallelLinear,
)
from vllm.model_executor.models.transformers.utils import init_on_device_without_buffers
from vllm.model_executor.models.utils import maybe_prefix

from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.quantization import build_quant_config

if TYPE_CHECKING:
    from transformers import PretrainedConfig, PreTrainedModel
    from transformers.models.auto.auto_factory import _BaseAutoModelClass
    from vllm.model_executor.layers.quantization.base_config import (
        QuantizationConfig,
    )


Style = Literal["colwise", "colwise_rep", "rowwise", "rowwise_rep", "replicate"]


def make_attention_mask(
    hidden_states: torch.Tensor,
    seq_lengths: list[int],
) -> torch.Tensor | None:
    """Build a padding mask only when the batch contains padding.

    ``seq_lengths`` must be non-empty and contain valid sequence lengths.
    Returns a boolean mask of shape ``(len(seq_lengths), max(seq_lengths))``
    on the device of ``hidden_states``, with True marking valid tokens,
    or None when all sequence lengths are equal.

    Passing an all-true CUDA mask makes attention backends inspect a device
    scalar to select the dense path, which introduces a host-device sync.
    """
    max_seq_len = max(seq_lengths)
    if all(seq_len == max_seq_len for seq_len in seq_lengths):
        return None

    attention_mask = hidden_states.new_zeros(
        len(seq_lengths),
        max_seq_len,
        dtype=torch.bool,
    )
    for i, seq_len in enumerate(seq_lengths):
        attention_mask[i, :seq_len] = True
    return attention_mask


def replace_linear_class(
    linear: nn.Linear,
    style: Style = "replicate",
    quant_config: QuantizationConfig | None = None,
    *,
    prefix: str = "",
) -> ColumnParallelLinear | RowParallelLinear | ReplicatedLinear:
    """
    Replace nn.Linear with one of vLLM's tensor parallel linear classes.

    Args:
        linear: `nn.Linear` to be replaced.
        style: Tensor parallel style of the new linear, e.g. "colwise".
        quant_config: Quantization config for the new linear.
    Returns:
        The new linear.
    """

    if not isinstance(style, str):
        raise ValueError(f"Unsupported parallel style type {type(style)}, expected str")

    vllm_linear_maps = {
        "colwise": (ColumnParallelLinear, {}),
        "colwise_rep": (ColumnParallelLinear, {"gather_output": True}),
        "rowwise": (RowParallelLinear, {}),
        "rowwise_rep": (RowParallelLinear, {"input_is_parallel": False}),
        "replicate": (ReplicatedLinear, {}),
    }
    vllm_linear_cls, vllm_linear_kwargs = vllm_linear_maps[style]

    return vllm_linear_cls(
        input_size=linear.in_features,
        output_size=linear.out_features,
        bias=linear.bias is not None,
        quant_config=quant_config,
        prefix=prefix,
        return_bias=False,
        **vllm_linear_kwargs,
    )


def recursive_replace_linear(model: nn.Module, od_config: OmniDiffusionConfig):
    """Recursively replace modules in the model as needed.
    Currently, this replaces:
    - `nn.Linear` with vLLM's tensor parallel linear classes
    """
    # Prefix the patterns because we always start from `self.model`
    quant_config = build_quant_config(od_config.quantization_config)

    def _recursive_replace(module: nn.Module, prefix: str):
        for child_name, child_module in module.named_children():
            new_module = child_module
            qual_name = maybe_prefix(prefix, child_name)
            # Replace modules as needed
            if isinstance(child_module, nn.Linear):
                style = "replicate"
                new_module = replace_linear_class(child_module, style, quant_config, prefix=qual_name)
            else:
                _recursive_replace(child_module, prefix=qual_name)
            if new_module is not child_module:
                setattr(module, child_name, new_module)

    _recursive_replace(model, prefix="")


def recursive_replace_linear_with_quantization_config(
    model: nn.Module,
    quant_config: QuantizationConfig | None,
    prefix: str = "",
    skip_modules: Sequence[str] = (),
) -> None:
    for child_name, child_module in model.named_children():
        qual_name = maybe_prefix(prefix, child_name)
        if any(qual_name == path or qual_name.startswith(path + ".") for path in skip_modules):
            continue
        if isinstance(child_module, nn.Linear):
            setattr(
                model,
                child_name,
                replace_linear_class(child_module, quant_config=quant_config, prefix=qual_name),
            )
        else:
            recursive_replace_linear_with_quantization_config(
                child_module, quant_config, prefix=qual_name, skip_modules=skip_modules
            )


def init_parameters(
    module: nn.Module,
    dtype: torch.dtype | None,
    device: torch.device | None = None,
):
    for name, param in module.named_parameters(recurse=False):
        if param.device == torch.device("meta"):
            new_param = nn.Parameter(
                torch.empty_like(
                    param.data,
                    dtype=dtype,
                    device=device,
                ),
                requires_grad=param.requires_grad,
            )
            setattr(module, name, new_param)
    for child in module.children():
        init_parameters(child, dtype, device)


def release_module_parameters_to_meta(module: nn.Module) -> None:
    """Release a fully constructed module's parameter storage to meta.

    The module must finish construction first so initialization code and vLLM
    weight-loader metadata follow the ordinary CPU path. ``swap_tensors`` keeps
    each registered Parameter object's identity while dropping only its storage.
    Persistent buffers intentionally remain materialized.
    """
    for param in module.parameters():
        if param.device.type == "meta":
            continue

        meta_tensor = torch.empty_like(param, device="meta")
        if type(param) is nn.Parameter:
            meta_param = nn.Parameter(meta_tensor, requires_grad=param.requires_grad)
        else:
            meta_param = param.__class__.__new__(param.__class__, meta_tensor)
            meta_param.requires_grad_(param.requires_grad)
        meta_param.__dict__.update(getattr(param, "__dict__", {}))
        torch.utils.swap_tensors(param, meta_param)


def init_parameters_preserving_vllm_linear(
    module: nn.Module,
    dtype: torch.dtype | None,
    device: torch.device | None = None,
) -> None:
    if isinstance(module, LinearBase):
        return
    for name, param in module.named_parameters(recurse=False):
        if param.is_meta:
            setattr(
                module,
                name,
                nn.Parameter(torch.empty_like(param, dtype=dtype, device=device), requires_grad=param.requires_grad),
            )
    for child in module.children():
        init_parameters_preserving_vllm_linear(child, dtype, device)


def create_transformers_model(
    auto_cls: _BaseAutoModelClass,
    od_config: OmniDiffusionConfig,
    hf_config: PretrainedConfig,
    dtype: torch.dtype | None = None,
    device: torch.device | None = None,
) -> PreTrainedModel:
    """Create a HuggingFace model using the given auto class and model name."""
    dtype = dtype or od_config.dtype
    device = device or torch.get_default_device()
    with init_on_device_without_buffers("meta"):
        model = auto_cls.from_config(hf_config)
    recursive_replace_linear(model, od_config)
    init_parameters(model, dtype=dtype, device=device)
    return model


def create_transformers_model_with_vllm_linears(
    auto_cls: _BaseAutoModelClass,
    hf_config: PretrainedConfig,
    quant_config: QuantizationConfig | None,
    dtype: torch.dtype,
    device: torch.device,
    prefix: str = "",
    skip_modules: Sequence[str] = (),
) -> PreTrainedModel:
    """Build a Transformers model with vLLM linears."""
    with init_on_device_without_buffers("meta"):
        model = auto_cls.from_config(hf_config, dtype=dtype)
    recursive_replace_linear_with_quantization_config(model, quant_config, prefix=prefix, skip_modules=skip_modules)
    init_parameters_preserving_vllm_linear(model, dtype=dtype, device=device)
    return model


def _load_json(model_path: str, filename: str, local_files_only: bool = True) -> dict:
    """Load a JSON config file from a local path or HuggingFace Hub repo."""
    if local_files_only:
        path = os.path.join(model_path, *filename.split("/"))
        with open(path) as f:
            return json.load(f)
    else:
        from vllm_omni.transformers_utils.repo_utils import hf_api

        cached = hf_api().hf_hub_download(repo_id=model_path, filename=filename)
        with open(cached) as f:
            return json.load(f)