# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

import copy
import logging
import uuid
from typing import Any, Optional

import torch
from vllm.inputs.data import TokensPrompt
from vllm.v1.engine.async_llm import AsyncLLM

from dynamo.common.memory.multimodal_embedding_cache_manager import (
    MultimodalEmbeddingCacheManager,
)
from dynamo.common.multimodal.embedding_transfer import (
    LocalEmbeddingReceiver,
    NixlReadEmbeddingReceiver,
    NixlWriteEmbeddingReceiver,
)
from dynamo.common.utils import nvtx_utils as _nvtx
from dynamo.common.utils.otel_tracing import build_trace_headers
from dynamo.common.utils.time_section import time_and_log_code_section
from dynamo.runtime import Client, DistributedRuntime

from ..args import Config
from ..constants import DisaggregationMode, EmbeddingTransferMode
from ..handlers import BaseWorkerHandler, build_sampling_params
from ..multimodal_utils import (
    MyRequestOutput,
    PatchedTokensPrompt,
    vLLMMultimodalRequest,
)
from ..multimodal_utils.model import is_qwen_vl_model
from ..multimodal_utils.prefill_worker_utils import MultiModalEmbeddingLoader

logger = logging.getLogger(__name__)

IMAGE_URL_KEY = "image_url"


class MultimodalPDWorkerHandler(BaseWorkerHandler[dict, dict]):
    """Prefill/Decode or Prefill-only worker for multimodal serving"""

    def __init__(
        self,
        runtime,
        engine_client: AsyncLLM,
        config: Config,
        encode_worker_client: Optional[Client] = None,
        decode_worker_client: Optional[Client] = None,
        shutdown_event=None,
        generate_endpoint=None,
    ):
        # Get default_sampling_params from config
        default_sampling_params = (
            config.engine_args.create_model_config().get_diff_sampling_param()
        )

        # Call BaseWorkerHandler.__init__ with proper parameters
        super().__init__(
            runtime,
            config,
            engine_client,
            default_sampling_params,
            enable_multimodal=config.enable_multimodal,
            generate_endpoint=generate_endpoint,
            shutdown_event=shutdown_event,
        )

        self.config = config
        self.decode_worker_client = decode_worker_client
        self.enable_disagg = config.disaggregation_mode == DisaggregationMode.PREFILL

        # Initialize multimodal-specific components
        logger.info("Multimodal PD Worker startup started.")

        # Embedding loader consist of two main components:
        # 1) An remote encode worker client and matching embedding receiver,
        #    which can request remote encode and handle the transfer of embeddings
        #    from the encode worker to this prefill worker.
        # 2) A local embedding cache manager, which can store previously fetched embeddings
        #    and used to determine whether remote encode is necessary for a given mm data.
        self.encode_worker_client = encode_worker_client  # type: ignore
        if config.embedding_transfer_mode == EmbeddingTransferMode.LOCAL:
            self.embedding_receiver = LocalEmbeddingReceiver()  # type: ignore
        elif config.embedding_transfer_mode == EmbeddingTransferMode.NIXL_WRITE:
            self.embedding_receiver = NixlWriteEmbeddingReceiver()  # type: ignore
        elif config.embedding_transfer_mode == EmbeddingTransferMode.NIXL_READ:
            # [gluo FIXME] can't use pre-registered tensor as NIXL requires descriptors
            # to be at matching size, need to overwrite nixl connect library
            self.embedding_receiver = NixlReadEmbeddingReceiver(max_items=0)  # type: ignore
        else:
            raise ValueError(
                f"Invalid embedding transfer mode: {config.embedding_transfer_mode}"
            )
        self.embedding_cache_manager: MultimodalEmbeddingCacheManager | None = None
        if config.multimodal_embedding_cache_capacity_gb > 0:
            capacity_bytes = int(
                config.multimodal_embedding_cache_capacity_gb * 1024**3
            )
            self.embedding_cache_manager = MultimodalEmbeddingCacheManager(
                capacity_bytes
            )
        self.embedding_loader: MultiModalEmbeddingLoader = MultiModalEmbeddingLoader(
            encode_worker_client=self.encode_worker_client,  # type: ignore
            receiver=self.embedding_receiver,
            embedding_cache_manager=self.embedding_cache_manager,
        )

        logger.info("Multimodal PD Worker has been initialized")

    async def async_init(self, runtime: DistributedRuntime):
        """Async initialization for connector that requires async setup"""
        logger.info("Multimodal PD Worker async initialization completed.")

    def _parse_frontend_request(
        self, raw_request: dict
    ) -> tuple[vLLMMultimodalRequest, list[str]]:
        """Parse a raw frontend dict into a vLLMMultimodalRequest and image URLs.

        The Rust frontend sends a dict with ``token_ids`` and
        ``multi_modal_data`` (containing image URLs). This method extracts
        those fields into a structured request. No I/O is performed here;
        embedding fetching is handled separately by ``_load_multimodal_data``.
        """
        request_id = str(uuid.uuid4().hex)

        image_urls: list[str] = []
        mm_data = raw_request.get("multi_modal_data")
        if mm_data is not None:
            for item in mm_data.get(IMAGE_URL_KEY, []):
                if isinstance(item, dict) and "Url" in item:
                    image_urls.append(item["Url"])
                elif isinstance(item, dict) and "Decoded" in item:
                    image_urls.append(item["Decoded"])

        sampling_params = build_sampling_params(
            raw_request, self.default_sampling_params
        )

        request = vLLMMultimodalRequest(
            engine_prompt=PatchedTokensPrompt(
                prompt_token_ids=raw_request["token_ids"]
            ),
            sampling_params=sampling_params,
            request_id=request_id,
            model=raw_request.get("model"),
        )

        return request, image_urls

    # ── Multimodal data loading ──────────────────────────────────────

    async def _load_multimodal_data(
        self, image_urls: list[str], request_id: str, context=None
    ) -> dict[str, Any]:
        """Fetch embeddings from encode workers and load into an engine-ready dict.

        Returns an empty dict when no encode worker is configured or no images
        are present.
        """

        return await self.embedding_loader.load_multimodal_embeddings(
            image_urls,
            request_id,
            model=self.config.model,
            context=context,
        )

    # ── Request metadata finalization ────────────────────────────────

    def _finalize_request_metadata(
        self,
        request: vLLMMultimodalRequest,
        multi_modal_data: dict[str, Any],
    ) -> None:
        """Attach model-specific metadata to the request for the decode worker.

        For Qwen VL (mRoPE) models, captures image grid dimensions and
        embedding shapes so the decode worker can reconstruct
        ``multi_modal_data`` consistently for multiple images.
        """
        if is_qwen_vl_model(self.config.model) and isinstance(
            multi_modal_data.get("image"), dict
        ):
            image_data = multi_modal_data["image"]
            image_grid_thw = image_data.get("image_grid_thw")
            image_embeds = image_data.get("image_embeds")
            if image_grid_thw is not None:
                request.image_grid_thw = (
                    image_grid_thw.tolist()
                    if isinstance(image_grid_thw, torch.Tensor)
                    else image_grid_thw
                )
            if image_embeds is not None:
                request.embeddings_shape = list(image_embeds.shape)
        # prune empty multimodal data, vLLM will expect multi_modal_uuids if the mm items are empty
        # i.e. ValueError: multi_modal_data['image'] is empty but multi_modal_uuids['image'] is missing.
        for key, value in multi_modal_data.items():
            if not isinstance(value, torch.Tensor):
                if not value:
                    del multi_modal_data[key]
                else:
                    logger.debug(
                        f"Prepared multimodal data key {key}, number of items: {len(multi_modal_data[key])}"
                    )

        logger.debug("Multimodal data keys: %s", list(multi_modal_data.keys()))

    @staticmethod
    def _format_engine_output(
        response, num_output_tokens_so_far: int
    ) -> dict[str, Any]:
        """Format a vLLM RequestOutput as an LLMEngineOutput-compatible dict.

        This produces the same incremental dict format that the regular
        (non-multimodal) handler yields, which the Rust frontend expects
        after model registration.
        """
        if not response.outputs:
            return {
                "finish_reason": "error: No outputs from vLLM engine",
                "token_ids": [],
            }

        output = response.outputs[0]
        out: dict[str, Any] = {
            "token_ids": output.token_ids[num_output_tokens_so_far:],
        }

        if output.finish_reason:
            # Inline normalization: map vLLM's "abort" to Dynamo's "cancelled"
            finish_reason = output.finish_reason
            if finish_reason.startswith("abort"):
                finish_reason = "cancelled"
            out["finish_reason"] = finish_reason
            out["completion_usage"] = BaseWorkerHandler._build_completion_usage(
                request_output=response,
            )
        if output.stop_reason:
            out["stop_reason"] = output.stop_reason

        return out

    # ── Aggregated generation (prefill + decode locally) ─────────────

    async def _generate_agg(
        self,
        request: vLLMMultimodalRequest,
        multi_modal_data: dict[str, Any],
        rng_ttft=None,
        context=None,
    ):
        """Run prefill and decode on this worker (aggregated mode)."""
        lora_request = self._resolve_lora_request(request.model)
        trace_headers = build_trace_headers(context) if context else None
        gen = self.engine_client.generate(
            prompt=TokensPrompt(
                prompt_token_ids=request.engine_prompt["prompt_token_ids"],
                multi_modal_data=multi_modal_data,
            ),
            sampling_params=request.sampling_params,
            request_id=request.request_id,
            lora_request=lora_request,
            trace_headers=trace_headers,
        )

        num_output_tokens_so_far = 0
        first_token = True
        try:
            async for response in gen:
                if first_token:
                    if rng_ttft is not None:
                        _nvtx.end_range(rng_ttft)
                    first_token = False
                logger.debug(
                    f"Response kv_transfer_params: {response.kv_transfer_params}"
                )
                logger.debug(
                    f"length of expanded prompt ids: {len(response.prompt_token_ids)}"
                )
                yield self._format_engine_output(response, num_output_tokens_so_far)
                if response.outputs:
                    num_output_tokens_so_far = len(response.outputs[0].token_ids)
        finally:
            if first_token:
                if rng_ttft is not None:
                    _nvtx.end_range(rng_ttft)

    # ── Disaggregated generation (prefill here, decode remote) ───────

    async def _generate_disagg(
        self,
        request: vLLMMultimodalRequest,
        multi_modal_data: dict[str, Any],
        rng_ttft=None,
        context=None,
    ):
        """Prefill locally, then forward to a remote decode worker."""
        with _nvtx.annotate(
            "mm:pd:disagg_prefill", color="darkred"
        ), time_and_log_code_section(
            f"[PREFILL] request: {request.request_id} prefill time"
        ):
            # Prepare prefill-only request
            prefill_only_request = copy.deepcopy(request)
            extra_args = prefill_only_request.sampling_params.extra_args or {}
            extra_args["kv_transfer_params"] = {"do_remote_decode": True}
            prefill_only_request.sampling_params.extra_args = extra_args
            prefill_only_request.sampling_params.max_tokens = 1
            prefill_only_request.sampling_params.min_tokens = 1
            logger.debug("Prefill request: %s", prefill_only_request)

            lora_request = self._resolve_lora_request(request.model)
            trace_headers = build_trace_headers(context) if context else None
            gen = self.engine_client.generate(
                prompt=TokensPrompt(
                    prompt_token_ids=prefill_only_request.engine_prompt[
                        "prompt_token_ids"
                    ],
                    multi_modal_data=multi_modal_data,
                ),
                sampling_params=prefill_only_request.sampling_params,
                request_id=prefill_only_request.request_id,
                lora_request=lora_request,
                trace_headers=trace_headers,
            )

            # Drain prefill generator (max_tokens=1, expect a single response)
            async for prefill_response in gen:
                pass
        if rng_ttft is not None:
            _nvtx.end_range(rng_ttft)

        # Qwen VL (mRoPE): keep the ORIGINAL unexpanded prompt.
        # The decode worker passes multi_modal_data which causes vLLM to
        # expand the prompt identically to prefill, ensuring block counts match.
        #
        # Other models: use the expanded prompt from prefill response.
        # They don't pass multi_modal_data in decode, so they need the
        # already-expanded prompt to match the KV cache layout.
        if not is_qwen_vl_model(self.config.model):
            request.engine_prompt[
                "prompt_token_ids"
            ] = prefill_response.prompt_token_ids

        logger.debug(
            f"Prefill response kv_transfer_params: {prefill_response.kv_transfer_params}"
        )
        extra_args = request.sampling_params.extra_args or {}
        extra_args["kv_transfer_params"] = prefill_response.kv_transfer_params
        extra_args.pop("serialized_request", None)
        request.sampling_params.extra_args = extra_args
        logger.debug("Decode request: %s", request)

        # Serialized request is lightweight: token IDs, sampling params with
        # kv_transfer_params, and small Qwen metadata (image_grid_thw,
        # embeddings_shape).  Heavy multimodal data was consumed locally by
        # engine_client.generate() and multimodal_inputs was cleared by
        # `_finalize_request_metadata`.
        #
        # request.model (LoRA name) is preserved in the serialized request
        # so the decode worker can resolve the same LoRA adapter.
        if lora_request and request.model:
            logger.debug(
                f"Forwarding disaggregated decode with LoRA '{request.model}' "
                f"— ensure the same adapter is loaded on the decode worker."
            )

        with (
            _nvtx.annotate("mm:pd:disagg_remote_decode", color="purple"),
            time_and_log_code_section(
                f"[PREFILL] request: {request.request_id} remote decode time"
            ) as decode_timer,
        ):
            num_output_tokens_so_far = 0
            if self.decode_worker_client is None:
                raise RuntimeError("Decode worker client is not configured.")
            async for (decode_response) in await self.decode_worker_client.round_robin(
                request.model_dump_json(), context=context
            ):
                output = MyRequestOutput.model_validate_json(decode_response.data())
                yield self._format_engine_output(output, num_output_tokens_so_far)
                if output.outputs:
                    if num_output_tokens_so_far == 0:
                        decode_timer.stop_interval()  # Log time to first decode response
                    num_output_tokens_so_far = len(output.outputs[0].token_ids)

    # ── Public entry point ───────────────────────────────────────────

    async def generate(self, raw_request: dict, context):
        """Parse the request, load multimodal data, and run inference."""
        rng_pd = _nvtx.start_range("mm:pd_worker_generate", color="green")
        rng_ttft = _nvtx.start_range("mm:pd:ttft", color="orange")

        with time_and_log_code_section("[REQUEST] embedding processing time"):
            rng_parse = _nvtx.start_range("mm:pd:parse_request", color="cyan")
            request, image_urls = self._parse_frontend_request(raw_request)
            logger.debug(f"Received PD request: {{ id: {request.request_id} }}.")
            _nvtx.end_range(rng_parse)

            rng_load = _nvtx.start_range("mm:pd:load_multimodal", color="yellow")
            multi_modal_data = await self._load_multimodal_data(
                image_urls, request.request_id, context
            )
            _nvtx.end_range(rng_load)

            self._finalize_request_metadata(request, multi_modal_data)

        if self.enable_disagg and self.decode_worker_client:
            rng_disagg = _nvtx.start_range("mm:pd:generate_disagg", color="red")
            async for chunk in self._generate_disagg(
                request, multi_modal_data, rng_ttft, context=context
            ):
                yield chunk
            _nvtx.end_range(rng_disagg)
        else:
            rng_agg = _nvtx.start_range("mm:pd:generate_agg", color="red")
            async for chunk in self._generate_agg(
                request, multi_modal_data, rng_ttft, context=context
            ):
                yield chunk
            _nvtx.end_range(rng_agg)

        _nvtx.end_range(rng_pd)
