# SPDX-License-Identifier: Apache-2.0
"""vLLM reasoning parser for Aria v1's legacy think-tag format.

Aria v1 was trained with a prompt-prefilled ``<think>`` tag and a generated
``</think>`` delimiter before the final answer.  Gemma's tokenizer does not
represent these strings as single special tokens, so vLLM's built-in
``qwen3``/``deepseek_r1`` parsers cannot be used directly.

Usage:

    vllm serve xlr8harder/aria-gemma4-31b-v1 \
      --reasoning-parser-plugin ./aria_v1_reasoning_parser.py \
      --reasoning-parser aria_v1
"""

from __future__ import annotations

from collections.abc import Iterable, Sequence

from vllm.entrypoints.openai.engine.protocol import DeltaMessage
from vllm.reasoning.abs_reasoning_parsers import (
    ReasoningParser,
    ReasoningParserManager,
)


THINK_OPEN = "<think>"
THINK_CLOSE = "</think>"


def _find_subsequence(haystack: Sequence[int], needle: Sequence[int]) -> int:
    if not needle:
        return -1
    limit = len(haystack) - len(needle) + 1
    for i in range(max(0, limit)):
        if list(haystack[i : i + len(needle)]) == list(needle):
            return i
    return -1


def _suffix_prefix_len(text: str, prefix: str) -> int:
    max_len = min(len(text), len(prefix) - 1)
    for length in range(max_len, 0, -1):
        if text[-length:] == prefix[:length]:
            return length
    return 0


@ReasoningParserManager.register_module("aria_v1")
class AriaV1ReasoningParser(ReasoningParser):
    """Parser for outputs shaped as ``reasoning</think>answer``."""

    @property
    def reasoning_start_str(self) -> str | None:
        return THINK_OPEN

    @property
    def reasoning_end_str(self) -> str | None:
        return THINK_CLOSE

    def __init__(self, tokenizer, *args, **kwargs):
        super().__init__(tokenizer, *args, **kwargs)
        self.end_token_ids = tokenizer.encode(THINK_CLOSE, add_special_tokens=False)
        self.open_token_ids = tokenizer.encode(THINK_OPEN, add_special_tokens=False)
        self._streaming_reasoning_done = False
        self._streaming_buffer = ""
        self._streaming_started = False
        self._strip_next_content_delta = False

    def _strip_optional_open(self, text: str) -> str:
        text = text.lstrip()
        if text.startswith(THINK_OPEN):
            return text[len(THINK_OPEN) :].lstrip("\r\n ")
        return text

    def is_reasoning_end(self, input_ids: Sequence[int]) -> bool:
        return _find_subsequence(input_ids, self.end_token_ids) >= 0

    def is_reasoning_end_streaming(
        self, input_ids: Sequence[int], delta_ids: Iterable[int]
    ) -> bool:
        return self.is_reasoning_end(input_ids)

    def extract_content_ids(self, input_ids: list[int]) -> list[int]:
        end_index = _find_subsequence(input_ids, self.end_token_ids)
        if end_index < 0:
            return []
        return input_ids[end_index + len(self.end_token_ids) :]

    def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int:
        end_index = _find_subsequence(token_ids, self.end_token_ids)
        return len(token_ids) if end_index < 0 else end_index

    def extract_reasoning(self, model_output: str, request) -> tuple[str | None, str | None]:
        model_output = self._strip_optional_open(model_output)
        if THINK_CLOSE not in model_output:
            return model_output, None
        reasoning, _, content = model_output.partition(THINK_CLOSE)
        return reasoning.strip(), content.lstrip("\r\n ") or None

    def extract_reasoning_streaming(
        self,
        previous_text: str,
        current_text: str,
        delta_text: str,
        previous_token_ids: Sequence[int],
        current_token_ids: Sequence[int],
        delta_token_ids: Sequence[int],
    ) -> DeltaMessage | None:
        if self._streaming_reasoning_done:
            if self._strip_next_content_delta:
                delta_text = delta_text.lstrip("\r\n ")
                if delta_text:
                    self._strip_next_content_delta = False
                else:
                    return None
            return DeltaMessage(content=delta_text) if delta_text else None

        combined = self._streaming_buffer + delta_text
        if not self._streaming_started:
            combined = self._strip_optional_open(combined)
            self._streaming_started = True

        if THINK_CLOSE in combined:
            reasoning, _, content = combined.partition(THINK_CLOSE)
            self._streaming_buffer = ""
            self._streaming_reasoning_done = True
            stripped_content = content.lstrip("\r\n ")
            if not stripped_content:
                self._strip_next_content_delta = True
            return DeltaMessage(
                reasoning=reasoning if reasoning else None,
                content=stripped_content or None,
            )

        keep = _suffix_prefix_len(combined, THINK_CLOSE)
        if keep:
            emit = combined[:-keep]
            self._streaming_buffer = combined[-keep:]
        else:
            emit = combined
            self._streaming_buffer = ""

        return DeltaMessage(reasoning=emit) if emit else None