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

import copy
import pprint
from dataclasses import asdict, dataclass, field
from typing import Any, TypeAlias

import torch
from typing_extensions import NotRequired, TypedDict
from vllm.inputs import EmbedsPrompt, PromptType, TextPrompt, TokensPrompt
from vllm.inputs.engine import TokensInput
from vllm.sampling_params import SamplingParams

from vllm_omni.lora.request import LoRARequest

DIFFUSION_QUALITY_LEVELS: tuple[str, ...] = ("lossless", "high")


class OmniTextPrompt(TextPrompt):
    """Text prompt with optional embeddings and stage payloads.

    Extends TextPrompt to support prompt embeddings and additional
    information payloads for direct transfer between pipeline stages.

    Attributes:
        prompt_embeds: Optional tensor containing prompt embeddings
        additional_information: Legacy optional dictionary containing
            additional information to pass along with the prompt.
        model_intermediate_buffer: Optional runner-owned stage payload.
    """

    negative_prompt: NotRequired[str]
    # Using modalities field to differentiate between different tasks for the same pipeline
    # for example Cosmos3OmniDiffusersPipeline handles t2i and t2v in the same pipeline.
    modalities: NotRequired[list[str]]
    prompt_embeds: NotRequired[torch.Tensor]
    negative_prompt_embeds: NotRequired[torch.Tensor]
    additional_information: NotRequired[dict[str, Any]]
    model_intermediate_buffer: NotRequired[dict[str, Any]]


class OmniInteractionEvent(OmniTextPrompt, total=False):
    """Prompt-event payload for mid-stream interactions."""

    pass


class OmniInteractionPrompt(TypedDict):
    """Normalized mid-stream interaction payload passed to ``submit_interaction``.

    WebSocket clients payload may omit ``event_id``; the API layer assigns a fallback ID in this case.
    """

    event_id: str
    event: OmniInteractionEvent
    transition_chunks: NotRequired[int]


class OmniTokensPrompt(TokensPrompt):
    """Tokens prompt with optional embeddings and stage payloads.

    Extends TokensPrompt to support prompt embeddings and additional
    information payloads for direct transfer between pipeline stages.

    Attributes:
        prompt_embeds: Optional tensor containing prompt embeddings
        additional_information: Legacy optional dictionary containing
            additional information to pass along with the prompt.
        model_intermediate_buffer: Optional runner-owned stage payload.
    """

    negative_prompt: NotRequired[str]
    prompt_embeds: NotRequired[torch.Tensor]
    negative_prompt_embeds: NotRequired[list[torch.Tensor] | None]
    """The embeddings of the prompt."""

    # New: optional additional information dictionary
    # Values may be torch.Tensor or list
    additional_information: NotRequired[dict[str, Any]]
    model_intermediate_buffer: NotRequired[dict[str, Any]]


class OmniTokenInputs(TokensInput):
    """Token inputs with optional embeddings and stage payloads.

    Extends TokensInput to support prompt embeddings and additional
    information payloads for direct transfer between pipeline stages.

    Attributes:
        prompt_embeds: Optional tensor containing prompt embeddings
            aligned with token IDs
        additional_information: Legacy optional dictionary containing
            additional information to pass along with the inputs.
        model_intermediate_buffer: Optional runner-owned stage payload.
    """

    # New: optional prompt embeddings aligned with token ids
    negative_prompt: NotRequired[str]
    prompt_embeds: NotRequired[torch.Tensor]
    negative_prompt_embeds: NotRequired[list[torch.Tensor] | None]

    # New: optional additional information dictionary
    # Values may be torch.Tensor or list
    additional_information: NotRequired[dict[str, Any]]
    model_intermediate_buffer: NotRequired[dict[str, Any]]


