"""과도 반복(폭주) 롤아웃 조기 종료용 vLLM V1 커스텀 logits processor.

배경: long-CoT eval 실측에서 롤아웃의 ~35%가 20480 캡까지 도달하며, 상당수가
n-gram 반복 루프(어차피 reward에서 오답 처리됨)다. 이들이 가장 비싼(긴 컨텍스트)
디코드 토큰을 소모한다.

동작: 요청별로 생성 토큰 목록(vLLM이 라이브 참조로 제공)을 감시하다가
n-gram 반복 비율이 임계 이상이면 해당 요청의 로짓을 EOS만 남기고 -inf로 마스킹
→ 그 요청만 즉시 종료. **트리거 전에는 로짓을 전혀 건드리지 않으므로 정상
롤아웃의 출력은 비트 단위로 동일하다** (출력 불변 원칙 하에 폭주 커팅만 승인됨,
2026-07-09 결정). 트리거된 롤아웃은 짧게 잘린 채 수집되고 기존 repetition
필터/reward가 오답 처리한다.

사용법:
  1) 엔진: vllm_kwargs.logits_processors: ["ralo.vllm_repetition_abort:RepetitionEosLogitsProcessor"]
  2) 요청: SamplingParams.extra_args = {"repetition_abort": {
         "eos_token_id": <int>,        # 필수 — 강제할 EOS
         "ngram": 7,                    # n-gram 크기 (dapo 필터와 동일 계열)
         "threshold": 0.9,              # 반복 비율 임계 (필터 0.8보다 보수적 기본값 —
                                        #  출력이 실제로 잘리는 개입이므로)
         "min_tokens": 2048,            # 이 길이 전에는 검사하지 않음
         "check_interval": 512,         # 검사 주기 (토큰)
     }}
  extra_args에 repetition_abort가 없는 요청은 완전히 무시된다.

주의 (2026-07-16 실측): vLLM V1의 async scheduling(기본 자동 활성화)에서는 워커가
output_tok_ids에 실제 토큰 대신 -1 플레이스홀더를 채워 이 프로세서가 무력화된다.
반드시 엔진에 `async_scheduling=False`를 함께 전달할 것 (think_format·boxed_eos 공통).
"""

from typing import Optional, Sequence

import torch

try:
    from vllm.sampling_params import SamplingParams
    from vllm.v1.sample.logits_processor import BatchUpdate, LogitsProcessor
    from vllm.v1.sample.logits_processor.builtin import process_dict_updates
    _VLLM_OK = True
except ImportError:  # vLLM 미설치 환경(트레이너 전용 노드 등)에서도 임포트 가능하게
    _VLLM_OK = False
    LogitsProcessor = object  # type: ignore
    BatchUpdate = None  # type: ignore


def ngram_repetition_ratio(token_ids: Sequence[int], n: int) -> float:
    """반복 n-gram에 속하는 "위치"의 비율 (0..1).

    주의: dapo/sampler.py의 필터 함수와 공식이 다르다. 필터는 "2회 이상 나온 고유
    n-gram 개수 ÷ 위치 수"인데, 이 정의는 순수 반복 루프에서 오히려 0에 수렴한다
    (고유 n-gram이 몇 개 안 되므로) — 재보정 주석의 실측치(정상 max 0.58, 폭주
    0.975)와 부합하는 것은 여기 구현한 "반복 n-gram이 차지하는 위치 비율" 쪽이다.
    순수 루프 → ~1.0, 무반복 텍스트 → 0.0.
    """
    if n <= 0 or len(token_ids) < n * 2:
        return 0.0
    counts: dict = {}
    total = 0
    for i in range(len(token_ids) - n + 1):
        total += 1
        key = tuple(token_ids[i : i + n])
        counts[key] = counts.get(key, 0) + 1
    if total == 0:
        return 0.0
    repeated_positions = sum(c for c in counts.values() if c > 1)
    return repeated_positions / max(total, 1)


