"""P0 worker: record actual rank-side target execution phases."""

from __future__ import annotations

from typing import Any

from vllm.v1.worker.gpu_worker import Worker

from .common import log_event, plan_summary


class P0Worker(Worker):
    """Log rank-local execution without modifying model or collective calls."""

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        super().__init__(*args, **kwargs)
        self._p0_execute_epoch = 0
        self._p0_batch_epoch = 0

    def init_device(self):  # type: ignore[no-untyped-def]
        result = super().init_device()
        model_runner = self.model_runner
        if getattr(model_runner, "_collectivespec_p0_wrapped", False):
            return result
        original = model_runner._determine_batch_execution_and_padding

        def traced_determine(*args: Any, **kwargs: Any):  # type: ignore[no-untyped-def]
            output = original(*args, **kwargs)
            self._p0_batch_epoch += 1
            scheduler_output = kwargs.get("scheduler_output")
            if scheduler_output is None and args:
                scheduler_output = args[0]
            cudagraph_mode, batch_desc, should_ubatch, rows_across_dp, _ = output
            log_event(
                "worker",
                "batch_execution_plan",
                batch_phase_epoch=self._p0_batch_epoch,
                **self._p0_rank_fields(),
                **plan_summary(scheduler_output),
                cudagraph_mode=getattr(cudagraph_mode, "value", str(cudagraph_mode)),
                physical_batch_rows=getattr(batch_desc, "num_tokens", None),
                should_ubatch=should_ubatch,
                rows_across_dp=rows_across_dp,
            )
            return output

        model_runner._determine_batch_execution_and_padding = traced_determine
        model_runner._collectivespec_p0_wrapped = True
        return result

    def _p0_rank_fields(self) -> dict[str, Any]:
        config = self.vllm_config.parallel_config
        return {
            "global_rank": getattr(self, "rank", None),
            "local_rank": getattr(self, "local_rank", None),
            "data_parallel_rank": getattr(config, "data_parallel_rank", None),
            "tensor_parallel_rank": getattr(config, "tensor_parallel_rank", None),
            "expert_parallel_size": getattr(config, "expert_parallel_size", None),
            "tensor_parallel_size": getattr(config, "tensor_parallel_size", None),
            "data_parallel_size": getattr(config, "data_parallel_size", None),
            "speculative_kmax": getattr(
                getattr(self, "speculative_config", None), "num_speculative_tokens", None
            ),
        }

    def execute_model(self, scheduler_output: Any):  # type: ignore[no-untyped-def]
        self._p0_execute_epoch += 1
        fields = {
            "worker_phase_epoch": self._p0_execute_epoch,
            **self._p0_rank_fields(),
            **plan_summary(scheduler_output),
        }
        log_event("worker", "target_execute_begin", **fields)
        output = super().execute_model(scheduler_output)
        log_event("worker", "target_execute_end", **fields)
        return output