"""P0 scheduler: inject a pre-fixed per-request verification horizon."""

from __future__ import annotations

from typing import Any

from vllm.v1.core.sched.scheduler import Scheduler
from vllm.v1.outputs import DraftTokenIds

from .common import digest, histogram, log_event, plan_summary, requested_k


class P0Scheduler(Scheduler):
    """A Scheduler that changes only post-drafter candidate visibility.

    EAGLE has already produced its Kmax candidates when this method runs.  The
    P0 experiment therefore tests verifier-side ragged execution and collective
    phase behavior, not a drafter speedup.
    """

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

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

    def schedule(self, *args: Any, **kwargs: Any):  # type: ignore[no-untyped-def]
        output = super().schedule(*args, **kwargs)
        self._p0_schedule_epoch += 1
        log_event(
            "scheduler",
            "schedule",
            core_local_epoch=self._p0_schedule_epoch,
            **self._p0_rank_fields(),
            **plan_summary(output),
        )
        return output

    def update_draft_token_ids(self, draft_token_ids: DraftTokenIds) -> None:
        self._p0_candidate_epoch += 1
        before = [list(ids) for ids in draft_token_ids.draft_token_ids]
        chosen: list[int] = []
        after: list[list[int]] = []
        structured_output_request_ids: list[str] = []
        for request_id, ids in zip(draft_token_ids.req_ids, before):
            request = self.requests.get(request_id)
            if request is not None and getattr(request, "structured_output_request", None) is not None:
                # Do not make grammar behavior part of this premise experiment.
                chosen.append(len(ids))
                after.append(ids)
                structured_output_request_ids.append(str(request_id))
                continue
            k = requested_k(str(request_id), len(ids))
            chosen.append(k)
            after.append(ids[:k])
        pairs = [
            {"request_id": str(request_id), "before_k": len(ids), "after_k": after_k}
            for request_id, ids, after_k in zip(draft_token_ids.req_ids, before, chosen)
        ]
        log_event(
            "scheduler",
            "candidate_truncate",
            candidate_epoch=self._p0_candidate_epoch,
            **self._p0_rank_fields(),
            candidate_count=len(pairs),
            candidate_digest=digest(pairs),
            before_k_histogram=histogram([item["before_k"] for item in pairs]),
            after_k_histogram=histogram(chosen),
            structured_output_request_ids=structured_output_request_ids,
        )
        super().update_draft_token_ids(
            DraftTokenIds(req_ids=draft_token_ids.req_ids, draft_token_ids=after)
        )