from dataclasses import dataclass
from typing import Optional

from vllm.v1.core.sched.output import NewRequestData
from vllm.v1.request import Request

from vllm_omni.engine import AdditionalInformationPayload, PromptEmbedsPayload


@dataclass
class OmniNewRequestData(NewRequestData):
    """New request data for omni models with embeddings support.

    Extends NewRequestData to include prompt embeddings and additional
    information for direct transfer between pipeline stages.

    Args:
        prompt_embeds: Optional serialized prompt embeddings payload
        additional_information: Optional serialized additional information
            dictionary containing tensors or lists
    """

    # Optional serialized prompt embeddings
    prompt_embeds: Optional[PromptEmbedsPayload] = None
    # Optional serialized additional information
    additional_information: Optional[AdditionalInformationPayload] = None

    @classmethod
    def from_request(
        cls,
        request: Request,
        block_ids: tuple[list[int], ...],
    ) -> "OmniNewRequestData":
        """Create OmniNewRequestData from a Request object.

        Args:
            request: Request object to convert
            block_ids: Tuple of block ID lists for KV cache allocation

        Returns:
            OmniNewRequestData instance with data from the request
        """
        return cls(
            req_id=request.request_id,
            prompt_token_ids=request.prompt_token_ids,
            mm_features=request.mm_features,
            sampling_params=request.sampling_params,
            pooling_params=request.pooling_params,
            block_ids=block_ids,
            num_computed_tokens=request.num_computed_tokens,
            lora_request=request.lora_request,
            prompt_embeds=request.prompt_embeds,
            additional_information=request.additional_information,
        )