#
# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved.
# This file is a part of the vllm-ascend project.
#
# 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.
#
# MiniMax-M2 usage accounting: backport reasoning-token usage details.
#

from __future__ import annotations

import json
from collections.abc import AsyncIterator, Sequence
from dataclasses import dataclass
from types import MethodType
from typing import Any

from vllm.entrypoints.openai.chat_completion import protocol as chat_protocol
from vllm.entrypoints.openai.chat_completion import serving as chat_serving
from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat
from vllm.entrypoints.openai.engine import protocol as engine_protocol
from vllm.reasoning import minimax_m2_reasoning_parser as minimax_parser

_MINIMAX_REASONING_PARSER_TYPES = (
    minimax_parser.MiniMaxM2ReasoningParser,
    minimax_parser.MiniMaxM2AppendThinkReasoningParser,
)


class CompletionTokenUsageInfo(engine_protocol.OpenAIBaseModel):
    reasoning_tokens: int | None = None
    audio_tokens: int | None = None
    accepted_prediction_tokens: int | None = None
    rejected_prediction_tokens: int | None = None


class UsageInfo(engine_protocol.UsageInfo):
    completion_tokens_details: CompletionTokenUsageInfo | None = None


CompletionTokenUsageInfo.__module__ = engine_protocol.__name__
UsageInfo.__module__ = engine_protocol.__name__

# The OpenAI usage schema is process-wide. Keep only this schema backfill
# global; the expensive token tracking below is bound to MiniMax instances.
engine_protocol.CompletionTokenUsageInfo = CompletionTokenUsageInfo
engine_protocol.UsageInfo = UsageInfo
chat_protocol.UsageInfo = UsageInfo
chat_serving.CompletionTokenUsageInfo = CompletionTokenUsageInfo
chat_serving.UsageInfo = UsageInfo


def _rebuild_model_field(model_cls, field_name: str, annotation) -> None:
    model_cls.__annotations__[field_name] = annotation
    model_cls.model_fields[field_name].annotation = annotation
    model_cls.model_rebuild(force=True)


_rebuild_model_field(chat_protocol.ChatCompletionResponse, "usage", UsageInfo)
_rebuild_model_field(chat_protocol.ChatCompletionStreamResponse, "usage", UsageInfo | None)
_rebuild_model_field(engine_protocol.RequestResponseMetadata, "final_usage_info", UsageInfo | None)


def _count_minimax_reasoning_tokens(
    token_ids: Sequence[int],
    end_token_id: int | None,
) -> int:
    if end_token_id is None:
        return 0

    for idx, token_id in enumerate(token_ids):
        if token_id == end_token_id:
            return idx
    return len(token_ids)


def _patched_count_reasoning_tokens(self, token_ids: Sequence[int]) -> int:
    return _count_minimax_reasoning_tokens(token_ids, self.end_token_id)


minimax_parser.MiniMaxM2ReasoningParser.count_reasoning_tokens = _patched_count_reasoning_tokens
minimax_parser.MiniMaxM2AppendThinkReasoningParser.count_reasoning_tokens = _patched_count_reasoning_tokens


def _count_minimax_reasoning_tokens_for_usage(
    token_ids: Sequence[int],
    reasoning_parser,
) -> int | None:
    reasoning_parser = _resolve_reasoning_parser(reasoning_parser)
    if reasoning_parser is None or not _is_minimax_reasoning_parser(reasoning_parser):
        return None

    count_reasoning_tokens = getattr(reasoning_parser, "count_reasoning_tokens", None)
    if count_reasoning_tokens is None:
        return None
    return count_reasoning_tokens(token_ids)


def _resolve_reasoning_parser(reasoning_parser):
    if reasoning_parser is None:
        return None
    return getattr(reasoning_parser, "reasoning_parser", reasoning_parser)


def _is_minimax_reasoning_parser(reasoning_parser) -> bool:
    return isinstance(
        _resolve_reasoning_parser(reasoning_parser),
        _MINIMAX_REASONING_PARSER_TYPES,
    )


def _clamp_reasoning_tokens(
    reasoning_tokens: int | None,
    completion_tokens: int,
) -> int | None:
    if reasoning_tokens is None:
        return None
    return max(0, min(reasoning_tokens, completion_tokens))


def _make_usage_info(
    self,
    *,
    prompt_tokens: int,
    completion_tokens: int,
    num_cached_tokens: int | None = None,
    reasoning_tokens: int | None = None,
) -> UsageInfo:
    usage = UsageInfo(
        prompt_tokens=prompt_tokens,
        completion_tokens=completion_tokens,
        total_tokens=prompt_tokens + completion_tokens,
    )
    reasoning_tokens = _clamp_reasoning_tokens(reasoning_tokens, completion_tokens)
    if reasoning_tokens is not None:
        usage.completion_tokens_details = CompletionTokenUsageInfo(reasoning_tokens=reasoning_tokens)
    if self.enable_prompt_tokens_details and num_cached_tokens is not None:
        usage.prompt_tokens_details = chat_serving.PromptTokenUsageInfo(cached_tokens=num_cached_tokens)
    return usage


