# Copyright 2025 SGLang Team
# 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.
# ==============================================================================
# Adapted from https://huggingface.co/nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16/blob/cb5a65ff10232128389d882d805fa609427544f1/configuration.py

from typing import Any

from transformers.configuration_utils import PretrainedConfig

from sglang.srt.configs.nemotron_h import NemotronHConfig
from sglang.srt.configs.radio import RadioConfig
from sglang.srt.multimodal.internvl_utils import IMAGENET_MEAN, IMAGENET_STD


def float_triplet(seq: Any):
    a, b, c = tuple(seq)
    assert (
        isinstance(a, float) and isinstance(b, float) and isinstance(c, float)
    ), "expected three floats"
    return a, b, c


class NemotronH_Nano_VL_V2_Config(PretrainedConfig):
    model_type = "NemotronH_Nano_VL_V2"
    is_composition = True

    def __init__(
        self,
        vision_config=None,
        llm_config=None,
        sound_config=None,
        force_image_size: int = 512,
        patch_size: int = 16,
        downsample_ratio=0.5,
        template=None,
        ps_version="v2",
        image_tag_type="internvl",
        projector_hidden_size=4096,
        vit_hidden_size=1280,
        video_pruning_rate: float = 0.0,
        video_context_token: str = "<video>",
        img_context_token: str = "<image>",
        img_start_token: str = "<img>",
        img_end_token: str = "</img>",
        audio_context_token: str = "<so_embedding>",
        audio_start_token: str = "<so_start>",
        audio_end_token: str = "<so_end>",
        norm_mean: tuple[float, float, float] | list[float] = IMAGENET_MEAN,
        norm_std: tuple[float, float, float] | list[float] = IMAGENET_STD,
        use_thumbnail: bool = True,
        **kwargs,
    ):
        # Round-trip: `to_dict()` emits `raw_vision_config` (V2's storage
        # name) but `from_dict()` rebuilds via this `vision_config` kwarg.
        # Without this alias, the V3->V2 alias rebuild in `get_config` loses
        # the vision config across the round-trip.
        if vision_config is None:
            vision_config = kwargs.pop("raw_vision_config", None)

        super().__init__(**kwargs)

        # Handle both cases: when loading from JSON (llm_config is dict) and when called internally by transformers (llm_config; vision_config are None)
        if llm_config is not None:
            self.llm_config = NemotronHConfig(**llm_config)
            assert isinstance(vision_config, dict), "vision_config must be a dictionary"
            self.raw_vision_config = vision_config
        else:
            assert vision_config is None
            self.llm_config = NemotronHConfig()
            self.raw_vision_config = {}

        # Audio (Parakeet) config: stored as a PretrainedConfig sub-object
        if sound_config is not None and isinstance(sound_config, dict):
            self.sound_config = PretrainedConfig.from_dict(sound_config)
        else:
            self.sound_config = sound_config

        # Assign configuration values
        vision_image_size = self.raw_vision_config.get("image_size", force_image_size)
        vision_patch_size = self.raw_vision_config.get("patch_size", patch_size)
        self.image_size = int(
            vision_image_size[0]
            if isinstance(vision_image_size, list)
            else vision_image_size
        )
        self.patch_size = int(
            vision_patch_size[0]
            if isinstance(vision_patch_size, list)
            else vision_patch_size
        )

        self.downsample_ratio = downsample_ratio
        self.video_context_token = video_context_token
        self.img_context_token = img_context_token
        self.template = template  # TODO move out of here and into the tokenizer
        self.ps_version = ps_version  # Pixel shuffle version
        self.image_tag_type = image_tag_type  # TODO: into the tokenizer too?
        self.projector_hidden_size = projector_hidden_size
        self.vit_hidden_size = vit_hidden_size
        self.video_pruning_rate = video_pruning_rate

        self.norm_mean = float_triplet(norm_mean)
        self.norm_std = float_triplet(norm_std)
        self.use_thumbnail = use_thumbnail
        self.img_start_token = img_start_token
        self.img_end_token = img_end_token
        self.audio_context_token = audio_context_token
        self.audio_start_token = audio_start_token
        self.audio_end_token = audio_end_token

        # Dynamic resolution: from vision_config top-level
        self.min_num_patches = self.raw_vision_config.get("min_num_patches", 0)
        self.max_num_patches = self.raw_vision_config.get("max_num_patches", 0)
        self.dynamic_resolution = self.min_num_patches > 0

        # Video temporal compression: from vision_config top-level
        self.video_temporal_patch_size = self.raw_vision_config.get(
            "video_temporal_patch_size", 1
        )
        self.separate_video_embedder = self.raw_vision_config.get(
            "separate_video_embedder", True
        )
        self.video_target_num_patches = self.raw_vision_config.get(
            "video_target_num_patches", 0
        )
        self.video_maintain_aspect_ratio = self.raw_vision_config.get(
            "video_maintain_aspect_ratio", True
        )

    def create_radio_config(self):
        config = self.raw_vision_config
        model_name = config["args"]["model"]
        reg_tokens = config["args"].get("register_multiple")
        image_size = config.get("preferred_resolution", [224])[0]
        radio_config = RadioConfig(
            patch_size=self.patch_size,
            norm_mean=self.norm_mean,
            norm_std=self.norm_std,
            model_name=model_name,
            reg_tokens=reg_tokens,
            image_size=image_size,
            min_num_patches=self.min_num_patches,
            max_num_patches=self.max_num_patches,
            video_temporal_patch_size=self.video_temporal_patch_size,
            separate_video_embedder=self.separate_video_embedder,
            video_target_num_patches=self.video_target_num_patches,
            video_maintain_aspect_ratio=self.video_maintain_aspect_ratio,
        )
        return radio_config


class NemotronH_Nano_Omni_Reasoning_V3_Config(NemotronH_Nano_VL_V2_Config):
    model_type = "NemotronH_Nano_Omni_Reasoning_V3"

    def __init__(self, *args, **kwargs):
        # Explicit __init__ prevents PretrainedConfig.__init_subclass__ from
        # replacing the parent's custom __init__ with a dataclass-generated one.
        super().__init__(*args, **kwargs)