class _ReqState:
    __slots__ = ("out", "eos", "ngram", "threshold", "interval", "next_check", "triggered")

    def __init__(self, out, eos, ngram, threshold, min_tokens, interval):
        self.out = out              # vLLM이 넘겨주는 라이브 출력 토큰 리스트 참조
        self.eos = int(eos)
        self.ngram = int(ngram)
        self.threshold = float(threshold)
        self.interval = max(1, int(interval))
        self.next_check = max(1, int(min_tokens))
        self.triggered = False

    def should_trigger(self) -> bool:
        """검사 시점 도달 시 반복 비율 평가. 트리거되면 True (한 번만)."""
        if self.triggered:
            return False
        if len(self.out) < self.next_check:
            return False
        ratio = ngram_repetition_ratio(self.out, self.ngram)
        if ratio >= self.threshold:
            self.triggered = True
            return True
        self.next_check = len(self.out) + self.interval
        return False


class RepetitionEosLogitsProcessor(LogitsProcessor):
    """반복 폭주 감지 시 해당 요청 로짓을 EOS 단일화 — 다른 요청/트리거 전 요청은 무변경."""

    def __init__(self, vllm_config, device: torch.device, is_pin_memory: bool):
        if not _VLLM_OK:
            raise RuntimeError("vLLM V1 logits processor API를 찾을 수 없음")
        self.device = device
        self.pin_memory = is_pin_memory
        self.states: dict[int, _ReqState] = {}
        # 트리거된 (배치 인덱스, eos id) 마스크 텐서
        self._rows = self._tensor([], torch.int64)
        self._eos = self._tensor([], torch.int64)
        self._have_triggered = False

    def _tensor(self, data, dtype):
        return torch.tensor(data, device="cpu", dtype=dtype, pin_memory=self.pin_memory).to(
            device=self.device, non_blocking=True
        )

    def is_argmax_invariant(self) -> bool:
        return False  # 트리거 시 argmax를 EOS로 바꾼다

    @staticmethod
    def add_request(params: "SamplingParams", _prompt, output_tok_ids) -> Optional[_ReqState]:
        cfg = (getattr(params, "extra_args", None) or {}).get("repetition_abort")
        if not cfg or cfg.get("eos_token_id") is None:
            return None
        return _ReqState(
            out=output_tok_ids,
            eos=cfg["eos_token_id"],
            ngram=cfg.get("ngram", 7),
            threshold=cfg.get("threshold", 0.9),
            min_tokens=cfg.get("min_tokens", 2048),
            interval=cfg.get("check_interval", 512),
        )

    def update_state(self, batch_update: "BatchUpdate | None") -> None:
        changed = process_dict_updates(self.states, batch_update, self.add_request)

        newly = False
        for state in self.states.values():
            if state.should_trigger():
                newly = True

        if changed or newly:
            rows, eos = [], []
            for idx, state in self.states.items():
                if state.triggered:
                    rows.append(idx)
                    eos.append(state.eos)
            self._rows = self._tensor(rows, torch.int64)
            self._eos = self._tensor(eos, torch.int64)
            self._have_triggered = bool(rows)
            if newly:
                try:
                    print(f"[RepetitionAbort] {len(rows)} request(s) forced to EOS "
                          f"(runaway n-gram repetition)", flush=True)
                except Exception:
                    pass

    def apply(self, logits: torch.Tensor) -> torch.Tensor:
        if self._have_triggered and self._rows.numel() > 0:
            # 트리거 행 전체 -inf 후 EOS만 0 → softmax에서 EOS 확률 1
            logits[self._rows] = float("-inf")
            logits[self._rows, self._eos] = 0.0
        return logits


def build_repetition_abort_extra_args(algo_cfg: dict, eos_token_id) -> Optional[dict]:
    """dapo_kwargs/lrs_kwargs의 repetition_abort 블록 → SamplingParams.extra_args.

    블록이 없거나 enabled가 아니면 None (프로세서가 해당 요청을 완전히 무시).
    """
    cfg = (algo_cfg or {}).get("repetition_abort") or {}
    if not cfg.get("enabled", False) or eos_token_id is None:
        return None
    return {
        "repetition_abort": {
            "eos_token_id": int(eos_token_id),
            "ngram": int(cfg.get("ngram", 7)),
            "threshold": float(cfg.get("threshold", 0.9)),
            "min_tokens": int(cfg.get("min_tokens", 2048)),
            "check_interval": int(cfg.get("check_interval", 512)),
        }
    }