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

"""Request normalization shared by LTX pipeline variants."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any

import torch

from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch


@dataclass
class LTXRequestInputs:
    """Resolved request values shared by all LTX execution variants."""

    prompt: str | list[str] | None
    negative_prompt: str | list[str] | None
    height: int
    width: int
    num_frames: int
    frame_rate: float
    num_inference_steps: int
    guidance_scale: float
    guidance_rescale: float
    num_videos_per_prompt: int
    generator: torch.Generator | list[torch.Generator] | None
    latents: torch.Tensor | None
    audio_latents: torch.Tensor | None
    prompt_embeds: torch.Tensor | None
    negative_prompt_embeds: torch.Tensor | None
    prompt_attention_mask: torch.Tensor | None
    negative_prompt_attention_mask: torch.Tensor | None
    decode_timestep: float | list[float]
    decode_noise_scale: float | list[float] | None
    output_type: str
    max_sequence_length: int


def _unwrap_request_tensor(value: Any) -> Any:
    if isinstance(value, list):
        return value[0] if value else None
    return value


def _get_prompt_field(prompt: Any, *keys: str) -> Any:
    if isinstance(prompt, str):
        return None
    additional = prompt.get("additional_information")
    for key in keys:
        value = prompt.get(key)
        if value is None and isinstance(additional, dict):
            value = additional.get(key)
        if value is not None:
            return _unwrap_request_tensor(value)
    return None


def _get_audio_latents_from_sampling(sampling: Any) -> torch.Tensor | None:
    if sampling.audio_latents is not None:
        return sampling.audio_latents
    return sampling.extra_args.get("audio_latents")


class LTXRequestMixin:
    """Normalize serving requests without coupling them to a model version."""

    supports_request_batch = False
    supports_guidance_rescale = False

    @staticmethod
    def _resolve_request_sigmas(
        req: DiffusionRequestBatch,
        fallback: list[float] | None,
    ) -> list[float] | None:
        request_sigmas = [sampling.sigmas for sampling in req.sampling_params_list]
        first = request_sigmas[0] if request_sigmas else None
        if any(item != first for item in request_sigmas[1:]):
            raise ValueError("Batched LTX requests must use identical custom sigmas.")
        return first if first is not None else fallback

    def check_inputs(
        self,
        prompt: str | list[str] | None,
        height: int,
        width: int,
        prompt_embeds: torch.Tensor | None = None,
        negative_prompt_embeds: torch.Tensor | None = None,
        prompt_attention_mask: torch.Tensor | None = None,
        negative_prompt_attention_mask: torch.Tensor | None = None,
    ) -> None:
        if height % 32 != 0 or width % 32 != 0:
            raise ValueError(f"`height` and `width` must be divisible by 32 but are {height} and {width}.")
        if prompt is not None and prompt_embeds is not None:
            raise ValueError("Cannot forward both `prompt` and `prompt_embeds`.")
        if prompt is None and prompt_embeds is None:
            raise ValueError("Provide either `prompt` or `prompt_embeds`.")
        if prompt is not None and not isinstance(prompt, (str, list)):
            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
        if prompt_embeds is not None and prompt_attention_mask is None:
            raise ValueError("Must provide `prompt_attention_mask` when specifying `prompt_embeds`.")
        if negative_prompt_embeds is not None and negative_prompt_attention_mask is None:
            raise ValueError("Must provide `negative_prompt_attention_mask` when specifying `negative_prompt_embeds`.")
        if prompt_embeds is not None and negative_prompt_embeds is not None:
            if prompt_embeds.shape != negative_prompt_embeds.shape:
                raise ValueError(
                    "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
                    f" got {prompt_embeds.shape} and {negative_prompt_embeds.shape}."
                )
            if prompt_attention_mask.shape != negative_prompt_attention_mask.shape:
                raise ValueError(
                    "`prompt_attention_mask` and `negative_prompt_attention_mask` must have the same shape when"
                    f" passed directly, but got {prompt_attention_mask.shape} and"
                    f" {negative_prompt_attention_mask.shape}."
                )

    def _resolve_request_inputs(
        self,
        req: DiffusionRequestBatch,
        *,
        prompt: str | list[str] | None,
        negative_prompt: str | list[str] | None,
        height: int | None,
        width: int | None,
        num_frames: int | None,
        frame_rate: float | None,
        num_inference_steps: int | None,
        timesteps: list[int] | None,
        guidance_scale: float,
        num_videos_per_prompt: int | None,
        generator: torch.Generator | list[torch.Generator] | None,
        latents: torch.Tensor | None,
        audio_latents: torch.Tensor | None,
        prompt_embeds: torch.Tensor | None,
        negative_prompt_embeds: torch.Tensor | None,
        prompt_attention_mask: torch.Tensor | None,
        negative_prompt_attention_mask: torch.Tensor | None,
        decode_timestep: float | list[float],
        decode_noise_scale: float | list[float] | None,
        output_type: str,
        max_sequence_length: int | None,
        guidance_rescale: float = 0.0,
    ) -> LTXRequestInputs:
        sampling_params_list = req.sampling_params_list
        sampling = sampling_params_list[0]
        prompt = [item if isinstance(item, str) else (item.get("prompt") or "") for item in req.prompts] or prompt
        if all(isinstance(item, str) or item.get("negative_prompt") is None for item in req.prompts):
            negative_prompt = None
        elif req.prompts:
            negative_prompt = [
                "" if isinstance(item, str) else (item.get("negative_prompt") or "") for item in req.prompts
            ]

        height = sampling.height or height or self.one_stage_recipe.height
        width = sampling.width or width or self.one_stage_recipe.width
        num_frames = sampling.num_frames or num_frames or self.one_stage_recipe.num_frames
        frame_rate = sampling.resolved_frame_rate or frame_rate or self.one_stage_recipe.frame_rate
        num_inference_steps = (
            sampling.num_inference_steps or num_inference_steps or self.one_stage_recipe.num_inference_steps
        )
        if timesteps is None:
            num_inference_steps = max(int(num_inference_steps), 2)
        elif len(timesteps) < 2:
            raise ValueError("`timesteps` must contain at least 2 values for FlowMatchEulerDiscreteScheduler.")

        num_videos_per_prompt = (
            sampling.num_outputs_per_prompt if sampling.num_outputs_per_prompt > 0 else num_videos_per_prompt or 1
        )
        max_sequence_length = sampling.max_sequence_length or max_sequence_length or self.tokenizer_max_length
        if sampling.guidance_scale_provided:
            guidance_scale = sampling.guidance_scale
        if self.supports_guidance_rescale and sampling.guidance_rescale is not None:
            guidance_rescale = sampling.guidance_rescale

        if self.supports_request_batch:
            if generator is None:
                generator = req.collate_request_generators(num_videos_per_prompt, generator)
            latents = req.collate_request_tensors("latents", latents)
            audio_latents = DiffusionRequestBatch.collate_tensors(
                [_get_audio_latents_from_sampling(item) for item in sampling_params_list],
                "audio_latents",
                audio_latents,
            )
            prompt_fields = DiffusionRequestBatch.collate_prompt_field_map(
                req.prompts,
                {
                    "prompt_embeds": prompt_embeds,
                    "negative_prompt_embeds": negative_prompt_embeds,
                    "prompt_attention_mask": prompt_attention_mask,
                    "negative_prompt_attention_mask": negative_prompt_attention_mask,
                },
                field_aliases={
                    "prompt_attention_mask": ("prompt_attention_mask", "attention_mask"),
                    "negative_prompt_attention_mask": (
                        "negative_prompt_attention_mask",
                        "negative_attention_mask",
                    ),
                },
            )
            prompt_embeds = prompt_fields["prompt_embeds"]
            negative_prompt_embeds = prompt_fields["negative_prompt_embeds"]
            prompt_attention_mask = prompt_fields["prompt_attention_mask"]
            negative_prompt_attention_mask = prompt_fields["negative_prompt_attention_mask"]
            if prompt_embeds is not None:
                prompt = None
            if negative_prompt_embeds is not None:
                negative_prompt = None
        else:
            if generator is None:
                generator = sampling.generator
            if generator is None and sampling.seed is not None:
                generator = torch.Generator(device=self.device).manual_seed(sampling.seed)
            latents = sampling.latents if sampling.latents is not None else latents
            sampling_audio_latents = _get_audio_latents_from_sampling(sampling)
            if sampling_audio_latents is not None:
                audio_latents = sampling_audio_latents

            request_prompt_embeds = [_get_prompt_field(item, "prompt_embeds") for item in req.prompts]
            if any(value is not None for value in request_prompt_embeds):
                prompt_embeds = torch.stack(request_prompt_embeds)  # type: ignore[arg-type]
            request_negative_embeds = [_get_prompt_field(item, "negative_prompt_embeds") for item in req.prompts]
            if any(value is not None for value in request_negative_embeds):
                negative_prompt_embeds = torch.stack(request_negative_embeds)  # type: ignore[arg-type]
            request_prompt_masks = [
                _get_prompt_field(item, "prompt_attention_mask", "attention_mask") for item in req.prompts
            ]
            if any(value is not None for value in request_prompt_masks):
                prompt_attention_mask = torch.stack(request_prompt_masks)  # type: ignore[arg-type]
            request_negative_masks = [
                _get_prompt_field(item, "negative_prompt_attention_mask", "negative_attention_mask")
                for item in req.prompts
            ]
            if any(value is not None for value in request_negative_masks):
                negative_prompt_attention_mask = torch.stack(request_negative_masks)  # type: ignore[arg-type]

        if sampling.decode_timestep is not None:
            decode_timestep = sampling.decode_timestep
        if sampling.decode_noise_scale is not None:
            decode_noise_scale = sampling.decode_noise_scale
        if sampling.output_type is not None:
            output_type = sampling.output_type

        return LTXRequestInputs(
            prompt=prompt,
            negative_prompt=negative_prompt,
            height=int(height),
            width=int(width),
            num_frames=int(num_frames),
            frame_rate=float(frame_rate),
            num_inference_steps=int(num_inference_steps),
            guidance_scale=guidance_scale,
            guidance_rescale=guidance_rescale,
            num_videos_per_prompt=int(num_videos_per_prompt),
            generator=generator,
            latents=latents,
            audio_latents=audio_latents,
            prompt_embeds=prompt_embeds,
            negative_prompt_embeds=negative_prompt_embeds,
            prompt_attention_mask=prompt_attention_mask,
            negative_prompt_attention_mask=negative_prompt_attention_mask,
            decode_timestep=decode_timestep,
            decode_noise_scale=decode_noise_scale,
            output_type=output_type,
            max_sequence_length=int(max_sequence_length),
        )