"""Optional bounded generation for batch-size-one Transformers inference.

This is an inference safeguard, not a weight repair. A guard stop is incomplete,
never a successful answer. Time limits are cooperative, checked between tokens.
"""
import time
import torch
from transformers import StoppingCriteria, StoppingCriteriaList


def repeated_block(tokens, block_size=48, occurrences=3):
    """Detect a long suffix repeated in generated text only, not in the prompt."""
    if len(tokens) < block_size * occurrences:
        return False
    suffix = tokens[-block_size:]
    hits = 1
    end = len(tokens) - block_size
    while end >= block_size:
        found = False
        for start in range(end - block_size, -1, -1):
            if tokens[start:start + block_size] == suffix:
                hits += 1
                end = start
                found = True
                break
        if hits >= occurrences:
            return True
        if not found:
            break
    return False


class GenerationGuard(StoppingCriteria):
    def __init__(self, prompt_length, max_seconds=45, check_every=16):
        self.prompt_length = prompt_length
        self.deadline = time.monotonic() + max_seconds
        self.check_every = check_every
        self.reason = None

    def __call__(self, input_ids, scores, **kwargs):
        if input_ids.shape[0] != 1:
            raise ValueError('Q36 GenerationGuard supports batch size one only')
        generated = input_ids.shape[1] - self.prompt_length
        if time.monotonic() >= self.deadline:
            self.reason = 'time_limit'
        elif generated >= 144 and generated % self.check_every == 0:
            tokens = input_ids[0, max(self.prompt_length, input_ids.shape[1]-1536):].tolist()
            if repeated_block(tokens):
                self.reason = 'repeated_block'
        return torch.tensor([self.reason is not None], device=input_ids.device)


def generate_bounded(model, tokenizer, inputs, max_new_tokens=1024, max_seconds=45):
    """Caller controls thinking via apply_chat_template(enable_thinking=False).

    Legitimate repeated text can trigger the heuristic; retain raw output and
    expose finish_reason to the caller. No automatic retry or tool execution.
    """
    if inputs['input_ids'].shape[0] != 1:
        raise ValueError('Only batch size one is supported')
    if max_new_tokens <= 0 or max_seconds <= 0:
        raise ValueError('Positive token and time budgets required')
    n = inputs['input_ids'].shape[1]
    guard = GenerationGuard(n, max_seconds)
    began = time.monotonic()
    with torch.inference_mode():
        output = model.generate(
            **inputs, max_new_tokens=max_new_tokens, do_sample=False,
            use_cache=True, stopping_criteria=StoppingCriteriaList([guard]),
        )
    tail = output[0, n:]
    eos = model.generation_config.eos_token_id
    eos = eos if isinstance(eos, list) else [eos]
    complete = bool(len(tail) and int(tail[-1]) in eos)
    return {
        'text': tokenizer.decode(tail, skip_special_tokens=True),
        'raw_output': tokenizer.decode(tail, skip_special_tokens=False),
        'completed': complete,
        'finish_reason': 'eos' if complete else guard.reason or 'token_limit',
        'output_tokens': len(tail), 'seconds': time.monotonic()-began,
    }