# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import asyncio
from abc import ABC, abstractmethod
from collections import defaultdict
from collections.abc import Awaitable, Callable
from functools import cached_property, partial
from typing import TYPE_CHECKING, Any, Final

import torch
from PIL import Image
from typing_extensions import override

from vllm import envs
from vllm.config import ModelConfig
from vllm.exceptions import VLLMValidationError
from vllm.multimodal.media import MEDIA_CONNECTOR_REGISTRY, MediaConnector
from vllm.renderers.embed_utils import (
    safe_load_prompt_embeds,
    safe_load_prompt_embeds_async,
)

from .types import InputAudio, ModalityStr, MultiModalEmbedsPayload

if TYPE_CHECKING:
    from .item_tracker import AsyncMultiModalItemTracker, MultiModalItemTracker

MODALITY_PLACEHOLDERS_MAP = {
    "image": "<##IMAGE##>",
    "audio": "<##AUDIO##>",
    "video": "<##VIDEO##>",
    "prompt_embeds": "<##PROMPT_EMBEDS##>",
}


PROMPT_EMBEDS_PLACEHOLDER_TOKEN: Final[str] = "<prompt_embeds>"
"""The special token used as a placeholder for each embedding
position during chat template rendering.

Registered as an additional special token when `--enable-prompt-embeds` is set.
See `_ensure_prompt_embeds_placeholder_token` in `vllm/renderers/hf.py`.
"""

_ENABLE_PROMPT_EMBEDS_ERROR: Final[str] = (
    "You must set `--enable-prompt-embeds` to input `prompt_embeds`"
)


class BaseMultiModalContentParser(ABC):
    def __init__(self) -> None:
        super().__init__()

        # stores model placeholders list with corresponding
        # general MM placeholder:
        # {
        #   "<##IMAGE##>": ["<image>", "<image>", "<image>"],
        #   "<##AUDIO##>": ["<audio>", "<audio>"],
        #   "<##PROMPT_EMBEDS##>": ["<prompt_embeds>", "<prompt_embeds>"]
        # }
        self._placeholder_storage: dict[str, list] = defaultdict(list)

    @property
    @abstractmethod
    def model_config(self) -> ModelConfig:
        raise NotImplementedError

    def _add_placeholder(self, modality: ModalityStr, placeholder: str | None):
        mod_placeholder = MODALITY_PLACEHOLDERS_MAP[modality]
        if placeholder:
            self._placeholder_storage[mod_placeholder].append(placeholder)

    def mm_placeholder_storage(self) -> dict[str, list]:
        return dict(self._placeholder_storage)

    @abstractmethod
    def parse_image(self, image_url: str | None, uuid: str | None = None) -> None:
        raise NotImplementedError

    @abstractmethod
    def parse_image_embeds(
        self,
        image_embeds: MultiModalEmbedsPayload | None,
        uuid: str | None = None,
    ) -> None:
        raise NotImplementedError

    @abstractmethod
    def parse_image_pil(
        self, image_pil: Image.Image | None, uuid: str | None = None
    ) -> None:
        raise NotImplementedError

    @abstractmethod
    def parse_audio(self, audio_url: str | None, uuid: str | None = None) -> None:
        raise NotImplementedError

    @abstractmethod
    def parse_input_audio(
        self, input_audio: InputAudio | None, uuid: str | None = None
    ) -> None:
        raise NotImplementedError

    @abstractmethod
    def parse_audio_embeds(
        self,
        audio_embeds: MultiModalEmbedsPayload | None,
        uuid: str | None = None,
    ) -> None:
        raise NotImplementedError

    @abstractmethod
    def parse_prompt_embeds(self, data: str) -> None:
        raise NotImplementedError

    @abstractmethod
    def parse_video(self, video_url: str | None, uuid: str | None = None) -> None:
        raise NotImplementedError

    @abstractmethod
    def parse_video_embeds(
        self,
        video_embeds: MultiModalEmbedsPayload | None,
        uuid: str | None = None,
    ) -> None:
        raise NotImplementedError


