"""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

    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