# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
#
# 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 typing import List, Optional, Tuple, Union

from ......utils.deps import is_dep_available

if all(map(is_dep_available, ("sglang", "torch"))):
    import asyncio
    import math

    import torch
    from PIL import Image
    from sglang.srt.multimodal.processors.base_processor import (
        BaseMultimodalProcessor,
        MultimodalSpecialTokens,
    )

    def smart_resize(
        height: int,
        width: int,
        factor: int = 28,
        min_pixels: int = 28 * 28 * 130,
        max_pixels: int = 28 * 28 * 1280,
    ):
        """Rescales the image so that the following conditions are met:

        1. Both dimensions (height and width) are divisible by 'factor'.

        2. The total number of pixels is within the range ['min_pixels', 'max_pixels'].

        3. The aspect ratio of the image is maintained as closely as possible.

        """
        if height < factor:
            width = round((width * factor) / height)
            height = factor

        if width < factor:
            height = round((height * factor) / width)
            width = factor

        if max(height, width) / min(height, width) > 200:
            raise ValueError(
                f"absolute aspect ratio must be smaller than 200, got {max(height, width) / min(height, width)}"
            )
        h_bar = round(height / factor) * factor
        w_bar = round(width / factor) * factor
        if h_bar * w_bar > max_pixels:
            beta = math.sqrt((height * width) / max_pixels)
            h_bar = math.floor(height / beta / factor) * factor
            w_bar = math.floor(width / beta / factor) * factor
        elif h_bar * w_bar < min_pixels:
            beta = math.sqrt(min_pixels / (height * width))
            h_bar = math.ceil(height * beta / factor) * factor
            w_bar = math.ceil(width * beta / factor) * factor
        return h_bar, w_bar

    def resize_image(image, min_pixels, max_pixels, factor) -> Image.Image:
        width, height = image.size
        resized_height, resized_width = smart_resize(
            height,
            width,
            factor=factor,
            min_pixels=min_pixels,
            max_pixels=max_pixels,
        )
        image = image.resize((resized_width, resized_height))
        return image

    async def resize_image_async(image, min_pixels, max_pixels, factor):
        return resize_image(image, min_pixels, max_pixels, factor)

    class PaddleOCRVLImageProcessor(BaseMultimodalProcessor):

        def __init__(self, hf_config, server_args, _processor, *args, **kwargs):
            super().__init__(hf_config, server_args, _processor, *args, **kwargs)

            image_processor_config = _processor.image_processor
            self.MIN_PIXELS = image_processor_config.min_pixels
            self.MAX_PIXELS = image_processor_config.max_pixels
            self.IMAGE_FACTOR = (
                image_processor_config.patch_size * image_processor_config.merge_size
            )

            self.vision_start_token_id = hf_config.vision_start_token_id
            self.mm_tokens = MultimodalSpecialTokens(
                image_token="<|IMAGE_START|><|IMAGE_PLACEHOLDER|><|IMAGE_END|>",
                image_token_id=hf_config.image_token_id,
                video_token_id=hf_config.video_token_id,
            ).build(_processor)

        async def process_mm_data_async(
            self,
            image_data: List[Union[str, bytes]],
            input_text,
            request_obj,
            *args,
            **kwargs,
        ):
            base_output = self.load_mm_data(
                prompt=input_text,
                image_data=image_data,
                multimodal_tokens=self.mm_tokens,
            )

            if base_output.images and isinstance(base_output.images[0], Image.Image):
                resize_tasks = [
                    resize_image_async(
                        image, self.MIN_PIXELS, self.MAX_PIXELS, self.IMAGE_FACTOR
                    )
                    for image in base_output.images
                ]
                base_output.images = await asyncio.gather(*resize_tasks)

            mm_items, input_ids, ret = self.process_and_combine_mm_data(
                base_output, self.mm_tokens
            )

            input_ids = input_ids.flatten()
            mrope_positions, mrope_position_delta = self.get_rope_index(
                spatial_merge_size=self.hf_config.vision_config.spatial_merge_size,
                image_token_id=self.mm_tokens.image_token_id,
                video_token_id=self.mm_tokens.video_token_id,
                vision_start_token_id=self.vision_start_token_id,
                model_type=self.hf_config.model_type,
                tokens_per_second=getattr(
                    self.hf_config.vision_config, "tokens_per_second", None
                ),
                input_ids=input_ids.unsqueeze(0),
                image_grid_thw=getattr(ret, "image_grid_thw", None),
            )
            mrope_positions = mrope_positions.squeeze(1)

            return {
                "mm_items": mm_items,
                "input_ids": input_ids.tolist(),
                "im_token_id": self.mm_tokens.image_token_id,
                "mrope_positions": mrope_positions,
                "mrope_position_delta": mrope_position_delta,
            }

        @staticmethod
        def get_rope_index(
            spatial_merge_size: int,
            image_token_id: int,
            video_token_id: int,
            vision_start_token_id: int,
            model_type: str,
            tokens_per_second: Optional[int] = None,
            input_ids: Optional[torch.LongTensor] = None,
            image_grid_thw: Optional[torch.LongTensor] = None,
            video_grid_thw: Optional[torch.LongTensor] = None,
            second_per_grid_ts: Optional[torch.Tensor] = None,
            **kwargs,
        ) -> Tuple[torch.Tensor, torch.Tensor]:
            mrope_position_deltas = []
            if input_ids is not None and (
                image_grid_thw is not None or video_grid_thw is not None
            ):
                total_input_ids = input_ids
                position_ids = torch.ones(
                    3,
                    input_ids.shape[0],
                    input_ids.shape[1],
                    dtype=input_ids.dtype,
                    device=input_ids.device,
                )
                image_index, video_index = 0, 0
                for i, input_ids in enumerate(total_input_ids):
                    image_nums, video_nums = 0, 0
                    vision_start_indices = torch.argwhere(
                        input_ids == vision_start_token_id
                    ).squeeze(1)
                    vision_tokens = input_ids[vision_start_indices + 1]
                    image_nums = (vision_tokens == image_token_id).sum()
                    video_nums = (vision_tokens == video_token_id).sum()
                    input_tokens = input_ids.tolist()
                    llm_pos_ids_list: list = []
                    st = 0
                    remain_images, remain_videos = image_nums, video_nums
                    for _ in range(image_nums + video_nums):
                        if image_token_id in input_tokens and remain_images > 0:
                            ed_image = input_tokens.index(image_token_id, st)
                        else:
                            ed_image = len(input_tokens) + 1
                        if video_token_id in input_tokens and remain_videos > 0:
                            ed_video = input_tokens.index(video_token_id, st)
                        else:
                            ed_video = len(input_tokens) + 1
                        if ed_image < ed_video:
                            t, h, w = (
                                image_grid_thw[image_index][0],
                                image_grid_thw[image_index][1],
                                image_grid_thw[image_index][2],
                            )
                            second_per_grid_t = 0
                            image_index += 1
                            remain_images -= 1
                            ed = ed_image
                        else:
                            t, h, w = (
                                video_grid_thw[video_index][0],
                                video_grid_thw[video_index][1],
                                video_grid_thw[video_index][2],
                            )
                            if second_per_grid_ts is not None:
                                second_per_grid_t = second_per_grid_ts[video_index]
                            else:
                                second_per_grid_t = 1.0
                            video_index += 1
                            remain_videos -= 1
                            ed = ed_video
                        llm_grid_t, llm_grid_h, llm_grid_w = (
                            t.item(),
                            h.item() // spatial_merge_size,
                            w.item() // spatial_merge_size,
                        )
                        text_len = ed - st

                        st_idx = (
                            llm_pos_ids_list[-1].max() + 1
                            if len(llm_pos_ids_list) > 0
                            else 0
                        )
                        llm_pos_ids_list.append(
                            torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx
                        )

                        range_tensor = torch.arange(llm_grid_t).view(-1, 1)
                        expanded_range = range_tensor.expand(
                            -1, llm_grid_h * llm_grid_w
                        )

                        time_tensor = (
                            expanded_range * second_per_grid_t * tokens_per_second
                        )

                        time_tensor_long = time_tensor.long()
                        t_index = time_tensor_long.flatten()

                        h_index = (
                            torch.arange(llm_grid_h)
                            .view(1, -1, 1)
                            .expand(llm_grid_t, -1, llm_grid_w)
                            .flatten()
                        )
                        w_index = (
                            torch.arange(llm_grid_w)
                            .view(1, 1, -1)
                            .expand(llm_grid_t, llm_grid_h, -1)
                            .flatten()
                        )
                        llm_pos_ids_list.append(
                            torch.stack([t_index, h_index, w_index]) + text_len + st_idx
                        )
                        st = ed + llm_grid_t * llm_grid_h * llm_grid_w

                    if st < len(input_tokens):
                        st_idx = (
                            llm_pos_ids_list[-1].max() + 1
                            if len(llm_pos_ids_list) > 0
                            else 0
                        )
                        text_len = len(input_tokens) - st
                        llm_pos_ids_list.append(
                            torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx
                        )

                    llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
                    position_ids[..., i, :] = llm_positions.to(position_ids.device)
                    mrope_position_deltas.append(
                        llm_positions.max() + 1 - len(total_input_ids[i])
                    )
                mrope_position_deltas = torch.tensor(
                    mrope_position_deltas, device=input_ids.device
                ).unsqueeze(1)
                return position_ids, mrope_position_deltas
            else:
                s = input_ids.shape[1]
                position_ids = torch.arange(s)
                position_ids = (
                    position_ids.unsqueeze(0).expand(3, -1, -1).to(input_ids.device)
                )
                max_position_ids = position_ids.max(0, keepdim=False)[0].max(
                    -1, keepdim=True
                )[0]
                mrope_position_deltas = max_position_ids + 1 - s
                return position_ids, mrope_position_deltas