class MultiModalContentParser(BaseMultiModalContentParser):
    def __init__(
        self,
        tracker: "MultiModalItemTracker",
        mm_processor_kwargs: dict[str, Any] | None = None,
    ) -> None:
        super().__init__()

        self._tracker = tracker
        self._mm_processor_kwargs = mm_processor_kwargs

    @cached_property
    def _connector(self) -> MediaConnector:
        # Connector setup may probe VLLM_MEDIA_CACHE. Defer it until a request
        # actually contains media so text-only parsing never blocks on that I/O.
        return MEDIA_CONNECTOR_REGISTRY.load(
            envs.VLLM_MEDIA_CONNECTOR,
            media_io_kwargs=self._tracker.media_io_kwargs,
            allowed_local_media_path=self._tracker.allowed_local_media_path,
            allowed_media_domains=self._tracker.allowed_media_domains,
            decoded_budget=self._tracker.decoded_budget,
            download_budget=self._tracker.download_budget,
        )

    @property
    def model_config(self) -> ModelConfig:
        return self._tracker.model_config

    @override
    def parse_prompt_embeds(self, data: str) -> None:
        """Decode a base64 prompt embeds tensor and store it in the tracker.

        Emits a single `PROMPT_EMBEDS_PLACEHOLDER_TOKEN` sentinel per
        content part. The renderer later expands each sentinel to a span of
        `tensor.shape[0]` placeholder tokens after tokenization.
        """
        if not self.model_config.enable_prompt_embeds:
            raise VLLMValidationError(
                _ENABLE_PROMPT_EMBEDS_ERROR, parameter="prompt_embeds"
            )

        tensor = safe_load_prompt_embeds(self.model_config, data.encode())
        self._tracker.add("prompt_embeds", (tensor, None))
        self._add_placeholder("prompt_embeds", PROMPT_EMBEDS_PLACEHOLDER_TOKEN)

    def parse_image(self, image_url: str | None, uuid: str | None = None) -> None:
        image = self._connector.fetch_image(image_url) if image_url else None

        placeholder = self._tracker.add("image", (image, uuid))
        self._add_placeholder("image", placeholder)

    def parse_image_embeds(
        self,
        image_embeds: MultiModalEmbedsPayload | None,
        uuid: str | None = None,
    ) -> None:
        mm_config = self.model_config.get_multimodal_config()
        if not mm_config.enable_mm_embeds:
            raise VLLMValidationError(
                "You must set `--enable-mm-embeds` to input `image_embeds`",
                parameter="image_embeds",
            )

        if isinstance(image_embeds, dict):
            embeds = {
                k: self._connector.fetch_image_embedding(v) if isinstance(v, str) else v
                for k, v in image_embeds.items()
            }
            placeholder = self._tracker.add("image_embeds", (embeds, uuid))

        if isinstance(image_embeds, str):
            embedding = self._connector.fetch_image_embedding(image_embeds)
            placeholder = self._tracker.add("image_embeds", (embedding, uuid))

        if image_embeds is None:
            placeholder = self._tracker.add("image_embeds", (None, uuid))

        self._add_placeholder("image", placeholder)

    def parse_audio_embeds(
        self,
        audio_embeds: MultiModalEmbedsPayload | None,
        uuid: str | None = None,
    ) -> None:
        mm_config = self.model_config.get_multimodal_config()
        if not mm_config.enable_mm_embeds:
            raise VLLMValidationError(
                "You must set `--enable-mm-embeds` to input `audio_embeds`",
                parameter="audio_embeds",
            )

        if isinstance(audio_embeds, dict):
            embeds = {
                k: self._connector.fetch_audio_embedding(v) if isinstance(v, str) else v
                for k, v in audio_embeds.items()
            }
            placeholder = self._tracker.add("audio_embeds", (embeds, uuid))
        elif isinstance(audio_embeds, str):
            embedding = self._connector.fetch_audio_embedding(audio_embeds)
            placeholder = self._tracker.add("audio_embeds", (embedding, uuid))
        else:
            placeholder = self._tracker.add("audio_embeds", (None, uuid))

        self._add_placeholder("audio", placeholder)

    def parse_image_pil(
        self, image_pil: Image.Image | None, uuid: str | None = None
    ) -> None:
        if image_pil is not None:
            self._tracker.decoded_budget.charge(
                image_pil.size[0] * image_pil.size[1] * 4
            )
        placeholder = self._tracker.add("image", (image_pil, uuid))
        self._add_placeholder("image", placeholder)

    def parse_audio(self, audio_url: str | None, uuid: str | None = None) -> None:
        audio = self._connector.fetch_audio(audio_url) if audio_url else None

        placeholder = self._tracker.add("audio", (audio, uuid))
        self._add_placeholder("audio", placeholder)

    def parse_input_audio(
        self, input_audio: InputAudio | None, uuid: str | None = None
    ) -> None:
        if input_audio:
            audio_data = input_audio.get("data", "")
            audio_format = input_audio.get("format", "")
            if audio_data:
                audio_url = f"data:audio/{audio_format};base64,{audio_data}"
            else:
                # If a UUID is provided, audio data may be empty.
                audio_url = None
        else:
            audio_url = None

        return self.parse_audio(audio_url, uuid)

    def parse_video(self, video_url: str | None, uuid: str | None = None) -> None:
        video = (
            self._connector.fetch_video(
                video_url=video_url,
                video_processor=self._tracker.video_processor_name,
            )
            if video_url
            else None
        )

        placeholder = self._tracker.add("video", (video, uuid))
        self._add_placeholder("video", placeholder)

        # Extract audio from video if use_audio_in_video is True
        if (
            video_url
            and self._mm_processor_kwargs
            and self._mm_processor_kwargs.get("use_audio_in_video", False)
        ):
            audio = self._connector.fetch_audio(video_url) if video_url else None
            audio_placeholder = self._tracker.add("audio", (audio, uuid))
            self._add_placeholder("audio", audio_placeholder)

    def parse_video_embeds(
        self,
        video_embeds: MultiModalEmbedsPayload | None,
        uuid: str | None = None,
    ) -> None:
        mm_config = self.model_config.get_multimodal_config()
        if not mm_config.enable_mm_embeds:
            raise VLLMValidationError(
                "You must set `--enable-mm-embeds` to input `video_embeds`",
                parameter="video_embeds",
            )

        if isinstance(video_embeds, dict):
            embeds = {
                k: self._connector.fetch_video_embedding(v) if isinstance(v, str) else v
                for k, v in video_embeds.items()
            }
            placeholder = self._tracker.add("video_embeds", (embeds, uuid))
        elif isinstance(video_embeds, str):
            embedding = self._connector.fetch_video_embedding(video_embeds)
            placeholder = self._tracker.add("video_embeds", (embedding, uuid))
        else:
            placeholder = self._tracker.add("video_embeds", (None, uuid))

        self._add_placeholder("video", placeholder)


