# Copyright 2025 - Oumi
#
# 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.

from __future__ import annotations

import functools
import json
from typing import Any, NamedTuple

import pydantic
from typing_extensions import override

from oumi.builders import (
    build_processor,
    build_tokenizer,
    is_image_text_llm,
)
from oumi.core.configs import (
    GenerationParams,
    ModelParams,
    RemoteParams,
)
from oumi.core.configs.internal.supported_models import (
    find_internal_model_config_using_model_name,
)
from oumi.core.processors.base_processor import BaseProcessor
from oumi.core.types.conversation import Conversation, Message, Role, Type
from oumi.inference.remote_inference_engine import RemoteInferenceEngine
from oumi.utils.conversation_utils import (
    base64encode_content_item_image_bytes,
    load_image_bytes_to_content_item,
)


class _SamplingParams(NamedTuple):
    """It's a clone of `sglang.lang.ir.SglSamplingParams`.

    Only includes a subset of parameters supported in oumi.
    Unsupported params are left commented out for reference.
    """

    max_new_tokens: int = 128
    # min_new_tokens: int = 0
    stop: str | list[str] = ""
    stop_token_ids: list[int] | None = None
    temperature: float = 1.0
    top_p: float = 1.0
    # top_k: int = -1  # -1 means disable
    min_p: float = 0.0
    frequency_penalty: float = 0.0
    presence_penalty: float = 0.0
    # ignore_eos: bool = False
    # return_logprob: bool | None = None
    # logprob_start_len: int | None = None
    # top_logprobs_num: int | None = None
    # return_text_in_logprobs: bool | None = None
    json_schema: str | None = None

    # For constrained generation:
    # dtype: str | None = None
    regex: str | None = None



