import copy
import math
import os
import PIL
from PIL import Image
from typing import Dict, List, Optional, Union

import numpy as np
import torch
from PIL import Image
import re
from torchvision.transforms.functional import to_tensor
from transformers import (
    AutoTokenizer,
    AutoFeatureExtractor,
    AutoImageProcessor,
    AutoVideoProcessor,
)
from transformers.audio_utils import AudioInput
from transformers.image_processing_utils import (
    BaseImageProcessor,
    BatchFeature,
    get_size_dict,
)
from transformers.image_transforms import (
    convert_to_rgb,
    get_resize_output_image_size,
    resize,
    to_channel_dimension_format,
)
from transformers.image_utils import (
    ImageInput,
)
from transformers.processing_utils import (
    ProcessingKwargs,
    ProcessorMixin,
    SpecificProcessorType,
    Unpack,
)

from transformers.tokenization_utils_base import PreTokenizedInput, TextInput
from transformers.utils import TensorType, logging
from transformers.video_utils import VideoInput
from typing_extensions import Unpack
import librosa
import requests
import io
import base64

logger = logging.get_logger(__name__)


class HyperCLOVAXOmniProcessorKwargs(ProcessingKwargs, total=False):
    _defaults = {
        "audio_kwargs": {
            "sample_rate": 16_000,
            "chunk_unit": 80,
            "min_chunk_size": 1_600,
        },
        "images_kwargs": {},
        "text_kwargs": {
            "padding": False,
            "return_mm_token_type_ids": False,
        },
        "videos_kwargs": {
            "max_num_frames": 120,
        },
    }