def _is_minimax_reasoning_parser_cls(reasoning_parser_cls) -> bool:
    return isinstance(reasoning_parser_cls, type) and issubclass(
        reasoning_parser_cls,
        _MINIMAX_REASONING_PARSER_TYPES,
    )


@dataclass
class _UsageTrackingState:
    completion_tokens: list[int]
    raw_output_token_ids: list[list[int]]
    reasoning_parser: Any
    enable_prompt_tokens_details: bool = False
    num_prompt_tokens: int = 0
    num_cached_tokens: int | None = None
    final_res: Any = None


def _create_usage_tracking_state(
    num_choices: int,
    reasoning_parser,
    enable_prompt_tokens_details: bool = False,
) -> _UsageTrackingState:
    return _UsageTrackingState(
        completion_tokens=[0] * num_choices,
        raw_output_token_ids=[[] for _ in range(num_choices)],
        reasoning_parser=reasoning_parser,
        enable_prompt_tokens_details=enable_prompt_tokens_details,
    )


def _update_usage_tracking_state(
    state: _UsageTrackingState,
    res,
) -> None:
    if res.prompt_token_ids is not None:
        num_prompt_tokens = len(res.prompt_token_ids)
        if res.encoder_prompt_token_ids is not None:
            num_prompt_tokens += len(res.encoder_prompt_token_ids)
        state.num_prompt_tokens = num_prompt_tokens

    if state.num_cached_tokens is None:
        state.num_cached_tokens = res.num_cached_tokens

    state.final_res = res

    for output in res.outputs:
        if 0 <= output.index < len(state.completion_tokens):
            token_ids = chat_serving.as_list(output.token_ids)
            state.completion_tokens[output.index] += len(token_ids)
            state.raw_output_token_ids[output.index].extend(token_ids)


async def _tracked_result_generator(
    result_generator: AsyncIterator,
    state: _UsageTrackingState,
):
    async for res in result_generator:
        _update_usage_tracking_state(state, res)
        yield res


def _sum_reasoning_tokens_for_usage(
    raw_output_token_ids: list[list[int]],
    reasoning_parser,
) -> int | None:
    if reasoning_parser is None:
        return None
    reasoning_token_counts = [
        _count_minimax_reasoning_tokens_for_usage(token_ids, reasoning_parser) for token_ids in raw_output_token_ids
    ]
    if all(reasoning_tokens is None for reasoning_tokens in reasoning_token_counts):
        return None
    return sum(reasoning_tokens or 0 for reasoning_tokens in reasoning_token_counts)


def _reasoning_tokens_for_choice(
    state: _UsageTrackingState,
    choice_index: int,
) -> int | None:
    if state.reasoning_parser is None:
        return None
    if not 0 <= choice_index < len(state.raw_output_token_ids):
        return None
    return _count_minimax_reasoning_tokens_for_usage(
        state.raw_output_token_ids[choice_index],
        state.reasoning_parser,
    )


def _make_full_response_usage(
    self,
    state: _UsageTrackingState,
) -> UsageInfo | None:
    if state.final_res is None:
        return None

    return self._make_usage_info(
        prompt_tokens=state.num_prompt_tokens,
        completion_tokens=sum(state.completion_tokens),
        num_cached_tokens=state.num_cached_tokens,
        reasoning_tokens=_sum_reasoning_tokens_for_usage(
            state.raw_output_token_ids,
            state.reasoning_parser,
        ),
    )


def _usage_reasoning_tokens_for_stream_chunk(
    state: _UsageTrackingState,
    chunk: dict[str, Any],
    completion_tokens: int,
) -> int | None:
    if state.reasoning_parser is None:
        return None

    choices = chunk.get("choices") or []
    if choices:
        choice_index = choices[0].get("index", 0)
        reasoning_tokens = _reasoning_tokens_for_choice(state, choice_index)
    else:
        reasoning_tokens = _sum_reasoning_tokens_for_usage(
            state.raw_output_token_ids,
            state.reasoning_parser,
        )
    return _clamp_reasoning_tokens(reasoning_tokens, completion_tokens)


def _inject_stream_usage_details(
    data: str,
    state: _UsageTrackingState,
) -> str:
    prefix = "data: "
    suffix = "\n\n"
    if not data.startswith(prefix):
        return data

    payload = data[len(prefix) :]
    if payload.endswith(suffix):
        payload = payload[: -len(suffix)]
    if payload == "[DONE]":
        return data

    try:
        chunk = json.loads(payload)
    except json.JSONDecodeError:
        return data

    usage = chunk.get("usage")
    if not isinstance(usage, dict):
        return data

    updated_usage = False
    if state.enable_prompt_tokens_details and state.num_cached_tokens is not None:
        usage["prompt_tokens_details"] = {
            "cached_tokens": state.num_cached_tokens,
        }
        updated_usage = True

    completion_tokens = usage.get("completion_tokens") or 0
    reasoning_tokens = _usage_reasoning_tokens_for_stream_chunk(
        state,
        chunk,
        completion_tokens,
    )
    if reasoning_tokens is not None:
        usage["completion_tokens_details"] = {
            "reasoning_tokens": reasoning_tokens,
        }
        updated_usage = True

    if not updated_usage:
        return data
    return f"{prefix}{json.dumps(chunk, ensure_ascii=False)}{suffix}"