[docs]
class SGLangInferenceEngine(RemoteInferenceEngine):
    """Engine for running SGLang inference."""

    def __init__(
        self,
        model_params: ModelParams,
        *,
        remote_params: RemoteParams | None = None,
        generation_params: GenerationParams | None = None,
    ):
        """Initializes the SGL inference Engine.

        Args:
            model_params: The model parameters to use for inference.
            remote_params: Remote server params.
            generation_params: The generation parameters to use for inference.
        """
        if remote_params is None:
            raise ValueError("remote_params is required")

        super().__init__(
            model_params=model_params,
            generation_params=generation_params,
            remote_params=remote_params,
        )

        self._tokenizer = build_tokenizer(self._model_params)
        self._processor: BaseProcessor | None = None
        self._supports_multiple_images: bool = False

        if is_image_text_llm(self._model_params):
            # Only enable Processor for vision language models for now.
            self._processor = build_processor(
                self._model_params.model_name,
                self._tokenizer,
                trust_remote_code=self._model_params.trust_remote_code,
                processor_kwargs=self._model_params.processor_kwargs,
            )
            internal_model_config = find_internal_model_config_using_model_name(
                self._model_params.model_name,
                trust_remote_code=self._model_params.trust_remote_code,
            )
            self._supports_multiple_images = (
                (internal_model_config is not None)
                and (internal_model_config.visual_config is not None)
                and internal_model_config.visual_config.supports_multiple_images
            )

        # TODO Launch a local SGLLang server if requested.

    def _create_sampling_params(
        self, generation_params: GenerationParams
    ) -> _SamplingParams:
        regex: str | None = None
        json_schema: str | None = None
        if generation_params.guided_decoding is not None:
            if generation_params.guided_decoding.regex is not None:
                regex = generation_params.guided_decoding.regex
            else:
                json_schema_value = None
                if generation_params.guided_decoding.json is not None:
                    json_schema_value = generation_params.guided_decoding.json
                elif (
                    generation_params.guided_decoding.choice is not None
                    and len(generation_params.guided_decoding.choice) > 0
                ):
                    json_schema_value = {
                        "enum": generation_params.guided_decoding.choice
                    }

                if isinstance(json_schema_value, str):
                    json_schema = json_schema_value
                elif isinstance(json_schema_value, dict):
                    json_schema = json.dumps(json_schema_value, ensure_ascii=False)
                elif isinstance(json_schema_value, pydantic.BaseModel) or (
                    isinstance(json_schema_value, type)
                    and issubclass(json_schema_value, pydantic.BaseModel)
                ):
                    json_schema = json.dumps(json_schema_value.model_json_schema())
                else:
                    raise ValueError(
                        "Unsupported type of generation_params.guided_decoding.json: "
                        f"{type(generation_params.guided_decoding.json)}"
                    )

        return _SamplingParams(
            max_new_tokens=generation_params.max_new_tokens,
            temperature=generation_params.temperature,
            top_p=generation_params.top_p,
            min_p=generation_params.min_p,
            frequency_penalty=generation_params.frequency_penalty,
            presence_penalty=generation_params.presence_penalty,
            stop=(generation_params.stop_strings or []),
            stop_token_ids=generation_params.stop_token_ids,
            regex=regex,
            json_schema=json_schema,
        )

    def _create_sampling_params_as_dict(
        self, generation_params: GenerationParams
    ) -> dict[str, Any]:
        return self._create_sampling_params(generation_params)._asdict()

    def _apply_chat_template_impl(self, conversation: Conversation) -> str:
        if self._processor is None:
            return self._tokenizer.apply_chat_template(
                conversation,  # type: ignore
                tokenize=False,
                add_generation_prompt=True,
            )
        return self._processor.apply_chat_template(
            conversation,  # type: ignore
            add_generation_prompt=True,
        )

    def _create_image_data_as_str_list(self, conversation: Conversation) -> list[str]:
        image_items = [
            item for m in conversation.messages for item in m.image_content_items
        ]
        num_images = len(image_items)
        if num_images <= 0:
            return []

        max_images = num_images if self._supports_multiple_images else 1

        if num_images > max_images:
            # If a conversation contains too many images, raise an error.
            # We can't silently discard extra images at this point
            # as many models verify that the actual number of images matches
            # the number of image tokens in text prompt.
            raise ValueError(
                conversation.append_id_to_string(
                    f"A conversation contains too many images ({num_images}). "
                    f"Max {max_images} image is allowed."
                )
            )

        result: list[str] = []
        for idx, image_item in enumerate(image_items):
            if image_item.type == Type.IMAGE_URL:
                # Preserve URL-s: leave them to SGLang server to download
                # to keep message payload size under control.
                # TODO Consider making this behaviour configurable.
                image_url = image_item.content
                if not image_url:
                    raise ValueError(
                        conversation.append_id_to_string(
                            f"Empty image URL in message: {image_item.type} "
                            f"in image item {idx + 1} of {num_images}!"
                        )
                    )
                result.append(image_url)
            else:
                image_item = load_image_bytes_to_content_item(image_item)
                if image_item.binary is None or len(image_item.binary) == 0:
                    raise ValueError(
                        conversation.append_id_to_string(
                            f"No image bytes in image item {idx + 1} of {num_images}!"
                        )
                    )
                result.append(base64encode_content_item_image_bytes(image_item))

        return result

    @override
    def _convert_conversation_to_api_input(
        self,
        conversation: Conversation,
        generation_params: GenerationParams,
        model_params: ModelParams,
    ) -> dict[str, Any]:
        """Converts a conversation to SGLang Native API input.

        See https://sgl-project.github.io/references/sampling_params.html for details.

        Args:
            conversation: The Oumi Conversation object to convert.
            generation_params: Parameters for text generation.
            model_params: Ignored.

        Returns:
            Dict[str, Any]: A dictionary containing the formatted input for the
            SGLang server native API, including the model, messages, generation params.
        """
        # Chat templates loaded by SGLang server are generally different from Oumi's
        # chat templates, hence, let's apply Oumi chat template here ourselves.
        prompt = self._apply_chat_template_impl(conversation)

        sampling_params_dict = self._create_sampling_params_as_dict(generation_params)
        body = {
            "text": prompt,
            "sampling_params": sampling_params_dict,
        }
        image_data: list[str] = self._create_image_data_as_str_list(conversation)
        if len(image_data) > 0:
            body["image_data"] = image_data if len(image_data) > 1 else image_data[0]
        return body

    @override
    def _convert_api_output_to_conversation(
        self, response: dict[str, Any], original_conversation: Conversation
    ) -> Conversation:
        """Converts an SGLang Native API response to a conversation."""
        new_message = Message(
            content=response["text"],
            role=Role.ASSISTANT,
        )
        return Conversation(
            messages=[*original_conversation.messages, new_message],
            metadata=original_conversation.metadata,
            conversation_id=original_conversation.conversation_id,
        )

    @override
    def _get_request_headers(self, remote_params: RemoteParams) -> dict[str, str]:
        return {
            "Content-Type": "application/json",
        }


[docs]
    @override
    @functools.cache
    def get_supported_params(self) -> set[str]:
        """Returns a set of supported generation parameters for this engine."""
        return {
            "frequency_penalty",
            "guided_decoding",
            "max_new_tokens",
            "min_p",
            "presence_penalty",
            "stop_strings",
            "stop_token_ids",
            "temperature",
            "top_p",
        }