# coding=utf-8
# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Classifier-Free Guidance (CFG) logits processor for vLLM v1.

Implements CFG by pairing conditional and unconditional requests in the same
batch.  The processor blends their logits before sampling so both requests
produce identical tokens.

Usage:
    Submit prompts in alternating pairs:
        [cond_0, uncond_0, cond_1, uncond_1, ...]

    Each request carries SamplingParams.extra_args:
        cond:   {"cfg_scale": 3.0, "cfg_role": "cond", "cfg_pair_id": "pair_0"}
        uncond: {"cfg_scale": 3.0, "cfg_role": "uncond", "cfg_pair_id": "pair_0"}

    Pass this processor to the vLLM engine:
        LLM(..., logits_processors=[CFGLogitsProcessor])
"""

from __future__ import annotations

import logging
from typing import Optional

import torch

from vllm.config import VllmConfig
from vllm.sampling_params import SamplingParams
from vllm.v1.sample.logits_processor import BatchUpdate, LogitsProcessor, MoveDirectionality

logger = logging.getLogger(__name__)


class CFGLogitsProcessor(LogitsProcessor):
    """Blend conditional + unconditional logits for classifier-free guidance.

    Pairs are matched by explicit ``cfg_pair_id`` in extra_args.  For each pair
    the processor computes:
        blended = uncond_logits + cfg_scale * (cond_logits - uncond_logits)
    and writes the result to *both* rows so the sampler picks the same token.

    On first instantiation in each worker process, this class also patches
    ``GPUModelRunner._sample`` to copy the conditional sampled token to the
    unconditional slot, guaranteeing identical token sequences.  This must
    live here (not in a separate monkey-patch module) because vLLM may
    ``spawn`` workers as fresh processes where main-process patches are lost.
    """

    _sample_patched = False

    @classmethod
    def validate_params(cls, params: SamplingParams) -> None:
        ea = params.extra_args
        if not ea:
            return
        role = ea.get("cfg_role")
        if role is not None and role not in ("cond", "uncond"):
            raise ValueError(f"cfg_role must be 'cond' or 'uncond', got '{role}'")
        scale = ea.get("cfg_scale")
        if scale is not None and (not isinstance(scale, (int, float)) or scale < 1.0):
            raise ValueError(f"cfg_scale must be >= 1.0, got {scale}")

    def __init__(
        self, vllm_config: VllmConfig, device: torch.device, is_pin_memory: bool
    ) -> None:
        self._info: dict[int, dict] = {}
        self._output_tokens: dict[int, list[int]] = {}
        self._pairs: list[tuple[int, int, float]] = []
        self._dirty = True
        self._ensure_sample_patched()

    @classmethod
    def _ensure_sample_patched(cls) -> None:
        """Patch ``GPUModelRunner._sample`` to sync tokens after sampling.

        Runs once per process (guarded by ``_sample_patched`` flag).
        Because vLLM may spawn workers via ``multiprocessing.Process``,
        main-process monkey-patches are invisible here -- so we patch
        from within ``__init__`` which vLLM calls in every worker.
        """
        if cls._sample_patched:
            return
        cls._sample_patched = True

        from vllm.v1.worker.gpu_model_runner import GPUModelRunner

        _orig_sample = GPUModelRunner._sample

        def _sample_with_cfg_sync(self, logits, spec_decode_metadata):
            sampler_output = _orig_sample(self, logits, spec_decode_metadata)
            for proc in self.input_batch.logitsprocs.all:
                if hasattr(proc, "_pairs") and proc._pairs:
                    for cond_idx, uncond_idx, _ in proc._pairs:
                        sampler_output.sampled_token_ids[uncond_idx] = (
                            sampler_output.sampled_token_ids[cond_idx]
                        )
                    break
            return sampler_output

        GPUModelRunner._sample = _sample_with_cfg_sync
        logger.info(
            "CFGLogitsProcessor: patched GPUModelRunner._sample "
            "for post-sampling token sync (pid=%d)", __import__("os").getpid()
        )

    def is_argmax_invariant(self) -> bool:
        return False


    def _reset(self) -> None:
        self._info.clear()
        self._output_tokens.clear()
        self._pairs.clear()
        self._dirty = True

    def update_state(self, batch_update: Optional[BatchUpdate]) -> None:
        if batch_update is None:
            return

        for idx in batch_update.removed:
            logger.debug("Removing idx=%d", idx)
            self._info.pop(idx, None)
            self._output_tokens.pop(idx, None)

        if not self._info and batch_update.added:
            self._reset()

        for idx, params, _, output_token_ids in batch_update.added:
            ea = params.extra_args if params else None
            logger.debug(
                "Adding idx=%d role=%s pair_id=%s output_len=%d",
                idx,
                ea.get("cfg_role") if ea else None,
                ea.get("cfg_pair_id") if ea else None,
                len(output_token_ids),
            )
            if ea and ea.get("cfg_role") in ("cond", "uncond"):
                self._info[idx] = {
                    "role": ea["cfg_role"],
                    "cfg_scale": float(ea.get("cfg_scale", 1.0)),
                    "pair_id": ea.get("cfg_pair_id"),
                }
                self._output_tokens[idx] = output_token_ids
            else:
                self._info.pop(idx, None)
                self._output_tokens.pop(idx, None)
        self._dirty = True

        if self._info:
            for adx, bdx, direction in batch_update.moved:
                logger.debug("Moving %d -> %d direction=%s", adx, bdx, direction)
                a_val = self._info.pop(adx, None)
                b_val = self._info.pop(bdx, None)
                a_tok = self._output_tokens.pop(adx, None)
                b_tok = self._output_tokens.pop(bdx, None)
                if a_val is not None:
                    self._info[bdx] = a_val
                if a_tok is not None:
                    self._output_tokens[bdx] = a_tok
                if direction == MoveDirectionality.SWAP:
                    if b_val is not None:
                        self._info[adx] = b_val
                    if b_tok is not None:
                        self._output_tokens[adx] = b_tok
            self._dirty = True

    def _rebuild_pairs(self) -> None:
        """Match cond/uncond by ``cfg_pair_id``."""
        by_pair: dict[str, dict[str, tuple[int, float]]] = {}
        for idx, info in self._info.items():
            pair_id = info.get("pair_id")
            if pair_id is None:
                continue
            by_pair.setdefault(pair_id, {})[info["role"]] = (
                idx,
                info["cfg_scale"],
            )

        self._pairs = [
            (roles["cond"][0], roles["uncond"][0], roles["cond"][1])
            for roles in by_pair.values()
            if "cond" in roles and "uncond" in roles
        ]
        self._dirty = False


    _apply_step = 0
    _LOG_EVERY = 200

    def apply(self, logits: torch.Tensor) -> torch.Tensor:
        if not self._info:
            return logits

        if self._dirty:
            self._rebuild_pairs()

        do_log = (CFGLogitsProcessor._apply_step % self._LOG_EVERY == 0
                  and self._pairs)
        CFGLogitsProcessor._apply_step += 1

        for i, (cond_idx, uncond_idx, cfg_scale) in enumerate(self._pairs):
            cond_toks = self._output_tokens.get(cond_idx)
            uncond_toks = self._output_tokens.get(uncond_idx)
            if cond_toks and uncond_toks and len(cond_toks) != len(uncond_toks):
                logger.debug(
                    "CFG pair (%d, %d) output length mismatch: %d vs %d",
                    cond_idx, uncond_idx, len(cond_toks), len(uncond_toks),
                )

            if do_log and i == 0:
                k = 5
                cond_top = torch.topk(logits[cond_idx], k)
                uncond_top = torch.topk(logits[uncond_idx], k)

            blended = logits[uncond_idx] + cfg_scale * (
                logits[cond_idx] - logits[uncond_idx]
            )
            logits[cond_idx] = blended
            logits[uncond_idx] = blended

            if do_log and i == 0:
                blended_top = torch.topk(blended, k)
                logger.warning(
                    "CFG probe step=%d scale=%.1f | "
                    "cond top%d: ids=%s vals=%s | "
                    "uncond top%d: ids=%s vals=%s | "
                    "blended top%d: ids=%s vals=%s",
                    CFGLogitsProcessor._apply_step - 1, cfg_scale,
                    k, cond_top.indices.tolist(),
                    [f"{v:.2f}" for v in cond_top.values.tolist()],
                    k, uncond_top.indices.tolist(),
                    [f"{v:.2f}" for v in uncond_top.values.tolist()],
                    k, blended_top.indices.tolist(),
                    [f"{v:.2f}" for v in blended_top.values.tolist()],
                )

        return logits