class HyperCLOVAXOmniProcessor(ProcessorMixin):
    attributes = [
        "audio_processor",
        "image_processor",
        "video_processor",
        "tokenizer",
    ]
    audio_processor_class = "AutoFeatureExtractor"
    image_processor_class = "AutoImageProcessor"
    tokenizer_class = ("PreTrainedTokenizer", "PreTrainedTokenizerFast")
    video_processor_class = "AutoVideoProcessor"

    def __init__(
        self,
        audio_processor: Optional[AutoFeatureExtractor] = None,
        chat_template: Optional[str] = None,
        image_processor: Optional[AutoImageProcessor] = None,
        video_processor: Optional[AutoVideoProcessor] = None,
        tokenizer: AutoTokenizer = None,
        **kwargs,
    ):
        # Prefer explicit chat_template; fall back to tokenizer's if available
        if chat_template is None and hasattr(tokenizer, "chat_template"):
            chat_template = tokenizer.chat_template
        # Call the shared mixin directly with all declared attributes, including audio
        ProcessorMixin.__init__(
            self,
            audio_processor,
            image_processor,
            video_processor,
            tokenizer,
            chat_template=chat_template,
        )

        self.modalities = list()
        if self.audio_processor is not None:
            self.modalities.append("audio")
            self.audio_token = self.audio_processor.audio_token
            self.audio_token_id = tokenizer.convert_tokens_to_ids(self.audio_processor.audio_token)
            self.audio_start_token_id = tokenizer.convert_tokens_to_ids(self.audio_processor.audio_start_token)
            self.audio_end_token_id = tokenizer.convert_tokens_to_ids(self.audio_processor.audio_end_token)
            
            self.discrete_audio_token_id = None
            self.discrete_audio_start_token_id = None
            self.discrete_audio_end_token_id = None
            if self.image_processor.use_discrete_image_token:
                self.discrete_audio_token_id = tokenizer.convert_tokens_to_ids(self.audio_processor.discrete_audio_token)
                self.discrete_audio_start_token_id = tokenizer.convert_tokens_to_ids(self.audio_processor.discrete_audio_start_token)
                self.discrete_audio_end_token_id = tokenizer.convert_tokens_to_ids(self.audio_processor.discrete_audio_end_token)

        if self.image_processor is not None:
            self.modalities.append("image")
            self.image_token = self.image_processor.image_token
            self.image_token_id = tokenizer.convert_tokens_to_ids(self.image_processor.image_token)
            self.image_start_token_id = tokenizer.convert_tokens_to_ids(self.image_processor.image_start_token)
            self.image_end_token_id = tokenizer.convert_tokens_to_ids(self.image_processor.image_end_token)
            
            self.discrete_image_token_id = None
            self.discrete_image_start_token_id = None
            self.discrete_image_end_token_id = None
            if self.image_processor.use_discrete_image_token:
                self.discrete_image_token_id = tokenizer.convert_tokens_to_ids(self.image_processor.discrete_image_token)
                self.discrete_image_start_token_id = tokenizer.convert_tokens_to_ids(self.image_processor.discrete_image_start_token)
                self.discrete_image_end_token_id = tokenizer.convert_tokens_to_ids(self.image_processor.discrete_image_end_token)

        if self.video_processor is not None:
            self.modalities.append("video")
            self.video_token = self.video_processor.video_token
            self.video_token_id = tokenizer.convert_tokens_to_ids(self.video_processor.video_token)
            self.video_start_token_id = tokenizer.convert_tokens_to_ids(self.video_processor.video_start_token)
            self.video_end_token_id = tokenizer.convert_tokens_to_ids(self.video_processor.video_end_token)
            self.video_audio_token = "<|VIDEO_AUDIO_PAD|>"
            self.video_audio_token_id = tokenizer.convert_tokens_to_ids("<|VIDEO_AUDIO_PAD|>")

    @classmethod
    def from_pretrained(
        cls: type[SpecificProcessorType],
        pretrained_model_name_or_path: Union[str, os.PathLike],
        **kwargs,
    ):
        audio_processor_kwargs = kwargs.pop("audio_processor_kwargs", dict())
        image_processor_kwargs = kwargs.pop("image_processor_kwargs", dict())
        video_processor_kwargs = kwargs.pop("video_processor_kwargs", dict())

        if "tokenizer" not in kwargs:
            kwargs["tokenizer"] = AutoTokenizer.from_pretrained(
                pretrained_model_name_or_path,
                **kwargs,
            )

        if not kwargs.get("audio_processor", None):
            kwargs["audio_processor"] = AutoFeatureExtractor.from_pretrained(
                pretrained_model_name_or_path,
                **audio_processor_kwargs,
                **kwargs,
            )

        if not kwargs.get("image_processor", None):
            kwargs["image_processor"] = AutoImageProcessor.from_pretrained(
                pretrained_model_name_or_path,
                **image_processor_kwargs,
                **kwargs,
            )

        if not kwargs.get("video_processor", None):
            kwargs["video_processor"] = AutoVideoProcessor.from_pretrained(
                pretrained_model_name_or_path,
                **video_processor_kwargs,
                **kwargs,
            )

        return super().from_pretrained(
            pretrained_model_name_or_path=pretrained_model_name_or_path,
            **kwargs,
        )

    def save_pretrained(
        self,
        save_directory: Union[str, os.PathLike],
        *args,
        **kwargs,
    ):
        original_attributes = list(self.__class__.attributes)
        try:
            audio_processor = getattr(self, "audio_processor", None)
            if audio_processor is None and "audio_processor" in self.__class__.attributes:
                self.__class__.attributes = [a for a in self.__class__.attributes if a != "audio_processor"]

            self.register_for_auto_class()
            super().save_pretrained(save_directory, *args, **kwargs)
        finally:
            self.__class__.attributes = original_attributes

    def process_mm_info(self, conversations: Union[List[Dict], List[List[Dict]]], use_audio_in_video: bool = False):
        """
        Conversation 리스트를 입력받아 실제 오디오(numpy), 이미지(PIL), 비디오 데이터를 로드하여 반환합니다.
        지원하는 입력 포맷: 로컬 파일, HTTP URL, Base64
        """
        import mimetypes

        if isinstance(conversations, dict):
            conversations = [conversations]

        if isinstance(conversations[0], dict):
            conversations = [conversations]

        audios = []
        images = []
        videos = []

        target_sr = 16_000
        if (
            self.audio_processor is not None 
            and hasattr(self.audio_processor, "sampling_rate")
        ):
            target_sr = self.audio_processor.sampling_rate

        for conversation in conversations:
            for message in conversation:
                # 메시지 레벨에 미디어 파일이 있는 경우 (content가 string이라 아래 list 순회에서 처리되지 않는 경우)
                if message.get("audio_files"):
                    for audio_path in message["audio_files"]:
                        audio_data = self._load_audio(audio_path, sr=target_sr)
                        audios.append(audio_data)

                if message.get("image_files"):
                    for image_path in message["image_files"]:
                        image_data = self._load_image(image_path)
                        images.append(image_data)

                if message.get("video_files"):
                    for video_path in message["video_files"]:
                        if use_audio_in_video:
                            try:
                                audio_data = self._load_audio_from_video(video_path, sr=target_sr)
                                audios.append(audio_data)
                            except Exception as e:
                                print(f"[Warning] Failed to extract audio from video {video_path}: {e}")
                        video_data = self._load_video(video_path)
                        videos.append(video_data)

                content = message.get("content", [])
                if not isinstance(content, list):
                    continue

                for ele in content:
                    type_ = ele.get("type")

                    if type_ == "audio":
                        path = ele.get("audio", ele.get("audio_url"))
                        if path:
                            # Add MIME type dynamically using mimetypes module (like Preprocessor)
                            if "mime_type" not in ele and isinstance(path, str):
                                # Get filename from path or use default
                                filename = ele.get("filename", path if not path.startswith("http") else "a.wav")
                                mime_type = mimetypes.guess_type(filename)[0]
                                if mime_type:
                                    ele["mime_type"] = mime_type

                            audio_data = self._load_audio(
                                path, sr=target_sr, start=ele.get("audio_start", 0.0), end=ele.get("audio_end", None)
                            )
                            audios.append(audio_data)

                    elif type_ == "image":
                        path = ele.get("image", ele.get("image_url"))
                        if path:
                            # Add MIME type dynamically using mimetypes module (like Preprocessor)
                            if "mime_type" not in ele and isinstance(path, str):
                                filename = ele.get("filename", path if not path.startswith("http") else "a.jpg")
                                mime_type = mimetypes.guess_type(filename)[0]
                                if mime_type:
                                    ele["mime_type"] = mime_type

                            image_data = self._load_image(path)
                            images.append(image_data)

                    elif type_ == "video":
                        path = ele.get("video", ele.get("video_url"))
                        if path:
                            # Add MIME type dynamically using mimetypes module (like Preprocessor)
                            if "mime_type" not in ele and isinstance(path, str):
                                filename = ele.get("filename", path if not path.startswith("http") else "a.mp4")
                                mime_type = mimetypes.guess_type(filename)[0]
                                if mime_type:
                                    ele["mime_type"] = mime_type

                            if use_audio_in_video:
                                try:
                                    audio_data = self._load_audio_from_video(
                                        path,
                                        sr=target_sr,
                                        start=ele.get("video_start", 0.0),
                                        end=ele.get("video_end", None),
                                    )
                                    audios.append(audio_data)
                                except Exception as e:
                                    print(f"[Warning] Failed to extract audio from video {path}: {e}")

                            video_data = self._load_video(
                                path,
                                start=ele.get("video_start", 0.0),
                                end=ele.get("video_end", None),
                                max_num_frames=ele.get("max_num_frames", None),
                            )
                            videos.append(video_data)

        return (audios if audios else None, images if images else None, videos if videos else None)

    def _load_audio(self, path, sr=16000, start=0.0, end=None):
        if isinstance(path, np.ndarray):
            if path.ndim > 1:
                path = path.mean(axis=1)
            start_idx = int(sr * start)
            end_idx = int(sr * end) if end is not None else None
            return path[start_idx:end_idx]

        audio_source = path

        if isinstance(path, str):
            if path.startswith("data:audio"):
                _, base64_data = path.split("base64,", 1)
                audio_source = io.BytesIO(base64.b64decode(base64_data))
            elif path.startswith("http://") or path.startswith("https://"):
                response = requests.get(path)
                response.raise_for_status()
                audio_source = io.BytesIO(response.content)

        duration = (end - start) if end is not None else None

        y, _ = librosa.load(audio_source, sr=sr, offset=start, duration=duration)
        return y

    def _load_audio_from_video(self, path, sr=16000, start=0.0, end=None):
        """decord.AudioReader를 사용하여 비디오 파일에서 오디오 추출 (preprocessor 방식)"""
        from decord import AudioReader
        from decord import cpu as decord_cpu

        ar = AudioReader(path, ctx=decord_cpu(0), sample_rate=sr, mono=True)
        total_samples = ar.shape[1]

        start_sample = int(start * sr)
        end_sample = int(end * sr) if end is not None else total_samples
        end_sample = min(end_sample, total_samples)

        audio = ar[start_sample:end_sample].asnumpy().flatten().astype(np.float32)
        return audio

    def _load_image(self, path):
        image_source = path

        if isinstance(path, str):
            if path.startswith("data:image"):
                _, base64_data = path.split("base64,", 1)
                image_source = io.BytesIO(base64.b64decode(base64_data))
            elif path.startswith("http://") or path.startswith("https://"):  # URL
                image_source = requests.get(path, stream=True).raw

        image = Image.open(image_source)

        if image.mode != "RGB":
            image = image.convert("RGB")

        return image

    def _load_video(self, path, start=0.0, end=None, max_num_frames=None, fps=2.0):
        """Load video frames from a local file, HTTP URL, or base64 string.

        Returns a list of PIL.Image.Image frames that the video_processor can consume.
        """
        video_source = path

        if isinstance(path, str):
            if path.startswith("data:video"):
                _, base64_data = path.split("base64,", 1)
                video_source = io.BytesIO(base64.b64decode(base64_data))
            elif path.startswith("http://") or path.startswith("https://"):
                response = requests.get(path)
                response.raise_for_status()
                video_source = io.BytesIO(response.content)

        try:
            import decord
            from decord import cpu as decord_cpu

            vr = decord.VideoReader(video_source, ctx=decord_cpu(0))
            total_frames = len(vr)
            video_fps = vr.get_avg_fps()

            # Determine frame range from start/end times
            start_frame = int(start * video_fps) if start else 0
            end_frame = int(end * video_fps) if end is not None else total_frames - 1
            end_frame = min(end_frame, total_frames - 1)
            if start_frame >= end_frame:
                start_frame, end_frame = 0, total_frames - 1
            available_frames = end_frame - start_frame + 1

            # Sample frames at the target fps
            nframes = max(2, round(available_frames / video_fps * fps))
            if max_num_frames is not None:
                nframes = min(nframes, max_num_frames)
            # Ensure even number of frames (FRAME_FACTOR=2)
            nframes = max(2, nframes - (nframes % 2))
            nframes = min(nframes, available_frames)

            idx = torch.linspace(start_frame, end_frame, nframes).round().long().tolist()
            frames_np = vr.get_batch(idx).asnumpy()  # (T, H, W, C)

            # Convert to list of PIL images
            frames = [Image.fromarray(frames_np[i]) for i in range(frames_np.shape[0])]

            del vr
            return frames

        except ImportError:
            pass

        # Fallback: torchvision
        import torchvision

        if isinstance(video_source, io.BytesIO):
            # torchvision needs a file path; write to a temp file
            import tempfile

            with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as tmp:
                tmp.write(video_source.getvalue())
                tmp_path = tmp.name
            video_source = tmp_path
        else:
            tmp_path = None

        try:
            video_tensor, _, info = torchvision.io.read_video(
                video_source, start_pts=start, end_pts=end, pts_unit="sec"
            )
            # video_tensor shape: (T, H, W, C)
            total_frames = video_tensor.shape[0]
            video_fps = info.get("video_fps", 24.0)

            nframes = max(2, round(total_frames / video_fps * fps))
            if max_num_frames is not None:
                nframes = min(nframes, max_num_frames)
            nframes = max(2, nframes - (nframes % 2))
            nframes = min(nframes, total_frames)

            idx = torch.linspace(0, total_frames - 1, nframes).round().long()
            sampled = video_tensor[idx].numpy()  # (T, H, W, C)

            frames = [Image.fromarray(sampled[i]) for i in range(sampled.shape[0])]
            return frames
        finally:
            if tmp_path is not None:
                os.remove(tmp_path)

    def __call__(
        self,
        text: Union[TextInput, PreTokenizedInput, list[TextInput], list[PreTokenizedInput]] = None,
        audios: AudioInput | None = None,
        images: ImageInput | None = None,
        videos: VideoInput | None = None,
        **kwargs: Unpack[HyperCLOVAXOmniProcessorKwargs],
    ) -> BatchFeature:
        """
        Main method to prepare for the model one or several sequences(s) and image(s). This method forwards the `text`
        and `kwargs` arguments to Qwen2TokenizerFast's [`~Qwen2TokenizerFast.__call__`] if `text` is not `None` to encode
        the text. To prepare the vision inputs, this method forwards the `vision_infos` and `kwrags` arguments to
        Qwen2VLImageProcessor's [`~Qwen2VLImageProcessor.__call__`] if `vision_infos` is not `None`.

        Args:
            images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `list[PIL.Image.Image]`, `list[np.ndarray]`, `list[torch.Tensor]`):
                The image or batch of images to be prepared. Each image can be a PIL image, NumPy array or PyTorch
                tensor. Both channels-first and channels-last formats are supported.
            text (`str`, `list[str]`, `list[list[str]]`):
                The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings
                (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set
                `is_split_into_words=True` (to lift the ambiguity with a batch of sequences).
            videos (`np.ndarray`, `torch.Tensor`, `list[np.ndarray]`, `list[torch.Tensor]`):
                The image or batch of videos to be prepared. Each video can be a 4D NumPy array or PyTorch
                tensor, or a nested list of 3D frames. Both channels-first and channels-last formats are supported.
            return_tensors (`str` or [`~utils.TensorType`], *optional*):
                If set, will return tensors of a particular framework. Acceptable values are:
                - `'tf'`: Return TensorFlow `tf.constant` objects.
                - `'pt'`: Return PyTorch `torch.Tensor` objects.
                - `'np'`: Return NumPy `np.ndarray` objects.
                - `'jax'`: Return JAX `jnp.ndarray` objects.

        Returns:
            [`BatchFeature`]: A [`BatchFeature`] with the following fields:

            - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`.
            - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when
              `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not
              `None`).
            - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`.
            - **pixel_values_videos** -- Pixel values of videos to be fed to a model. Returned when `videos` is not `None`.
            - **image_grid_thw** -- List of image 3D grid in LLM. Returned when `images` is not `None`.
            - **video_grid_thw** -- List of video 3D grid in LLM. Returned when `videos` is not `None`.
        """
        output_kwargs = self._merge_kwargs(
            HyperCLOVAXOmniProcessorKwargs,
            tokenizer_init_kwargs=self.tokenizer.init_kwargs,
            **kwargs,
        )

        # [Text Processing] (Placeholder Replacement)
        if text is None:
            pass
        else:
            if isinstance(text, str):
                text = [
                    text,
                ]
            # below lines change text in-place
            text = copy.deepcopy(text)

        # [Audio Processing]
        audio_inputs = dict()
        discrete_audio_inputs = dict()
        if audios is not None and self.audio_processor is not None:
            if (
                len(audios) > 0 
                and isinstance(audios[0], np.ndarray)
            ):
                audios = [audios, ]

            _no_concat_keys = {"num_audio_tokens", "num_discrete_audio_tokens"}
            for _audio_clips in audios:
                _audio_features = self.audio_processor(
                    audios=_audio_clips,
                    **output_kwargs.get("audio_kwargs", {}),
                )
                for _k, _v in _audio_features.items():
                    if _k in ["discrete_audio_values", "num_discrete_audio_tokens"]:
                        if _k not in discrete_audio_inputs:
                            discrete_audio_inputs[_k] = list()
                        discrete_audio_inputs[_k].append(_v)
                    else:
                        if _k not in audio_inputs:
                            audio_inputs[_k] = list()
                        audio_inputs[_k].append(_v)
            audio_inputs = {
                _k: torch.cat(_v, dim=0) if isinstance(_v[0], torch.Tensor) and _k not in _no_concat_keys else _v
                for _k, _v in audio_inputs.items()
            }
            if discrete_audio_inputs:
                discrete_audio_inputs = {
                    _k: torch.cat(_v, dim=0) if isinstance(_v[0], torch.Tensor) and _k not in _no_concat_keys else _v
                    for _k, _v in discrete_audio_inputs.items()
                }

        # [Image Processing]
        image_inputs, image_grid_thw = dict(), list()
        discrete_image_inputs, discrete_image_ratios = dict(), list()
        if images is not None and self.image_processor is not None:
            # if len(images) > 0 and isinstance(images[0], PIL.Image.Image):
            #     images = [images]

            image_features = self.image_processor(images=images, **output_kwargs.get("images_kwargs", {}))
            for k, v in image_features.items():
                if k in ["discrete_pixel_values", "discrete_image_ratios"]:
                    discrete_image_inputs[k] = v
                else:
                    image_inputs[k] = v

        # [Video Processing]
        video_inputs, video_grid_thw = dict(), list()
        if videos is not None and self.video_processor is not None:
            if len(videos) > 0 and isinstance(videos[0], np.ndarray):  # sample to batch if a single item is given
                videos = [
                    videos,
                ]

            # Video feature extraction
            video_inputs = dict()
            video_grid_thw = list()
            # HyperCLOVAXOmni 전용 kwargs를 Qwen video processor에 전달하지 않도록 분리
            _videos_kwargs = {
                k: v
                for k, v in output_kwargs["videos_kwargs"].items()
                if k not in ("max_num_frames", "use_video_audio")
            }
            for _videos in videos:
                _video_inputs = self.video_processor(
                    videos=_videos,
                    **_videos_kwargs,
                )
                _video_grid_thw = _video_inputs["video_grid_thw"]
                for _k, _v in _video_inputs.items():
                    if _k not in video_inputs:
                        video_inputs[_k] = list()
                    video_inputs[_k].append(_v)
                video_grid_thw.append(_video_grid_thw)
            video_inputs = {
                _k: torch.cat(_v, dim=0) if isinstance(_v[0], torch.Tensor) else _v for _k, _v in video_inputs.items()
            }

        # [Duration Replacement] - <|audio_duration|> / <|video_duration|> 플레이스홀더를 실제 값으로 치환
        if text is not None and audios is not None:
            sr = self.audio_processor.sampling_rate if hasattr(self.audio_processor, "sampling_rate") else 16000
            # audios는 [batch][audio_idx] 또는 [audio_idx] 형태
            flat_audios = audios
            if len(audios) > 0 and isinstance(audios[0], list):
                flat_audios = [a for batch in audios for a in batch]
            audio_dur_idx = 0
            for _sample_idx, _text in enumerate(text):
                while "<|audio_duration|>" in _text and audio_dur_idx < len(flat_audios):
                    audio_data = flat_audios[audio_dur_idx]
                    duration_sec = len(audio_data) / sr
                    # Add quotes around duration to maintain valid JSON format
                    _text = _text.replace("<|audio_duration|>", f'"{duration_sec:.2f}s"', 1)
                    audio_dur_idx += 1
                text[_sample_idx] = _text

        if text is not None and videos is not None:
            # videos는 [batch] 형태, 각 batch는 list of PIL images
            fps = output_kwargs.get("videos_kwargs", {}).get("fps", 2.0)
            flat_videos = videos
            if (
                len(videos) > 0
                and isinstance(videos[0], list)
                and len(videos[0]) > 0
                and isinstance(videos[0][0], list)
            ):
                flat_videos = [v for batch in videos for v in batch]
            video_dur_idx = 0
            for _sample_idx, _text in enumerate(text):
                while "<|video_duration|>" in _text and video_dur_idx < len(flat_videos):
                    video_frames = flat_videos[video_dur_idx]
                    num_frames = len(video_frames) if isinstance(video_frames, list) else video_frames.shape[0]
                    duration_sec = round(num_frames / fps, 2)
                    _text = _text.replace("<|video_duration|>", f"{duration_sec}s", 1)
                    video_dur_idx += 1
                text[_sample_idx] = _text

        # [Expansion] - Audio (discrete)
        if (
            text is not None
            and discrete_audio_inputs
            and self.audio_processor is not None
            and self.audio_processor.use_discrete_audio_token
        ):
            for _sample_idx, (_text_before, _num_discrete_audio_tokens) in enumerate(
                zip(text, discrete_audio_inputs["num_discrete_audio_tokens"])
            ):
                discrete_audio_block_pattern = (
                    re.escape(self.audio_processor.discrete_audio_start_token)
                    + r".*?"
                    + re.escape(self.audio_processor.discrete_audio_token)
                    + r".*?"
                    + re.escape(self.audio_processor.discrete_audio_end_token)
                )
                _find_iters = list(re.finditer(discrete_audio_block_pattern, _text_before))
                if len(_find_iters) > 0:
                    _text_after = ""
                    _prev_end_idx = 0
                    for _idx, _match in enumerate(_find_iters):
                        _discrete_replacement = self.audio_processor.discrete_audio_token * int(
                            _num_discrete_audio_tokens[_idx]
                        )
                        _discrete_replacement = f"{self.audio_processor.discrete_audio_start_token}{_discrete_replacement}{self.audio_processor.discrete_audio_end_token}"
                        _text_after += _text_before[_prev_end_idx : _match.start()]
                        _text_after += _discrete_replacement
                        _prev_end_idx = _match.end()
                    _text_after += _text_before[_prev_end_idx:]
                    text[_sample_idx] = _text_after

        # [Expansion] - Audio (continuous)
        if text is not None and audio_inputs:
            for _sample_idx, (_text_before, _num_audio_tokens) in enumerate(
                zip(text, audio_inputs["num_audio_tokens"])
            ):
                cont_audio_block_pattern = (
                    re.escape(self.audio_processor.audio_start_token)
                    + r".*?"
                    + re.escape(self.audio_processor.audio_token)
                    + r".*?"
                    + re.escape(self.audio_processor.audio_end_token)
                )
                _find_iters = list(re.finditer(cont_audio_block_pattern, _text_before))
                if len(_find_iters) > 0:
                    _text_after = ""
                    _prev_end_idx = 0
                    for _idx, _match in enumerate(_find_iters):
                        _cont_replacement = self.audio_processor.audio_token * int(_num_audio_tokens[_idx])
                        _cont_replacement = f"{self.audio_processor.audio_start_token}{_cont_replacement}{self.audio_processor.audio_end_token}"
                        _text_after += _text_before[_prev_end_idx : _match.start()]
                        _text_after += _cont_replacement
                        _prev_end_idx = _match.end()
                    _text_after += _text_before[_prev_end_idx:]
                    text[_sample_idx] = _text_after

        # [Expansion] - Image (discrete)
        if (
            text is not None
            and discrete_image_inputs
            and self.image_processor is not None
            and self.image_processor.use_discrete_image_token
        ):
            for _sample_idx, (_text_before, _sample_discrete_ratios) in enumerate(
                zip(text, discrete_image_inputs["discrete_image_ratios"])
            ):
                discrete_image_block_pattern = (
                    re.escape(self.image_processor.discrete_image_start_token)
                    + r".*?"
                    + re.escape(self.image_processor.discrete_image_token)
                    + r".*?"
                    + re.escape(self.image_processor.discrete_image_end_token)
                )
                _find_iters = list(re.finditer(discrete_image_block_pattern, _text_before))
                if len(_find_iters) > 0:
                    _text_after = ""
                    _prev_end_idx = 0
                    discrete_token_size = self.image_processor.discrete_token_size
                    for _idx, _match in enumerate(_find_iters):
                        _row_str = f"{(self.image_processor.discrete_image_token * discrete_token_size)}{self.image_processor.vision_eol_token}"
                        if isinstance(_sample_discrete_ratios, (list, tuple)):
                            ratio_key = f"{int(_sample_discrete_ratios[0])}:{int(_sample_discrete_ratios[1])}"
                        else:
                            ratio_key = f"{_sample_discrete_ratios[0].item()}:{_sample_discrete_ratios[1].item()}"
                        _ratio_token = self.image_processor.discrete_image_ratio_tokens[ratio_key]
                        _discrete_replacement = f"{self.image_processor.discrete_image_start_token}{_ratio_token}{(_row_str * discrete_token_size)}{self.image_processor.vision_eof_token}{self.image_processor.discrete_image_end_token}"
                        _text_after += _text_before[_prev_end_idx : _match.start()]
                        _text_after += _discrete_replacement
                        _prev_end_idx = _match.end()
                    _text_after += _text_before[_prev_end_idx:]
                    text[_sample_idx] = _text_after

        # [Expansion] - Image (continuous)
        if text is not None and image_inputs:
            for _sample_idx, (_text_before, _sample_token_counts) in enumerate(
                zip(text, image_inputs["vision_query_lengths"])
            ):
                cont_image_block_pattern = (
                    re.escape(self.image_processor.image_start_token)
                    + r".*?"
                    + re.escape(self.image_processor.image_token)
                    + r".*?"
                    + re.escape(self.image_processor.image_end_token)
                )
                _find_iters = list(re.finditer(cont_image_block_pattern, _text_before))
                if len(_find_iters) > 0:
                    _text_after = ""
                    _prev_end_idx = 0
                    for _idx, _match in enumerate(_find_iters):
                        _cont_replacement = self.image_processor.image_token * int(_sample_token_counts)
                        _cont_replacement = f"{self.image_processor.image_start_token}{_cont_replacement}{self.image_processor.image_end_token}"
                        _text_after += _text_before[_prev_end_idx : _match.start()]
                        _text_after += _cont_replacement
                        _prev_end_idx = _match.end()
                    _text_after += _text_before[_prev_end_idx:]
                    text[_sample_idx] = _text_after

        # [Expansion] - Video
        if text is not None and video_inputs:
            for _sample_idx, (_text_before, _video_grid_thw) in enumerate(zip(text, video_inputs["video_grid_thw"])):
                video_block_pattern = (
                    re.escape(self.video_processor.video_start_token)
                    + r".*?"
                    + re.escape(self.video_processor.video_token)
                    + r".*?"
                    + re.escape(self.video_processor.video_end_token)
                )
                _find_iters = list(re.finditer(video_block_pattern, _text_before))
                if len(_find_iters) > 0:
                    _text_after = ""
                    _prev_end_idx = 0
                    for _idx, _continuous_video_match in enumerate(_find_iters):
                        _cur_start_idx = _continuous_video_match.start()
                        _inplace_str = self.get_video_token_replacement(
                            video_grid_thw=_video_grid_thw[_idx],
                            include_boundary_tokens=True,
                            tokenize=False,
                        )
                        _text_after += _text_before[_prev_end_idx:_cur_start_idx]
                        _text_after += _inplace_str
                        _prev_end_idx = _continuous_video_match.end()
                    _text_after += _text_before[_prev_end_idx:]
                    text[_sample_idx] = _text_after

        return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None)
        return_mm_token_type_ids = output_kwargs["text_kwargs"].pop("return_mm_token_type_ids", False)
        use_video_audio = output_kwargs.get("videos_kwargs", {}).get("use_video_audio", False)

        text_inputs = dict()
        if text is not None:
            text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"], return_tensors=None)
            self._check_special_mm_tokens(
                text,
                text_inputs,
                modalities=self.modalities,
            )

            # [Post-tokenization] - Video audio interleaving
            # If videos have audio, interleave VIDEO_PAD and VIDEO_AUDIO_PAD tokens
            if use_video_audio and video_inputs and hasattr(self, "video_audio_token_id"):
                text_inputs = self._interleave_video_audio_tokens(
                    text_inputs,
                    video_inputs.get("video_grid_thw", []),
                    audios,
                    output_kwargs.get("audio_kwargs", {}),
                )

        if return_mm_token_type_ids:
            array_ids = np.array(text_inputs["input_ids"])
            mm_token_type_ids = np.zeros_like(text_inputs["input_ids"])
            mm_token_type_ids[array_ids == self.image_processor.image_token_id] = 1
            if text_inputs:
                text_inputs["mm_token_type_ids"] = mm_token_type_ids.tolist()

        data = {
            **audio_inputs,
            **image_inputs,
            **text_inputs,
            **video_inputs,
        }
        if (
            discrete_audio_inputs
            and self.audio_processor.use_discrete_audio_token
        ):
            data.update(discrete_audio_inputs)
        if (
            discrete_image_inputs
            and self.image_processor.use_discrete_image_token
        ):
            data.update(discrete_image_inputs)

        model_inputs = BatchFeature(data=data, tensor_type=return_tensors)
        return model_inputs

    def get_audio_placeholder(
        self,
        tokenize: bool = False,
    ):
        audio_placeholder = ""
        if self.audio_processor.use_discrete_audio_token:
            audio_placeholder += f'{self.audio_processor.discrete_audio_start_token}{self.audio_processor.discrete_audio_token}{self.audio_processor.discrete_audio_end_token}\n'
        audio_placeholder += f'{self.audio_processor.audio_start_token}{self.audio_processor.audio_token}{self.audio_processor.audio_end_token}'
        if tokenize:
            audio_placeholder = self.tokenizer.encode(audio_placeholder)
        return audio_placeholder

    def get_audio_token_replacement(
        self,
        num_audio_tokens: int,
        num_discrete_audio_tokens: Optional[int] = None,
        include_boundary_tokens: Optional[bool] = False,
        tokenize: Optional[bool] = False,
        return_tuple: Optional[bool] = None,
    ):
        conitnuous_replacement, discrete_replacement = "", ""
        
        if (
            isinstance(num_audio_tokens, (list, tuple))
            or (isinstance(num_audio_tokens, torch.Tensor) and num_audio_tokens.dim() >= 1)
        ):
            num_audio_tokens = num_audio_tokens[0]
        conitnuous_replacement = self.audio_processor.audio_token * num_audio_tokens
        if include_boundary_tokens:
            conitnuous_replacement = f"{self.audio_processor.audio_start_token}{conitnuous_replacement}{self.audio_processor.audio_end_token}"
        
        if self.audio_processor.use_discrete_audio_token:
            if (
                isinstance(num_discrete_audio_tokens, (list, tuple))
                or (isinstance(num_discrete_audio_tokens, torch.Tensor) and num_discrete_audio_tokens.dim() >= 1)
            ):
                num_discrete_audio_tokens = num_discrete_audio_tokens[0]    
            discrete_replacement = self.audio_processor.discrete_audio_token * num_discrete_audio_tokens
            if include_boundary_tokens:
                discrete_replacement = f"{self.audio_processor.discrete_audio_start_token}{discrete_replacement}{self.audio_processor.discrete_audio_end_token}"
            discrete_replacement = f'{discrete_replacement}\n'
        
        if return_tuple:
            if tokenize:
                conitnuous_replacement = self.tokenizer.encode(conitnuous_replacement)
                discrete_replacement = self.tokenizer.encode(discrete_replacement)
            return (conitnuous_replacement, discrete_replacement)
        else:
            replacement = f'{discrete_replacement}{conitnuous_replacement}'
            if tokenize:
                replacement = self.tokenizer.encode(replacement)
            return replacement

    def get_image_placeholder(
        self,
        tokenize: bool = False,
    ):
        image_placeholder = ""
        if self.image_processor.use_discrete_audio_token:
            image_placeholder += f'{self.image_processor.discrete_image_start_token}{self.image_processor.discrete_image_token}{self.image_processor.discrete_image_end_token}\n'
        image_placeholder += f'{self.image_processor.image_start_token}{self.image_processor.image_token}{self.image_processor.image_end_token}'
        if tokenize:
            image_placeholder = self.tokenizer.encode(image_placeholder)
        return image_placeholder

    def get_image_token_replacement(
        self,
        num_image_tokens: int,
        discrete_image_ratio: Optional[List[int]] = None,
        include_boundary_tokens: Optional[bool] = False,
        tokenize: Optional[bool] = False,
        return_tuple: Optional[bool] = None,
    ):
        conitnuous_replacement, discrete_replacement = "", ""
        
        if (
            isinstance(num_image_tokens, (list, tuple))
            or (isinstance(num_image_tokens, torch.Tensor) and num_image_tokens.dim() >= 1)
        ):
            num_image_tokens = num_image_tokens[0]
        discrete_token_size = self.image_processor.discrete_token_size
        continuous_replacement = self.image_processor.image_token * num_image_tokens
        if include_boundary_tokens:
            continuous_replacement = f"{self.image_processor.image_start_token}{continuous_replacement}{self.image_processor.image_end_token}"
        
        if self.image_processor.use_discrete_image_token:
            if (
                discrete_image_ratio
                and len(discrete_image_ratio) == 1
                and (
                    isinstance(discrete_image_ratio, (list, tuple)) 
                    or (isinstance(discrete_image_ratio, torch.Tensor) and discrete_image_ratio.dim() >= 2)
                )
            ): # [[16, 9]], or torch.Tensor([[16, 9]])
                discrete_image_ratio = discrete_image_ratio[0]
            row_str = self.image_processor.discrete_image_token * self.image_processor.discrete_token_size
            # row_str += self.image_processor.vision_eol_token
            discrete_replacement = row_str * self.image_processor.discrete_token_size
            if discrete_image_ratio:
                if isinstance(discrete_image_ratio, (list, tuple)):
                    ratio_key = f"{int(discrete_image_ratio[0])}:{int(discrete_image_ratio[1])}"
                elif isinstance(discrete_image_ratio, torch.Tensor):
                    ratio_key = f"{discrete_image_ratio[0].item()}:{discrete_image_ratio[1].item()}"
                discrete_image_ratio_token = self.image_processor.discrete_image_ratio_tokens[ratio_key]
                discrete_replacement = f"{discrete_image_ratio_token}{discrete_replacement}"
            # discrete_replacement = f"{discrete_replacement}{self.image_processor.vision_eof_token}"
            if include_boundary_tokens:
                discrete_replacement = f"{self.image_processor.discrete_image_start_token}{discrete_replacement}{self.image_processor.discrete_image_end_token}"
            discrete_replacement = f'{discrete_replacement}\n'

        if return_tuple:
            if tokenize:
                conitnuous_replacement = self.tokenizer.encode(conitnuous_replacement)
                discrete_replacement = self.tokenizer.encode(discrete_replacement)
            return (conitnuous_replacement, discrete_replacement)
        else:
            replacement = f'{discrete_replacement}{conitnuous_replacement}'
            if tokenize:
                replacement = self.tokenizer.encode(replacement)
            return replacement

    def get_video_placeholder(
        self,
        tokenize: bool = False,
    ):
        video_placeholder = f'{self.video_processor.video_start_token}{self.video_processor.video_token}{self.video_processor.video_end_token}'
        if tokenize:
            video_placeholder = self.tokenizer.encode(video_placeholder)
        return video_placeholder

    def get_video_token_replacement(
        self,
        num_video_tokens: int,
        include_boundary_tokens: Optional[bool] = False,
        tokenize: Optional[bool] = False,
        return_tuple: Optional[bool] = None,
    ):
        conitnuous_replacement, discrete_replacement = "", ""
        
        if (
            isinstance(num_video_tokens, (list, tuple))
            or (isinstance(num_video_tokens, torch.Tensor) and num_video_tokens.dim() >= 1)
        ):
            num_video_tokens = num_video_tokens[0]
        merge_length = self.video_processor.video_merge_size**2
        conitnuous_replacement = self.video_processor.video_token * int(num_video_tokens)
        if include_boundary_tokens:
            conitnuous_replacement = f"{self.video_processor.video_start_token}{conitnuous_replacement}{self.video_processor.video_end_token}"
            
        if return_tuple:
            if tokenize:
                conitnuous_replacement = self.tokenizer.encode(conitnuous_replacement)
                discrete_replacement = self.tokenizer.encode(discrete_replacement)
            return (conitnuous_replacement, discrete_replacement)
        else:
            replacement = f'{discrete_replacement}{conitnuous_replacement}'
            if tokenize:
                replacement = self.tokenizer.encode(replacement)
            return replacement

    def _interleave_video_audio_tokens(
        self,
        text_inputs: dict,
        video_grid_thw: list,
        audios: Optional[list] = None,
        audio_kwargs: Optional[dict] = None,
    ):
        """
        Interleave VIDEO_PAD and VIDEO_AUDIO_PAD tokens in input_ids.
        Based on preprocessor.py logic (lines 3515-3569).
        """
        if video_grid_thw is None or len(video_grid_thw) == 0 or audios is None:
            return text_inputs

        # Process each sample in the batch
        for sample_idx, input_ids in enumerate(text_inputs["input_ids"]):
            # Find all VIDEO_PAD token positions
            video_pad_positions = [i for i, token_id in enumerate(input_ids) if token_id == self.video_token_id]

            if not video_pad_positions:
                continue

            # Calculate video audio query lengths for each video
            # This follows preprocessor.py's video audio processing logic (lines 3115-3179)
            video_audio_query_lengths = []
            if audios and sample_idx < len(audios):
                sample_audios = audios[sample_idx] if isinstance(audios[0], list) else [audios[sample_idx]]

                for audio_data in sample_audios:
                    if audio_data is not None and len(audio_data) > 0:
                        # Calculate audio query length using the same logic as preprocessor
                        sr = audio_kwargs.get("sample_rate", 16000) if audio_kwargs else 16000
                        chunk_unit = audio_kwargs.get("chunk_unit", 80) if audio_kwargs else 80
                        chunk_size = chunk_unit * sr

                        audio_length = len(audio_data)
                        total_audio_tokens = 0
                        pool_size = 25

                        for start in range(0, audio_length, chunk_size):
                            end = min(start + chunk_size, audio_length)
                            chunked_length = end - start
                            num_mel_frames = chunked_length // 160
                            num_mel_frames_conv1 = (num_mel_frames + 2 * 1 - 1 * (3 - 1) - 1) // 2 + 1
                            num_audio_tokens = (num_mel_frames_conv1 + 2 * 1 - 1 * (3 - 1) - 1) // 2 + 1
                            total_audio_tokens += num_audio_tokens

                        total_audio_query_length = (total_audio_tokens + pool_size - 1) // pool_size
                        video_audio_query_lengths.append(total_audio_query_length)
                    else:
                        video_audio_query_lengths.append(0)

            # Interleave tokens for each video
            new_input_ids = []
            prev_idx = 0

            for vid_idx, (thw, vid_start_pos) in enumerate(zip(video_grid_thw, video_pad_positions)):
                # Copy tokens before this video
                new_input_ids.extend(input_ids[prev_idx:vid_start_pos])

                # Get video parameters
                # video_grid_thw는 shape에 따라 [T,H,W], [[T,H,W]], 등 다양한 형태일 수 있음
                if isinstance(thw, torch.Tensor):
                    thw_flat = thw.flatten()
                    num_frames = int(thw_flat[0])
                else:
                    num_frames = int(thw[0]) if not hasattr(thw[0], "__len__") else int(thw[0][0])

                # Count consecutive VIDEO_PAD tokens
                frame_query_length = 0
                pos = vid_start_pos
                while pos < len(input_ids) and input_ids[pos] == self.video_token_id:
                    frame_query_length += 1
                    pos += 1

                # Check if this video has audio
                has_video_audio = vid_idx < len(video_audio_query_lengths) and video_audio_query_lengths[vid_idx] > 0

                if has_video_audio and num_frames > 0:
                    total_audio_tokens = video_audio_query_lengths[vid_idx]

                    # Calculate token distribution per frame
                    frame_base = frame_query_length // num_frames
                    audio_base = total_audio_tokens // num_frames
                    audio_remainder = total_audio_tokens % num_frames

                    # Interleave frame and audio tokens
                    token_sequence = []
                    for frame_idx in range(num_frames):
                        # Add frame tokens
                        frame_tokens = frame_base
                        token_sequence.extend([self.video_token_id] * frame_tokens)

                        # Add audio tokens
                        audio_tokens = audio_base + (1 if frame_idx < audio_remainder else 0)
                        if audio_tokens > 0:
                            token_sequence.extend([self.video_audio_token_id] * audio_tokens)

                    new_input_ids.extend(token_sequence)
                else:
                    # No audio, just use frame tokens
                    new_input_ids.extend([self.video_token_id] * frame_query_length)

                prev_idx = pos

            # Copy remaining tokens
            new_input_ids.extend(input_ids[prev_idx:])
            text_inputs["input_ids"][sample_idx] = new_input_ids

        return text_inputs