from __future__ import annotations

import math
from typing import Any

import torch
from transformers.generation.logits_process import LogitsProcessor


class ReasoningBudgetLogitsProcessor(LogitsProcessor):
    """Apply an optional reasoning-token budget for one sequence."""

    def __init__(
        self,
        tokenizer: Any,
        prompt_length: int,
        reasoning_budget: int,
        grace_fraction: float = 0.1,
        early_exit: str = ".\n</think>\n\n",
    ) -> None:
        if reasoning_budget < 1:
            raise ValueError("reasoning_budget must be positive")

        self.tokenizer = tokenizer
        self.prompt_length = prompt_length
        self.reasoning_budget = reasoning_budget
        self.hard_limit = reasoning_budget + max(
            1, math.ceil(reasoning_budget * grace_fraction)
        )
        self.reasoning_end_ids = tokenizer.encode(
            "</think>", add_special_tokens=False
        )
        self.early_exit_ids = tokenizer.encode(
            early_exit, add_special_tokens=False
        )
        self.forced_index: int | None = None
        self.done = False

    def _force(self, scores: torch.FloatTensor, token_id: int) -> torch.FloatTensor:
        scores.fill_(-float("inf"))
        scores[:, token_id] = 0
        return scores

    def __call__(
        self,
        input_ids: torch.LongTensor,
        scores: torch.FloatTensor,
    ) -> torch.FloatTensor:
        if self.done:
            return scores
        if input_ids.shape[0] != 1:
            raise ValueError("Reasoning budget control requires batch size 1")

        generated_ids = input_ids[0, self.prompt_length :].tolist()
        if self.forced_index is not None:
            self.forced_index += 1
            if self.forced_index >= len(self.early_exit_ids):
                self.done = True
                return scores
            return self._force(scores, self.early_exit_ids[self.forced_index])

        if generated_ids[-len(self.reasoning_end_ids) :] == self.reasoning_end_ids:
            self.done = True
            return scores

        generated_tokens = len(generated_ids)
        if generated_tokens < self.reasoning_budget:
            return scores

        ended_line = bool(
            generated_ids
            and "\n"
            in self.tokenizer.decode(
                [generated_ids[-1]],
                skip_special_tokens=False,
                clean_up_tokenization_spaces=False,
            )
        )
        if not ended_line and generated_tokens < self.hard_limit:
            return scores

        self.forced_index = 0
        return self._force(scores, self.early_exit_ids[0])