class OmniEmbedsPrompt(EmbedsPrompt):
    """Embeddings prompt with optional stage payloads.

    Extends EmbedsPrompt to support additional information payloads
    for direct transfer between pipeline stages.

    Attributes:
        prompt_embeds: Optional tensor containing prompt embeddings
        additional_information: Legacy optional dictionary containing
            additional information to pass along with the prompt.
        model_intermediate_buffer: Optional runner-owned stage payload.
    """

    # New: optional prompt embeddings aligned with token ids
    prompt_embeds: NotRequired[torch.Tensor]
    negative_prompt_embeds: NotRequired[list[torch.Tensor] | None]

    # New: optional additional information dictionary
    # Values may be torch.Tensor or list
    additional_information: NotRequired[dict[str, Any]]
    model_intermediate_buffer: NotRequired[dict[str, Any]]


class OmniCustomPrompt(TypedDict, total=False):
    """Custom prompt type for diffusion pipelines with pre-tokenized inputs.

    Allows users to pass pre-tokenized prompt IDs, attention masks, and extra
    arguments directly, bypassing the tokenization stage in the pipeline.

    Attributes:
        prompt_ids: Pre-tokenized prompt token IDs (single or batched)
        negative_prompt_ids: Pre-tokenized negative prompt token IDs
        prompt_mask: Attention mask tensor for the prompt
        negative_prompt_mask: Attention mask tensor for the negative prompt
        extra_args: Additional pipeline-specific arguments
    """

    prompt_ids: list[int] | list[list[int]]
    negative_prompt_ids: list[int] | list[list[int]]
    prompt_mask: torch.Tensor
    negative_prompt_mask: torch.Tensor
    extra_args: dict[str, Any]


# Must ensure that all additional prompt types are inherited from vLLM prompt types
# Because TypedDict doesn't support isinstance and are dict. Cannot distinguish them in runtime.
# Inheritance ensure that there are only additional fields but not removing fields--safe to route to LLM.generate()
OmniSingletonPrompt: TypeAlias = str | list[int] | OmniTextPrompt | OmniTokensPrompt | OmniEmbedsPrompt
"""Omni singleton prompt type extending vLLM's SingletonPrompt."""

OmniPromptType: TypeAlias = PromptType | OmniTextPrompt | OmniTokensPrompt | OmniEmbedsPrompt | OmniCustomPrompt


def token_inputs_omni(
    prompt_token_ids: list[int],
    prompt: str | None = None,
    cache_salt: str | None = None,
    prompt_embeds: torch.Tensor | None = None,
    additional_information: dict[str, Any] | None = None,
    model_intermediate_buffer: dict[str, Any] | None = None,
) -> OmniTokenInputs:
    """Construct token inputs with optional embeddings and metadata.

    Creates an OmniTokenInputs object with token IDs and optional
    embeddings and additional information for pipeline stage transfer.

    Args:
        prompt_token_ids: List of token IDs for the prompt
        prompt: Optional prompt string
        cache_salt: Optional cache salt for prefix caching
        prompt_embeds: Optional tensor containing prompt embeddings
        additional_information: Legacy optional dictionary containing
            additional information (tensors or lists).
        model_intermediate_buffer: Optional runner-owned stage payload.

    Returns:
        OmniTokenInputs instance with the provided data
    """
    inputs = OmniTokenInputs(type="token", prompt_token_ids=prompt_token_ids)

    if prompt is not None:
        inputs["prompt"] = prompt
    if cache_salt is not None:
        inputs["cache_salt"] = cache_salt
    if prompt_embeds is not None:
        inputs["prompt_embeds"] = prompt_embeds
    if additional_information is not None:
        inputs["additional_information"] = additional_information
    if model_intermediate_buffer is not None:
        inputs["model_intermediate_buffer"] = model_intermediate_buffer

    return inputs


