from typing import Any, Optional

try:
    from typing import NotRequired
except ImportError:
    # Python < 3.11: use typing_extensions
    from typing_extensions import NotRequired

import torch
from vllm.inputs.data import EmbedsPrompt, TokenInputs, TokensPrompt


class OmniTokensPrompt(TokensPrompt):
    prompt_embeds: NotRequired[torch.Tensor]
    """The embeddings of the prompt."""

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


class OmniTokenInputs(TokenInputs):
    # New: optional prompt embeddings aligned with token ids
    prompt_embeds: NotRequired[torch.Tensor]

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


class OmniEmbedsPrompt(EmbedsPrompt):
    # New: optional prompt embeddings aligned with token ids
    prompt_embeds: NotRequired[torch.Tensor]

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


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

    return inputs