async def _wrapped_chat_completion_stream_generator(
    self,
    request: chat_protocol.ChatCompletionRequest,
    result_generator: AsyncIterator,
    request_id: str,
    model_name: str,
    conversation,
    tokenizer,
    request_metadata: engine_protocol.RequestResponseMetadata,
    reasoning_parser=None,
    **extra_kwargs: Any,
):
    original_stream_generator = self._ascend_original_chat_completion_stream_generator
    num_choices = 1 if request.n is None else request.n
    state = _create_usage_tracking_state(
        num_choices,
        reasoning_parser,
        enable_prompt_tokens_details=self.enable_prompt_tokens_details,
    )

    async for data in original_stream_generator(
        request,
        _tracked_result_generator(result_generator, state),
        request_id,
        model_name,
        conversation,
        tokenizer,
        request_metadata,
        reasoning_parser,
        **extra_kwargs,
    ):
        yield _inject_stream_usage_details(data, state)

    usage = _make_full_response_usage(self, state)
    if usage is not None:
        request_metadata.final_usage_info = usage


async def _wrapped_chat_completion_full_generator(
    self,
    request: chat_protocol.ChatCompletionRequest,
    result_generator: AsyncIterator,
    request_id: str,
    model_name: str,
    conversation,
    tokenizer,
    request_metadata: engine_protocol.RequestResponseMetadata,
    reasoning_parser=None,
):
    original_full_generator = self._ascend_original_chat_completion_full_generator
    num_choices = 1 if request.n is None else request.n
    state = _create_usage_tracking_state(
        num_choices,
        reasoning_parser,
        enable_prompt_tokens_details=self.enable_prompt_tokens_details,
    )

    response = await original_full_generator(
        request,
        _tracked_result_generator(result_generator, state),
        request_id,
        model_name,
        conversation,
        tokenizer,
        request_metadata,
        reasoning_parser,
    )

    if not isinstance(response, chat_protocol.ChatCompletionResponse):
        return response

    usage = _make_full_response_usage(self, state)
    if usage is None:
        return response

    response.usage = usage
    request_metadata.final_usage_info = usage
    return response


_wrapped_chat_completion_stream_generator.__module__ = OpenAIServingChat.__module__
_wrapped_chat_completion_stream_generator.__qualname__ = (
    f"{OpenAIServingChat.__qualname__}.chat_completion_stream_generator"
)
_wrapped_chat_completion_full_generator.__module__ = OpenAIServingChat.__module__
_wrapped_chat_completion_full_generator.__qualname__ = (
    f"{OpenAIServingChat.__qualname__}.chat_completion_full_generator"
)


def _should_patch_chat_usage_instance(self) -> bool:
    return _is_minimax_reasoning_parser_cls(self.reasoning_parser_cls)


def _patch_chat_usage_instance(self) -> None:
    if getattr(self, "_ascend_minimax_usage_patched", False):
        return
    self._make_usage_info = MethodType(_make_usage_info, self)
    self._ascend_original_chat_completion_stream_generator = MethodType(
        OpenAIServingChat.chat_completion_stream_generator,
        self,
    )
    self._ascend_original_chat_completion_full_generator = MethodType(
        OpenAIServingChat.chat_completion_full_generator,
        self,
    )
    self.chat_completion_stream_generator = MethodType(
        _wrapped_chat_completion_stream_generator,
        self,
    )
    self.chat_completion_full_generator = MethodType(
        _wrapped_chat_completion_full_generator,
        self,
    )
    self._ascend_minimax_usage_patched = True


class _ReasoningParserClsDescriptor:
    def __init__(self, default_value=None):
        self.default_value = default_value

    def __get__(self, instance, owner=None):
        if instance is None:
            return self.default_value
        return instance.__dict__.get("_ascend_reasoning_parser_cls", self.default_value)

    def __set__(self, instance, value) -> None:
        instance.__dict__["_ascend_reasoning_parser_cls"] = value
        if _is_minimax_reasoning_parser_cls(value):
            _patch_chat_usage_instance(instance)


_current_reasoning_parser_cls = OpenAIServingChat.__dict__.get("reasoning_parser_cls")
if not isinstance(_current_reasoning_parser_cls, _ReasoningParserClsDescriptor):
    OpenAIServingChat.reasoning_parser_cls = _ReasoningParserClsDescriptor(_current_reasoning_parser_cls)