@dataclass
class OmniDiffusionSamplingParams:
    """
    The collection of sampling parameters passed to diffusion pipelines.

    This dataclass contains all information needed during the diffusion pipeline
    execution, allowing methods to update specific components without needing
    to manage numerous individual parameters.
    """

    # Additional text-related parameters
    max_sequence_length: int | None = None
    prompt_template: dict[str, Any] | None = None
    do_classifier_free_guidance: bool = False
    output_type: str | None = None

    # Request-scoped quality intent. ``None`` delegates the default behavior
    # to the model; explicit levels select a model-owned quality policy.
    quality: str | None = None

    # Batch info
    num_outputs_per_prompt: int = 1
    seed: int | None = None
    generator: torch.Generator | list[torch.Generator] | None = None
    generator_device: str | None = None

    # layered info
    layers: int = 4

    # cfg info
    cfg_normalize: bool = False

    # caption language
    use_en_prompt: bool = False

    # different bucket in (640, 1024) to determine the condition and output resolution
    resolution: int = 640

    # Tracking if embeddings are already processed
    is_prompt_processed: bool = False

    # Latent tensors
    latents: torch.Tensor | None = None
    audio_latents: torch.Tensor | None = None
    raw_latent_shape: torch.Tensor | None = None
    noise_pred: torch.Tensor | None = None
    image_latent: torch.Tensor | None = None

    # Latent dimensions
    height_latents: list[int] | int | None = None
    width_latents: list[int] | int | None = None
    num_frames: int = 1  # Default for image models
    num_frames_round_down: bool = False  # Whether to round down num_frames if it's not divisible by num_gpus

    # Original dimensions (before VAE scaling)
    height: int | None = None
    width: int | None = None
    fps: float | None = None  # Request/output-encoding FPS; fractional rates are supported.
    frame_rate: float | None = None  # Floating-point rate used by the diffusion model when it differs from `fps`.
    height_not_provided: bool = False
    width_not_provided: bool = False
    enable_frame_interpolation: bool = False
    frame_interpolation_exp: int = 1
    frame_interpolation_scale: float = 1.0
    frame_interpolation_model_path: str | None = None

    # Timesteps
    timesteps: torch.Tensor | None = None
    timestep: torch.Tensor | float | int | None = None
    step_index: int | None = None
    boundary_ratio: float | None = None

    # Scheduler parameters – ``None`` means "not explicitly set by the caller";
    # each pipeline's ``forward()`` decides its own model-specific default.
    num_inference_steps: int | None = None
    guidance_scale: float | None = None
    guidance_scale_provided: bool = False
    guidance_scale_2: float | None = None
    guidance_scale_2_provided: bool = False
    guidance_rescale: float = 0.0
    strength: float | None = None  # I2I: Z-Image specific now, uses to control denoising start timestep
    decode_timestep: float | list[float] | None = None
    decode_noise_scale: float | list[float] | None = None
    eta: float = 0.0
    sigmas: list[float] | None = None

    true_cfg_scale: float | None = None  # qwen-image specific now

    n_tokens: int | None = None
    extra_step_kwargs: dict[str, Any] = field(default_factory=dict)

    # [Omni] KV Cache Transfer, for bagel model now
    past_key_values: Any | None = None  # Injected KV Cache
    kv_metadata: dict[str, Any] | None = None  # Metadata for KV Cache (e.g., kv_lens, ropes)
    need_kv_receive: bool = True  # Flag to indicate if this request expects KV transfer

    # [Omni] Multi-KV for CFG: populated by model-specific cfg_kv_collect_func
    cfg_text_past_key_values: Any | None = None
    cfg_img_past_key_values: Any | None = None
    cfg_text_kv_metadata: dict[str, Any] | None = None
    cfg_img_kv_metadata: dict[str, Any] | None = None
    cfg_kv_request_ids: dict[str, str] | None = None
    cfg_active_branch: str | None = None
    cfg_branch_roles: list[str] | None = None
    cfg_branch_past_key_values: dict[str, Any] | None = None
    cfg_branch_kv_metadata: dict[str, dict[str, Any]] | None = None

    # Component modules
    modules: dict[str, Any] = field(default_factory=dict)

    return_trajectory_latents: bool = False
    return_trajectory_decoded: bool = False
    trajectory_timesteps: list[torch.Tensor] | None = None
    trajectory_latents: torch.Tensor | None = None

    # Extra parameters that might be needed by specific pipeline implementations
    extra_args: dict[str, Any] = field(default_factory=dict)

    # Misc
    save_output: bool = True
    return_frames: bool = False

    # LoRA
    lora_request: LoRARequest | None = None
    lora_scale: float = 1.0

    # STA parameters
    STA_param: list | None = None
    is_cfg_negative: bool = False
    mask_search_final_result_pos: list[list] | None = None
    mask_search_final_result_neg: list[list] | None = None

    # VSA parameters
    VSA_sparsity: float = 0.0
    # perf_logger: PerformanceLogger | None = None

    # stage logging
    # logging_info: PipelineLoggingInfo = field(default_factory=PipelineLoggingInfo)

    # profile
    profile: bool = False
    num_profiled_timesteps: int = 8

    # debugging
    debug: bool = False

    # results
    output: torch.Tensor | None = None

    # Internal control-plane event used by asynchronous APIs. When enabled,
    # the diffusion engine emits a non-terminal output when the scheduler
    # first admits this request for execution.
    emit_request_lifecycle: bool = False

    def __post_init__(self) -> None:
        if self.quality is not None and self.quality not in DIFFUSION_QUALITY_LEVELS:
            raise ValueError(f"quality must be one of {list(DIFFUSION_QUALITY_LEVELS)}, got {self.quality!r}")

    @property
    def batch_size(self):
        # This class is changed to only represent a single prompt request
        # Only adjust batch size for number of videos per prompt
        return self.num_outputs_per_prompt

    @property
    def resolved_frame_rate(self) -> float | None:
        """Frame rate the model should run at, or ``None`` if the user did not provide one.

        Precedence is ``frame_rate`` (the model-internal rate) over ``fps`` (the
        request/output-encoding alias). Returns ``None`` when neither was set, which is
        the model-agnostic "fps not provided" signal the video serving layer emits when
        the request omits fps. Consumers MUST treat ``None`` as "use your own default" and
        guard the read (e.g. ``req.sampling_params.resolved_frame_rate or 24.0``); never
        feed it into arithmetic unguarded.
        """
        if self.frame_rate is not None:
            return float(self.frame_rate)

        fps = self.fps
        if isinstance(fps, list):
            fps = fps[0] if fps else None

        if fps is None:
            return None

        return float(fps)

    def __str__(self):
        return pprint.pformat(asdict(self), indent=2, width=120)

    def clone(self) -> "OmniDiffusionSamplingParams":
        return copy.deepcopy(self)

    @classmethod
    def from_params(cls, params: Any) -> "OmniDiffusionSamplingParams":
        """Normalize caller params for diffusion stages.

        Existing Omni params are returned unchanged. Plain vLLM SamplingParams
        copy seed and matching extra_args fields; remaining extra_args stay nested.
        Unsupported types raise TypeError.
        """
        if isinstance(params, cls):
            return params
        if isinstance(params, SamplingParams):
            extra = dict(getattr(params, "extra_args", None) or {})
            known = set(cls.__dataclass_fields__)
            mapped = {k: extra.pop(k) for k in list(extra) if k in known}
            if extra:
                mapped["extra_args"] = extra
            seed = getattr(params, "seed", None)
            if seed is not None:
                mapped.setdefault("seed", seed)
            return cls(**mapped)
        raise TypeError(
            "Diffusion stage requires OmniDiffusionSamplingParams or vllm.SamplingParams, "
            f"got {type(params).__name__!r}."
        )


OmniSamplingParams: TypeAlias = SamplingParams | OmniDiffusionSamplingParams