# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# 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.
import re
from typing import Any, cast

import httpx
from harbor.llms.base import (
    BaseLLM,
    ContextLengthExceededError,
    LLMResponse,
    OutputLengthExceededError,
)
from harbor.models.metric import UsageInfo
from tenacity import (
    retry,
    retry_if_exception_type,
    retry_if_not_exception_type,
    stop_after_attempt,
    wait_exponential,
)

from nemo_gym.openai_utils import NeMoGymResponseCreateParamsNonStreaming


_RoutedExperts = list[list[list[int]]]
_RolloutDetailsKey = tuple[tuple[int, ...], tuple[int, ...], tuple[float, ...] | None]


# Phrases in vLLM / OpenAI error bodies that signal context-length overflow.
_CONTEXT_LENGTH_ERROR_PHRASES = (
    "context length exceeded",
    "context_length_exceeded",
    "maximum context length",
    "`inputs` tokens + `max_new_tokens`",
)

_THINK_OPEN = "<think>"
_THINK_CLOSE = "</think>"
_THINK_PATTERN = re.compile(r"<think>(.*?)</think>", re.DOTALL)


class NemoGymLLM(BaseLLM):
    """LLM backend that calls NeMo Gym model servers via chat completions."""

    def __init__(
        self,
        model_name: str,
        api_base: str,
        collect_rollout_details: bool = False,
        model_info: dict[str, Any] | None = None,
        responses_create_params: dict[str, Any] | None = None,
        timeout_sec: float = 600.0,
        **kwargs: Any,
    ) -> None:
        super().__init__(**kwargs)
        self._model_name = model_name
        self._api_base = api_base.rstrip("/")
        self._collect_rollout_details = collect_rollout_details
        self._model_info = model_info or {}
        self._timeout_sec = timeout_sec

        # Accumulated token IDs from the most recent turn, used for
        # on-policy correction via _replace_prefix_tokens in vLLM.
        self._last_prompt_token_ids: list[int] | None = None
        self._last_completion_token_ids: list[int] | None = None
        self._last_logprobs: list[float] | None = None
        self._last_routed_experts: _RoutedExperts | None = None
        self._routed_experts_by_rollout_details: dict[_RolloutDetailsKey, _RoutedExperts] = {}
        self._ambiguous_routed_expert_keys: set[_RolloutDetailsKey] = set()

        # Set when the model hits the context length limit.
        self.context_length_exceeded = False

        # Pre-compute extra chat params from responses_create_params once,
        # since they don't change between calls.
        self._extra_chat_params = self._build_extra_chat_params(responses_create_params or {})

    @retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=4, max=15),
        retry=(
            retry_if_exception_type(Exception)
            & retry_if_not_exception_type(
                (
                    ContextLengthExceededError,
                    OutputLengthExceededError,
                )
            )
        ),
        reraise=True,
    )
    async def call(
        self,
        prompt: str,
        message_history: list[dict[str, Any]] | None = None,
        **kwargs: Any,
    ) -> LLMResponse:
        if message_history is None:
            message_history = []
        messages = message_history + [{"role": "user", "content": prompt}]

        # Attach token IDs from the previous turn to the last assistant
        # message so vLLM can perform on-policy correction via
        # _replace_prefix_tokens (see NeMoRLOpenAIChatRequestMixin).
        if self._last_prompt_token_ids is not None:
            for msg in reversed(messages):
                if msg.get("role") == "assistant":
                    msg["prompt_token_ids"] = self._last_prompt_token_ids
                    msg["generation_token_ids"] = self._last_completion_token_ids or []
                    msg["generation_log_probs"] = self._last_logprobs or []
                    if self._last_routed_experts is not None:
                        msg["routed_experts"] = self._last_routed_experts
                    break

        payload: dict[str, Any] = {
            "model": self._model_name,
            "messages": messages,
        }
        payload.update(self._extra_chat_params)

        response_dict = await self._post_chat_completions(payload)

        # Detect silently-swallowed context-length errors from the Gym proxy.
        # When vLLM returns 400 "maximum context length", the proxy catches it
        # and returns a fake 200 with id="chtcmpl-123" and content=None.
        if response_dict.get("id") == "chtcmpl-123":
            self.context_length_exceeded = True
            raise ContextLengthExceededError(
                f"Model {self._model_name} context length exceeded (detected fake response id='chtcmpl-123')"
            )

        choices = response_dict.get("choices", [])
        choice = choices[0] if isinstance(choices, list) and choices else {}
        message = choice.get("message", {}) if isinstance(choice, dict) else {}
        content = message.get("content", "") if isinstance(message, dict) else ""
        if content is None:
            content = ""
        reasoning_content = message.get("reasoning_content") if isinstance(message, dict) else None

        # Extract reasoning from the response content.  There are two cases:
        #
        # 1. Content has matched open+close tags (e.g. "<think>rc</think>text"):
        #    vllm_model app.py wraps reasoning this way when uses_reasoning_parser is true.
        #    We mirror vllm_model app.py's _parse_think_tags exactly: findall + sub to
        #    strip all <think> blocks, but only keep the FIRST match as reasoning_content.
        #    No .strip() — preserve whitespace so round-tripping is lossless.
        #
        # 2. Content has only a close tag (e.g. "rc</think>text"):
        #    The open tag was in the generation prompt (e.g. nano-v3 appends
        #    <think>\n to every prompt), so the model's output starts mid-think.
        if reasoning_content is None and isinstance(content, str):
            if _THINK_OPEN in content:
                # Case 1: matched open+close tags.
                matches = _THINK_PATTERN.findall(content)
                remaining = _THINK_PATTERN.sub("", content)
                if matches:
                    if remaining:
                        reasoning_content = matches[0]
                        content = remaining
                    else:
                        # Entire output classified as reasoning — model didn't
                        # generate the close tag.  Treat as content so the agent
                        # can act on it; leave reasoning_content None so the
                        # merge won't inject a close tag that was never generated
                        # (which would break token contiguity).
                        content = matches[0]
                        reasoning_content = None
            elif _THINK_CLOSE in content:
                # Case 2: unmatched close tag — open tag was in the generation
                # prompt (e.g. nanov3 appends <think>\n), so the model's output
                # starts mid-think.  Split on the first close tag.
                parts = content.split(_THINK_CLOSE, 1)
                reasoning_content = parts[0]
                content = parts[1] if len(parts) > 1 else ""

        if isinstance(choice, dict) and choice.get("finish_reason") == "length":
            raise OutputLengthExceededError(
                f"Model {self._model_name} hit max_tokens limit. "
                "Response was truncated. Consider increasing max_tokens if possible.",
                truncated_response=content,
            )

        usage = self._extract_usage_info(response_dict)
        prompt_token_ids = None
        completion_token_ids = None
        logprobs = None
        if self._collect_rollout_details:
            prompt_token_ids, completion_token_ids = self._extract_token_ids(response_dict)
            logprobs = self._extract_logprobs(response_dict)
            routed_experts = self._extract_routed_experts(response_dict)
            # Store for on-policy correction on the next turn.
            self._last_prompt_token_ids = prompt_token_ids
            self._last_completion_token_ids = completion_token_ids
            self._last_logprobs = logprobs
            self._last_routed_experts = routed_experts
            self._store_routed_experts_for_rollout_details(
                prompt_token_ids,
                completion_token_ids,
                logprobs,
                routed_experts,
            )

        return LLMResponse(
            content=content,
            reasoning_content=reasoning_content,
            usage=usage,
            prompt_token_ids=prompt_token_ids,
            completion_token_ids=completion_token_ids,
            logprobs=logprobs,
        )

    def get_model_context_limit(self) -> int:
        fallback_context_limit = 1000000

        try:
            max_input_tokens = self._model_info.get("max_input_tokens")
            if max_input_tokens is None:
                max_input_tokens = self._model_info.get("max_tokens")

            if isinstance(max_input_tokens, int) and max_input_tokens > 0:
                return max_input_tokens

            self._logger.warning(
                f"Model '{self._model_name}' info found but missing context limit fields. "
                f"Using fallback context limit: {fallback_context_limit}"
            )
        except Exception as e:
            self._logger.warning(
                f"Failed to retrieve model info for '{self._model_name}': {e}. "
                f"Using fallback context limit: {fallback_context_limit}"
            )

        return fallback_context_limit

    def get_model_output_limit(self) -> int | None:
        try:
            max_output_tokens = self._model_info.get("max_output_tokens")

            if max_output_tokens is None:
                self._logger.debug(f"Model '{self._model_name}' info found but missing max_output_tokens field.")

            if isinstance(max_output_tokens, int) and max_output_tokens > 0:
                return max_output_tokens

            return None
        except Exception as e:
            self._logger.debug(f"Failed to retrieve model info for '{self._model_name}': {e}.")
            return None

    async def _post_chat_completions(
        self, payload: dict[str, Any], timeout_sec: float | None = None
    ) -> dict[str, Any]:
        endpoint = self._chat_completions_endpoint()
        timeout = timeout_sec if timeout_sec is not None else self._timeout_sec
        async with httpx.AsyncClient(timeout=timeout) as client:
            response = await client.post(endpoint, json=payload)

        if response.status_code >= 400:
            error_text = response.text.lower()
            if any(phrase in error_text for phrase in _CONTEXT_LENGTH_ERROR_PHRASES):
                self.context_length_exceeded = True
                raise ContextLengthExceededError(f"Model {self._model_name} context length exceeded: {response.text}")
            response.raise_for_status()

        return response.json()

    def _chat_completions_endpoint(self) -> str:
        if self._api_base.endswith("/v1"):
            return f"{self._api_base}/chat/completions"
        return f"{self._api_base}/v1/chat/completions"

    def _extract_token_ids(self, response: dict[str, Any]) -> tuple[list[int] | None, list[int] | None]:
        choices = response.get("choices", [])
        choice = choices[0] if isinstance(choices, list) and choices else {}
        message = choice.get("message", {}) if isinstance(choice, dict) else {}

        prompt_token_ids = message.get("prompt_token_ids") if isinstance(message, dict) else None
        if prompt_token_ids is None:
            prompt_token_ids = response.get("prompt_token_ids")

        completion_token_ids = message.get("generation_token_ids") if isinstance(message, dict) else None

        return (
            self._normalize_token_ids(prompt_token_ids),
            self._normalize_token_ids(completion_token_ids),
        )

    def _build_extra_chat_params(self, responses_create_params: dict[str, Any]) -> dict[str, Any]:
        if not responses_create_params:
            return {}

        from responses_api_models.vllm_model.app import VLLMConverter

        params_for_conversion = {key: value for key, value in responses_create_params.items() if key != "input"}
        params_for_conversion["input"] = []
        responses_params = NeMoGymResponseCreateParamsNonStreaming.model_validate(params_for_conversion)

        converter = VLLMConverter(
            return_token_id_information=self._collect_rollout_details,
        )
        chat_params = converter.responses_to_chat_completion_create_params(responses_params).model_dump(
            exclude_unset=True
        )

        chat_params.pop("messages", None)
        return chat_params

    def _extract_logprobs(self, response: dict[str, Any]) -> list[float] | None:
        choices = response.get("choices", [])
        if not isinstance(choices, list) or not choices:
            return None

        choice = choices[0]
        if not isinstance(choice, dict):
            return None

        message = choice.get("message", {})
        if isinstance(message, dict):
            generation_log_probs = message.get("generation_log_probs")
            if isinstance(generation_log_probs, list):
                return [float(lp) for lp in generation_log_probs if isinstance(lp, (int, float))] or None

        logprobs_data = choice.get("logprobs")
        if isinstance(logprobs_data, dict):
            content = logprobs_data.get("content", [])
            extracted = [
                token_data["logprob"]
                for token_data in content
                if isinstance(token_data, dict) and "logprob" in token_data
            ]
            if extracted:
                return extracted

        return None

    def _extract_routed_experts(self, response: dict[str, Any]) -> _RoutedExperts | None:
        choices = response.get("choices", [])
        choice = choices[0] if isinstance(choices, list) and choices else {}
        message = choice.get("message", {}) if isinstance(choice, dict) else {}
        routed_experts = message.get("routed_experts") if isinstance(message, dict) else None
        if routed_experts is None:
            routed_experts = response.get("routed_experts")
        if not isinstance(routed_experts, list):
            return None
        return cast(_RoutedExperts, routed_experts)

    def _store_routed_experts_for_rollout_details(
        self,
        prompt_token_ids: list[int] | None,
        completion_token_ids: list[int] | None,
        logprobs: list[float] | None,
        routed_experts: _RoutedExperts | None,
    ) -> None:
        if routed_experts is None:
            return

        key = self._rollout_details_key(prompt_token_ids, completion_token_ids, logprobs)
        if key is None:
            return

        if key in self._routed_experts_by_rollout_details:
            self._ambiguous_routed_expert_keys.add(key)
            self._routed_experts_by_rollout_details.pop(key, None)
            return

        if key not in self._ambiguous_routed_expert_keys:
            self._routed_experts_by_rollout_details[key] = routed_experts

    def pop_routed_experts_for_rollout_details(
        self,
        prompt_token_ids: list[int] | None,
        completion_token_ids: list[int] | None,
        logprobs: list[float] | None,
    ) -> _RoutedExperts | None:
        key = self._rollout_details_key(prompt_token_ids, completion_token_ids, logprobs)
        if key is None or key in self._ambiguous_routed_expert_keys:
            return None
        return self._routed_experts_by_rollout_details.pop(key, None)

    @staticmethod
    def _rollout_details_key(
        prompt_token_ids: list[int] | None,
        completion_token_ids: list[int] | None,
        logprobs: list[float] | None,
    ) -> _RolloutDetailsKey | None:
        if prompt_token_ids is None or completion_token_ids is None:
            return None

        logprobs_key = tuple(logprobs) if logprobs is not None else None
        return (tuple(prompt_token_ids), tuple(completion_token_ids), logprobs_key)

    def _extract_usage_info(self, response: dict[str, Any]) -> UsageInfo | None:
        usage = response.get("usage")
        if not isinstance(usage, dict):
            return None

        prompt_tokens = usage.get("prompt_tokens", 0) or 0
        completion_tokens = usage.get("completion_tokens", 0) or 0
        prompt_tokens_details = usage.get("prompt_tokens_details") or {}
        cache_tokens = (
            prompt_tokens_details.get("cached_tokens", 0) if isinstance(prompt_tokens_details, dict) else 0
        ) or 0

        return UsageInfo(
            prompt_tokens=int(prompt_tokens),
            completion_tokens=int(completion_tokens),
            cache_tokens=int(cache_tokens),
            cost_usd=0.0,
        )

    def _normalize_token_ids(self, token_ids: Any) -> list[int] | None:
        if not isinstance(token_ids, list):
            return None

        normalized: list[int] = []
        for token_id in token_ids:
            if isinstance(token_id, int):
                normalized.append(token_id)
                continue
            if isinstance(token_id, str):
                stripped = token_id.removeprefix("token_id:")
                if stripped.isdigit():
                    normalized.append(int(stripped))
                    continue
            return None

        return normalized or None