class AsyncMultiModalContentParser(BaseMultiModalContentParser):
    def __init__(
        self,
        tracker: "AsyncMultiModalItemTracker",
        mm_processor_kwargs: dict[str, Any] | None = None,
    ) -> None:
        super().__init__()

        self._tracker = tracker
        self._mm_processor_kwargs: dict[str, Any] | None = mm_processor_kwargs

    @cached_property
    def _connector(self) -> MediaConnector:
        # Connector setup may probe VLLM_MEDIA_CACHE. Defer it until a request
        # actually contains media so text-only parsing never blocks on that I/O.
        return MEDIA_CONNECTOR_REGISTRY.load(
            envs.VLLM_MEDIA_CONNECTOR,
            media_io_kwargs=self._tracker.media_io_kwargs,
            allowed_local_media_path=self._tracker.allowed_local_media_path,
            allowed_media_domains=self._tracker.allowed_media_domains,
            decoded_budget=self._tracker.decoded_budget,
            download_budget=self._tracker.download_budget,
        )

    @property
    def model_config(self) -> ModelConfig:
        return self._tracker.model_config

    async def _item_with_uuid_async(self, item: object, uuid: str | None):
        return item, uuid

    @override
    def parse_prompt_embeds(self, data: str) -> None:
        """Schedule async prompt embeds decode and store the coroutine in the tracker.

        Like the sync variant, emits a single sentinel `PROMPT_EMBEDS_PLACEHOLDER_TOKEN`
        per content part. Unlike the sync variant, the tensor decode is deferred to a
        thread-pool executor via `safe_load_prompt_embeds_async`.
        """
        if not self.model_config.enable_prompt_embeds:
            raise VLLMValidationError(
                _ENABLE_PROMPT_EMBEDS_ERROR, parameter="prompt_embeds"
            )

        self._tracker.add(
            "prompt_embeds", partial(self._load_prompt_embeds_async, data.encode())
        )
        self._add_placeholder("prompt_embeds", PROMPT_EMBEDS_PLACEHOLDER_TOKEN)

    async def _load_prompt_embeds_async(
        self, data_bytes: bytes
    ) -> tuple[torch.Tensor, None]:
        # Second tuple slot fills the tracker's generic `(item, uuid | None)`
        # contract. prompt_embeds has no UUID concept, so it's always `None`.
        tensor = await safe_load_prompt_embeds_async(self.model_config, data_bytes)
        return tensor, None

    async def _image_with_uuid_async(self, image_url: str | None, uuid: str | None):
        image = (
            await self._connector.fetch_image_async(image_url) if image_url else None
        )
        return image, uuid

    def parse_image(self, image_url: str | None, uuid: str | None = None) -> None:
        placeholder = self._tracker.add(
            "image", partial(self._image_with_uuid_async, image_url, uuid)
        )
        self._add_placeholder("image", placeholder)

    def parse_image_embeds(
        self,
        image_embeds: MultiModalEmbedsPayload | None,
        uuid: str | None = None,
    ) -> None:
        mm_config = self.model_config.get_multimodal_config()
        if not mm_config.enable_mm_embeds:
            raise VLLMValidationError(
                "You must set `--enable-mm-embeds` to input `image_embeds`",
                parameter="image_embeds",
            )

        placeholder = self._tracker.add(
            "image_embeds",
            partial(self._image_embeds_with_uuid_async, image_embeds, uuid),
        )
        self._add_placeholder("image", placeholder)

    async def _image_embeds_with_uuid_async(
        self,
        image_embeds: MultiModalEmbedsPayload | None,
        uuid: str | None,
    ):
        if isinstance(image_embeds, dict):
            embeds = await _load_embeds_dict(
                image_embeds, self._connector.fetch_image_embedding_async
            )
        elif isinstance(image_embeds, str):
            embeds = await self._connector.fetch_image_embedding_async(image_embeds)
        else:
            embeds = None
        return embeds, uuid

    def parse_audio_embeds(
        self,
        audio_embeds: MultiModalEmbedsPayload | None,
        uuid: str | None = None,
    ) -> None:
        mm_config = self.model_config.get_multimodal_config()
        if not mm_config.enable_mm_embeds:
            raise VLLMValidationError(
                "You must set `--enable-mm-embeds` to input `audio_embeds`",
                parameter="audio_embeds",
            )

        placeholder = self._tracker.add(
            "audio_embeds",
            partial(self._audio_embeds_with_uuid_async, audio_embeds, uuid),
        )
        self._add_placeholder("audio", placeholder)

    async def _audio_embeds_with_uuid_async(
        self,
        audio_embeds: MultiModalEmbedsPayload | None,
        uuid: str | None,
    ):
        if isinstance(audio_embeds, dict):
            embeds = await _load_embeds_dict(
                audio_embeds, self._connector.fetch_audio_embedding_async
            )
        elif isinstance(audio_embeds, str):
            embeds = await self._connector.fetch_audio_embedding_async(audio_embeds)
        else:
            embeds = None
        return embeds, uuid

    def parse_image_pil(
        self,
        image_pil: Image.Image | None,
        uuid: str | None = None,
    ) -> None:
        if image_pil is not None:
            self._tracker.decoded_budget.charge(
                image_pil.size[0] * image_pil.size[1] * 4
            )
        placeholder = self._tracker.add(
            "image", partial(self._item_with_uuid_async, image_pil, uuid)
        )
        self._add_placeholder("image", placeholder)

    async def _audio_with_uuid_async(self, audio_url: str | None, uuid: str | None):
        audio = (
            await self._connector.fetch_audio_async(audio_url) if audio_url else None
        )
        return audio, uuid

    def parse_audio(self, audio_url: str | None, uuid: str | None = None) -> None:
        placeholder = self._tracker.add(
            "audio", partial(self._audio_with_uuid_async, audio_url, uuid)
        )
        self._add_placeholder("audio", placeholder)

    def parse_input_audio(
        self, input_audio: InputAudio | None, uuid: str | None = None
    ) -> None:
        if input_audio:
            audio_data = input_audio.get("data", "")
            audio_format = input_audio.get("format", "")
            if audio_data:
                audio_url = f"data:audio/{audio_format};base64,{audio_data}"
            else:
                # If a UUID is provided, audio data may be empty.
                audio_url = None
        else:
            audio_url = None

        return self.parse_audio(audio_url, uuid)

    async def _video_with_uuid_async(self, video_url: str | None, uuid: str | None):
        video = (
            await self._connector.fetch_video_async(
                video_url,
                video_processor=self._tracker.video_processor_name,
            )
            if video_url
            else None
        )
        return video, uuid

    def parse_video(self, video_url: str | None, uuid: str | None = None) -> None:
        placeholder = self._tracker.add(
            "video", partial(self._video_with_uuid_async, video_url, uuid)
        )
        self._add_placeholder("video", placeholder)

        # Extract audio from video if use_audio_in_video is True
        if (
            video_url
            and self._mm_processor_kwargs
            and self._mm_processor_kwargs.get("use_audio_in_video", False)
        ):
            audio_placeholder = self._tracker.add(
                "audio", partial(self._audio_with_uuid_async, video_url, uuid)
            )
            self._add_placeholder("audio", audio_placeholder)

    def parse_video_embeds(
        self,
        video_embeds: MultiModalEmbedsPayload | None,
        uuid: str | None = None,
    ) -> None:
        mm_config = self.model_config.get_multimodal_config()
        if not mm_config.enable_mm_embeds:
            raise VLLMValidationError(
                "You must set `--enable-mm-embeds` to input `video_embeds`",
                parameter="video_embeds",
            )

        placeholder = self._tracker.add(
            "video_embeds",
            partial(self._video_embeds_with_uuid_async, video_embeds, uuid),
        )
        self._add_placeholder("video", placeholder)

    async def _video_embeds_with_uuid_async(
        self,
        video_embeds: MultiModalEmbedsPayload | None,
        uuid: str | None,
    ):
        if isinstance(video_embeds, dict):
            embeds = await _load_embeds_dict(
                video_embeds, self._connector.fetch_video_embedding_async
            )
        elif isinstance(video_embeds, str):
            embeds = await self._connector.fetch_video_embedding_async(video_embeds)
        else:
            embeds = None
        return embeds, uuid


async def _load_embeds_dict(
    data: dict[str, str | list[int | float]],
    fetch: Callable[[str], Awaitable["torch.Tensor"]],
) -> dict[str, Any]:
    encoded = {key: value for key, value in data.items() if isinstance(value, str)}
    tensors = await asyncio.gather(*(fetch(value) for value in encoded.values()))
    return {**data, **dict(zip(encoded, tensors))}