from dataclasses import dataclass, field

from vllm.v1.core.sched.output import CachedRequestData, NewRequestData, SchedulerOutput
from vllm.v1.request import Request

from vllm_omni.engine import AdditionalInformationPayload


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

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

    Note: prompt_embeds is inherited from NewRequestData
    (torch.Tensor | None).

    Args:
        external_req_id: Optional external request ID for tracking
        additional_information: Optional serialized additional information
            dictionary containing tensors or lists
    """

    external_req_id: str | None = None
    additional_information: AdditionalInformationPayload | None = None

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

        Args:
            request: Request object to convert
            block_ids: Tuple of block ID lists for KV cache allocation
            prefill_token_ids: Optional prefill token IDs for v2 model runner

        Returns:
            OmniNewRequestData instance with data from the request
        """
        return cls(
            req_id=request.request_id,
            external_req_id=getattr(request, "external_req_id", None),
            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=getattr(request, "prompt_embeds", None),
            prompt_is_token_ids=getattr(request, "prompt_is_token_ids", None),
            prefill_token_ids=prefill_token_ids,
            additional_information=getattr(request, "additional_information", None),
        )


@dataclass
class OmniCachedRequestData(CachedRequestData):
    """Cached request data for omni models with embeddings support.

    Args:
        prompt_token_ids: Mapping from request ID to list of prompt token IDs
    """

    prompt_token_ids: dict[str, list[int]]
    additional_information: dict[str, dict | None]


@dataclass
class OmniChunkRecvHandle:
    """Minimal identifier carried from scheduler to runner for chunk-recv
    registration.

    The runner's ``register_chunk_recv`` only consumes ``request_id`` and
    ``external_req_id`` from each pending request, so we ship just those
    two fields instead of the full Request object.  Concrete typing
    keeps msgspec serialization deterministic across IPC (default,
    PD-disagg, multi-node executor variants) and avoids the
    ``list[Any]`` fallback path.
    """

    request_id: str
    external_req_id: str | None = None


@dataclass
class OmniSchedulerOutput(SchedulerOutput):
    """Scheduler output with omni-specific transfer metadata."""

    finished_requests_needing_kv_transfer: dict[str, dict] = field(default_factory=dict)
    pending_input_registrations: list[OmniChunkRecvHandle] = field(default_factory=list)