# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

from __future__ import annotations

from abc import ABC, abstractmethod
from collections.abc import Callable
from concurrent.futures import Future
from dataclasses import replace
from typing import TYPE_CHECKING, Any

from vllm.utils.import_utils import resolve_obj_by_qualname
from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec
from vllm.v1.outputs import KVConnectorOutput

from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
from vllm_omni.diffusion.sched.interface import CachedRequestData

if TYPE_CHECKING:
    from vllm_omni.diffusion.request import OmniDiffusionRequest
    from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput
    from vllm_omni.diffusion.worker.utils import BaseRunnerOutput


class DiffusionExecutor(ABC):
    """Abstract base class for Diffusion executors."""

    uses_multiproc: bool = False

    @staticmethod
    def get_class(od_config: OmniDiffusionConfig) -> type[DiffusionExecutor]:
        executor_class: type[DiffusionExecutor]
        distributed_executor_backend = od_config.distributed_executor_backend
        # Mirror vLLM's `world_size == 1 -> "uni"` default
        # (vllm/config/parallel.py). A single-GPU diffusion deployment has
        # nothing to distribute: spawning a worker only adds MessageQueues,
        # /dev/shm output segments, and a second model load. Explicit "mp"
        # keeps process isolation and RPC timeouts. Multi-GPU still defaults
        # to "mp".
        if distributed_executor_backend is None:
            num_gpus = od_config.num_gpus or 1
            distributed_executor_backend = "uni" if num_gpus == 1 else "mp"

        if isinstance(distributed_executor_backend, type):
            if not issubclass(distributed_executor_backend, DiffusionExecutor):
                raise TypeError(
                    "distributed_executor_backend must be a subclass of "
                    f"DiffusionExecutor. Got {distributed_executor_backend}."
                )
            executor_class = distributed_executor_backend
        elif distributed_executor_backend == "ray":
            raise NotImplementedError("ray backend is not yet supported.")
        elif distributed_executor_backend == "mp":
            from vllm_omni.diffusion.executor.multiproc_executor import MultiprocDiffusionExecutor

            executor_class = MultiprocDiffusionExecutor
        elif distributed_executor_backend == "uni":
            from vllm_omni.diffusion.executor.uniproc_executor import UniProcDiffusionExecutor

            executor_class = UniProcDiffusionExecutor
        elif distributed_executor_backend == "external_launcher":
            raise NotImplementedError("external_launcher backend is not yet supported.")
        elif isinstance(distributed_executor_backend, str):
            try:
                executor_class = resolve_obj_by_qualname(distributed_executor_backend)
            except (ImportError, ValueError) as e:
                raise ValueError(
                    f"Failed to load executor backend '{distributed_executor_backend}'. "
                    f"Ensure it is a valid python path. Error: {e}"
                ) from e

            if not issubclass(executor_class, DiffusionExecutor):
                raise TypeError(
                    f"distributed_executor_backend must be a subclass of DiffusionExecutor. Got {executor_class}."
                )
        else:
            raise ValueError(f"Unknown distributed executor backend: {distributed_executor_backend}")
        return executor_class

    def __init__(self, od_config: OmniDiffusionConfig) -> None:
        self.od_config = od_config
        self._init_executor()

    @abstractmethod
    def _init_executor(self) -> None:
        """Initialize the executor (e.g., launch workers, setup IPC)."""
        pass

    @property
    @abstractmethod
    def is_dead(self) -> bool:
        """Whether the executor is shut down or has failed fatally."""
        pass

    @abstractmethod
    def execute_request(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRunnerOutput:
        """Execute request-mode work from a scheduler output."""
        pass

    @abstractmethod
    def execute_batch(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRunnerOutput:
        """Execute request-mode work through the request-batch path."""
        pass

    @abstractmethod
    def execute_step(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRunnerOutput:
        """Execute step-mode work from a scheduler output."""
        pass

    @abstractmethod
    def collective_rpc(
        self,
        method: str,
        timeout: float | None = None,
        args: tuple = (),
        kwargs: dict | None = None,
        unique_reply_rank: int | None = None,
        exec_all_ranks: bool = False,
    ) -> Any:
        """Execute a method on workers."""
        pass

    @abstractmethod
    def check_health(self) -> None:
        """Check if the executor and workers are healthy."""
        pass

    def register_failure_callback(self, callback: Callable[[], None]) -> None:
        """Register a callback invoked when the executor fatally fails.

        Executors without a background failure monitor can keep the default
        no-op implementation.
        """
        return None

    def wait_output_ready(self, async_output_id: str) -> Future[DiffusionOutput]:
        """Resolve deferred output; only asynchronous executors implement this."""
        raise NotImplementedError(f"{type(self).__name__} does not support asynchronous output")

    def get_kv_cache_specs(self) -> list[dict[str, KVCacheSpec]]:
        """Collect rank-local native specs after every Worker loads its model."""

        result = self.collective_rpc(
            "get_kv_cache_specs",
            unique_reply_rank=0,
            exec_all_ranks=True,
        )
        if not isinstance(result, list):
            raise TypeError(f"get_kv_cache_specs returned {type(result).__name__}, expected list")
        return result

    def determine_available_kv_memory(self, profile_requests: list[OmniDiffusionRequest]) -> list[int]:
        """Profile and collect the KV memory budget on every Worker rank."""

        result = self.collective_rpc(
            "determine_available_kv_memory",
            args=(profile_requests,),
            unique_reply_rank=0,
            exec_all_ranks=True,
        )
        if not isinstance(result, list) or not all(isinstance(value, int) for value in result):
            raise TypeError("determine_available_kv_memory must return list[int]")
        return result

    def set_kv_cache_configs(self, kv_cache_configs: list[KVCacheConfig], resolved_max_model_len: int) -> None:
        """Send rank-local configs and the resolved model length to all Workers."""

        # The default control-plane RPC mode executes on every rank and has
        # rank 0 return the gathered rank statuses, so failures on nonzero
        # ranks are not silently dropped.
        self.collective_rpc("set_kv_cache_configs", args=(kv_cache_configs, resolved_max_model_len))

    def remove_diffusion_kv_requests(self, request_ids: list[str | tuple[str, int]]) -> None:
        """Clear request rows on every Worker after Scheduler retirement."""

        unique_request_ids = list(dict.fromkeys(request_ids))
        if not unique_request_ids:
            return
        self.collective_rpc(
            "remove_diffusion_kv_requests",
            args=(unique_request_ids,),
        )

    def prepare_kv_for_forward(self, scheduler_output: DiffusionSchedulerOutput) -> KVConnectorOutput | None:
        if scheduler_output.kv_prefetch_connector_metadata is not None:
            current = replace(
                scheduler_output,
                kv_transfer_request_ids=scheduler_output.kv_transfer_request_ids
                - scheduler_output.kv_prefetch_request_ids,
                kv_prefetch_connector_metadata=None,
                kv_prefetch_request_ids=set(),
            )
            # First complete the current request on every rank. Only then
            # submit B, preventing Mooncake from coalescing its bytes with A.
            self.prepare_kv_for_forward(current)
            return self.prepare_kv_for_forward(
                replace(
                    current,
                    kv_connector_metadata=scheduler_output.kv_prefetch_connector_metadata,
                    kv_transfer_request_ids=scheduler_output.kv_prefetch_request_ids,
                    kv_required_request_ids=set(),
                    kv_finished_request_ids=set(),
                )
            )
        if (
            scheduler_output.kv_connector_metadata is None
            and not scheduler_output.kv_required_request_ids
            and not scheduler_output.kv_poll_only
        ):
            return None
        transfer_output = replace(
            scheduler_output,
            scheduled_new_reqs=[],
            scheduled_cached_reqs=CachedRequestData.make_empty(),
            kv_prefetch_job=None,
        )
        outputs: list[KVConnectorOutput] = self.collective_rpc(
            "prepare_kv_for_forward",
            args=(transfer_output,),
            unique_reply_rank=0,
            exec_all_ranks=True,
        )
        if len(outputs) != self.od_config.num_gpus or any(output.invalid_block_ids for output in outputs):
            # Missing ranks / invalid pages cannot establish safe ownership.
            raise RuntimeError("Diffusion KV receive failed on one or more ranks")
        if scheduler_output.kv_required_request_ids is not None:
            # Prefetch workers retain completion events until retirement.
            # Their cumulative snapshots can be intersected directly.
            finished = set.intersection(*(set(output.finished_recving or ()) for output in outputs))
            if not scheduler_output.kv_required_request_ids.issubset(finished):
                raise RuntimeError("Required diffusion KV receive did not complete on every rank")
            outputs[0].finished_recving = finished
            return outputs[0]
        completed = getattr(self, "_kv_receive_completed_ranks", {})
        for rank, output in enumerate(outputs):
            for request_id in output.finished_recving or ():
                completed.setdefault(request_id, set()).add(rank)
        finished = {request_id for request_id, ranks in completed.items() if len(ranks) == len(outputs)}
        for request_id in finished:
            del completed[request_id]
        self._kv_receive_completed_ranks = completed
        outputs[0].finished_recving = finished
        return outputs[0]

    @abstractmethod
    def shutdown(self) -> None:
        """Shutdown the executor and release resources."""
        pass