import logging

import PIL
from transformers.image_utils import ChannelDimension

logger = logging.getLogger(__name__)

import torch
import numpy as np
from src.utils import print_master

PHI_IMAGE_TOKEN_MAX_INPUT_ID = int(1e9)
LLAVA_IMAGE_TOKEN_ID = 32000

PHI3V = 'phi3_v'
LLAVA_NEXT = 'llava_next'
QWEN2_VL = 'qwen2_vl'
QWEN2_VL_TOKENSELECTION = 'qwen2_vl'
QWEN2_5_VL = 'qwen2_5_vl'
QWEN2_VL_TOKENSELECTION = 'qwen2_vl_tokenselection'
QWEN2_5_VL_TOKENSELECTION = 'qwen2_5_vl_tokenselection'
QWEN2_5_OMNI = 'qwen2_5_omni'
INTERNVIDEO2 = 'internvideo2'
GME = 'gme'  # QWEN2-VL
LamRA = 'lamra'  # QWEN2-VL
LamRA_QWEN2_5 = 'lamra_qwen25'  # QWEN2.5-VL
COLPALI = 'colpali'  # PaliGemma-3B
WAVE = 'wave'  # Stock Qwen2.5-Omni Thinker (no custom modifications)
MODEL2BACKBONE = {  # keys are from hf_config.model_type or manually added if not provided
    'phi3_v': PHI3V,
    'llava_next': LLAVA_NEXT,
    'qwen2_vl': QWEN2_VL,
    'qwen2_vl_tokenselection': QWEN2_VL,
    'qwen2_5_vl': QWEN2_5_VL,
    'qwen2_vl_tokenselection': QWEN2_VL_TOKENSELECTION,
    'qwen2_5_vl_tokenselection': QWEN2_5_VL_TOKENSELECTION,
    'qwen2_5_omni': QWEN2_5_OMNI,
    'qwen2_5_omni_thinker': QWEN2_5_OMNI,  # Thinker model maps to same backbone
    'internvideo2': INTERNVIDEO2,
    'gme': GME, 
    'lamra': LamRA,
    'lamra_qwen25': LamRA,
    'colpali': COLPALI,
    'wave': WAVE,
}
SUPPORTED_MODELS = set(MODEL2BACKBONE.keys())

VLM_IMAGE_TOKENS = {
    PHI3V: "<|image_1|>",
    LLAVA_NEXT: "<image>",
    QWEN2_VL: "<|image_pad|>",
    QWEN2_5_VL: "<|image_pad|>",
    QWEN2_VL_TOKENSELECTION: "<|image_pad|>",
    QWEN2_5_VL_TOKENSELECTION: "<|image_pad|>",
    QWEN2_5_OMNI: "<|IMAGE|>",
    WAVE: "<|IMAGE|>",
    GME: "<|image_pad|>",
    LamRA: "<|image_pad|>",
    LamRA_QWEN2_5: "<|image_pad|>",
    INTERNVIDEO2: "",
    COLPALI: "",
}

VLM_VIDEO_TOKENS = {
    LLAVA_NEXT: "<image>",
    QWEN2_VL: "<|video_pad|>",
    QWEN2_5_VL: "<|video_pad|>",
    QWEN2_VL_TOKENSELECTION: "<|video_pad|>",
    QWEN2_5_VL_TOKENSELECTION: "<|video_pad|>",
    QWEN2_5_OMNI: "<|VIDEO|>",
    WAVE: "<|VIDEO|>",
    GME: "<|video_pad|>",
    LamRA: "<|video_pad|>",
    LamRA_QWEN2_5: "<|video_pad|>",
    INTERNVIDEO2: "",
    COLPALI: "",
}

def get_model_class(backbone_name):
    """Lazily import and return the model class for the given backbone.
    This avoids importing all backbone modules at the top level, which
    allows running with environments that only support a subset of backbones."""
    if backbone_name == PHI3V:
        from src.model.vlm_backbone.phi3_v.modeling_phi3_v import Phi3VForCausalLM
        return Phi3VForCausalLM
    elif backbone_name == LLAVA_NEXT:
        from src.model.vlm_backbone.llava_next import LlavaNextForConditionalGeneration
        return LlavaNextForConditionalGeneration
    elif backbone_name == QWEN2_VL:
        from src.model.vlm_backbone.qwen2_vl import Qwen2VLForConditionalGeneration
        return Qwen2VLForConditionalGeneration
    elif backbone_name == QWEN2_5_VL:
        from src.model.vlm_backbone.qwen2_5_vl import Qwen2_5_VLForConditionalGeneration
        return Qwen2_5_VLForConditionalGeneration
    elif backbone_name == QWEN2_VL_TOKENSELECTION:
        from src.model.vlm_backbone.qwen2_vl_tokenselection import \
            Qwen2VLForConditionalGeneration as Qwen2VLTokenSelectionForConditionalGeneration
        return Qwen2VLTokenSelectionForConditionalGeneration
    elif backbone_name == QWEN2_5_VL_TOKENSELECTION:
        from src.model.vlm_backbone.qwen2_5_vl_tokenselection import \
            Qwen2_5_VLForConditionalGeneration as Qwen2_5_VL_TokenSelectionForConditionalGeneration
        return Qwen2_5_VL_TokenSelectionForConditionalGeneration
    elif backbone_name == QWEN2_5_OMNI:
        from src.model.vlm_backbone.qwen2_5_omni.modeling_qwen2_5_omni import Qwen2_5OmniThinkerForConditionalGeneration
        return Qwen2_5OmniThinkerForConditionalGeneration
    elif backbone_name == INTERNVIDEO2:
        from src.model.vlm_backbone.internvideo2.modeling_internvideo2 import InternVideo2_Stage2
        return InternVideo2_Stage2
    elif backbone_name == WAVE:
        from src.model.vlm_backbone.wave.model.qwen2_5_omni.modeling_qwen2_5_omni import \
            Qwen2_5OmniThinkerForConditionalGeneration as WaveThinker
        return WaveThinker
    else:
        raise ValueError(f"No model class registered for backbone: {backbone_name}")


HF_LOCAL_FILES_ONLY = False


def load_processor(model_args, data_args=None):
    """
    Load processor based on VLM backbone.
    Note: due to this change, https://github.com/huggingface/transformers/commit/9215cc62d4366072aacafa4e44028c1ca187167b#diff-6505546ec5a9ab74b2ce6511681dd31194eb91e9fa3ce26282e487a5e61f9356L1102
    """
    model_name_or_path = model_args.checkpoint_path if model_args.checkpoint_path else model_args.model_name
    print_master(f'Loading processor from: {model_name_or_path}')
    if model_args.model_backbone == PHI3V:
        from src.model.vlm_backbone.phi3_v.processing_phi3_v import Phi3VProcessor
        processor = Phi3VProcessor.from_pretrained(
            model_name_or_path,
            trust_remote_code=True,
            num_crops=model_args.num_crops
        )
        processor.tokenizer.padding_side = "right"
    elif model_args.model_backbone == LLAVA_NEXT:
        from src.model.vlm_backbone.llava_next.processing_llava_next import LlavaNextProcessor
        processor = LlavaNextProcessor.from_pretrained(
            model_name_or_path,
            trust_remote_code=True
        )
    elif model_args.model_backbone in [QWEN2_VL, GME, LamRA]:
        from src.model.vlm_backbone.qwen2_vl.processing_qwen2_vl import Qwen2VLProcessor
        from src.model.vlm_backbone.qwen2_vl.image_processing_qwen2_vl import Qwen2VLImageProcessor
        from src.model.vlm_backbone.qwen2_vl.tokenization_qwen2_fast import Qwen2TokenizerFast
        min_pixels, max_pixels = None, None
        if data_args is not None:
            min_pixels, max_pixels = data_args.resize_min_pixels, data_args.resize_max_pixels
        size = {"shortest_edge": min_pixels, "longest_edge": max_pixels}
        image_processor = Qwen2VLImageProcessor.from_pretrained(model_name_or_path, size=size)
        tokenizer = Qwen2TokenizerFast.from_pretrained(model_name_or_path)
        processor = Qwen2VLProcessor.from_pretrained(
            model_name_or_path,
            image_processor=image_processor, tokenizer=tokenizer, size=size
        )
    elif model_args.model_backbone in [QWEN2_5_VL, LamRA_QWEN2_5]:
        from src.model.vlm_backbone.qwen2_5_vl.processing_qwen2_5_vl import Qwen2_5_VLProcessor
        from src.model.vlm_backbone.qwen2_5_vl.image_processing_qwen2_5_vl import Qwen2_5_VLImageProcessor
        from src.model.vlm_backbone.qwen2_vl.tokenization_qwen2_fast import Qwen2TokenizerFast
        min_pixels, max_pixels = None, None
        if data_args is not None:
            min_pixels, max_pixels = data_args.resize_min_pixels, data_args.resize_max_pixels
        size = {"shortest_edge": min_pixels, "longest_edge": max_pixels, "min_pixels": min_pixels, "max_pixels": max_pixels}
        image_processor = Qwen2_5_VLImageProcessor.from_pretrained(model_name_or_path, size=size)
        tokenizer = Qwen2TokenizerFast.from_pretrained(model_name_or_path)
        processor = Qwen2_5_VLProcessor.from_pretrained(model_name_or_path, image_processor=image_processor, tokenizer=tokenizer)
    elif model_args.model_backbone == QWEN2_5_OMNI:
        from src.model.vlm_backbone.qwen2_5_omni.processing_qwen2_5_omni import Qwen2_5OmniProcessor
        from src.model.vlm_backbone.qwen2_5_vl.image_processing_qwen2_5_vl import Qwen2_5_VLImageProcessor
        from src.model.vlm_backbone.qwen2_vl.tokenization_qwen2_fast import Qwen2TokenizerFast
        min_pixels, max_pixels = None, None
        if data_args is not None:
            min_pixels, max_pixels = data_args.resize_min_pixels, data_args.resize_max_pixels
        size = {"shortest_edge": min_pixels, "longest_edge": max_pixels, "min_pixels": min_pixels, "max_pixels": max_pixels}
        image_processor = Qwen2_5_VLImageProcessor.from_pretrained(model_name_or_path, size=size, local_files_only=HF_LOCAL_FILES_ONLY)
        tokenizer = Qwen2TokenizerFast.from_pretrained(model_name_or_path, local_files_only=HF_LOCAL_FILES_ONLY)
        processor = Qwen2_5OmniProcessor.from_pretrained(
            model_name_or_path,
            image_processor=image_processor, tokenizer=tokenizer,
            local_files_only=HF_LOCAL_FILES_ONLY,
        )
        # Enforce left padding for Qwen2.5 Omni Thinker (FlashAttention requirement)
        processor.tokenizer.padding_side = "left"
        if processor.tokenizer.pad_token is None:
            processor.tokenizer.pad_token = processor.tokenizer.eos_token
    elif model_args.model_backbone == WAVE:
        # WAVE baseline: uses WAVE's own processor (vlm_backbone/wave/)
        from src.model.vlm_backbone.wave.data.processing_qwen2_5_omni import Qwen2_5OmniProcessor as WaveProcessor
        from src.model.vlm_backbone.qwen2_vl.image_processing_qwen2_vl import Qwen2VLImageProcessor
        from src.model.vlm_backbone.qwen2_vl.tokenization_qwen2_fast import Qwen2TokenizerFast
        min_pixels, max_pixels = None, None
        if data_args is not None:
            min_pixels, max_pixels = data_args.resize_min_pixels, data_args.resize_max_pixels
        size = {"shortest_edge": min_pixels, "longest_edge": max_pixels, "min_pixels": min_pixels, "max_pixels": max_pixels}
        image_processor = Qwen2VLImageProcessor.from_pretrained(model_name_or_path, size=size)
        tokenizer = Qwen2TokenizerFast.from_pretrained(model_name_or_path)
        processor = WaveProcessor.from_pretrained(
            model_name_or_path,
            image_processor=image_processor, tokenizer=tokenizer
        )
        processor.tokenizer.padding_side = "left"
        if processor.tokenizer.pad_token is None:
            processor.tokenizer.pad_token = processor.tokenizer.eos_token
    else:
        from transformers import AutoProcessor
        processor = AutoProcessor.from_pretrained(
            model_args.processor_name if model_args.processor_name else model_args.model_name,
            trust_remote_code=True,
        )
    return processor


def get_backbone_name(hf_config, model_type=None):
    if model_type is not None:
        setattr(hf_config, 'model_type', model_type)
    assert hf_config.model_type in SUPPORTED_MODELS, f"Unknown backbone name {hf_config.model_type}.Supported models are {SUPPORTED_MODELS}"
    return MODEL2BACKBONE[hf_config.model_type]


def Phi3V_process_fn(model_inputs: dict, processor, max_length=None):
    input_ids, pixel_values, image_sizes, image_grid_thw = [], [], [], []
    texts, visual_inputs = model_inputs['text'], model_inputs['images']
    image_exists = False
    vlm_image_token = VLM_IMAGE_TOKENS[PHI3V]

    # 1. iterate each pair and process (since processors do not support batch processing)
    for text, images in zip(texts, visual_inputs):
        if images is None or (type(images)==list and any(i is None for i in images)):
            inputs = processor([text], None, return_tensors="np", max_length=max_length, truncation=True)
            input_id = inputs["input_ids"].squeeze().tolist()
            if isinstance(input_id, int):
                # in case of empty string, only BOS is included
                input_id = [input_id]
            input_ids.append(input_id)
            pixel_values.append(None)
            image_sizes.append(None)
            image_grid_thw.append(None)
        else:
            if vlm_image_token in text:
                if isinstance(images, PIL.Image.Image):
                    # images is a single image
                    images = [images]
                for iid, image in enumerate(images):
                    # rare case in MMEB eval: resize to 28*28 if either w or h is smaller than 28
                    if image.size[0] < 28 or image.size[1] < 28:
                        image = image.resize((56, 56))
                        images[iid] = image
                inputs = processor(text=text, images=images, return_tensors="np", max_length=max_length, truncation=True)
            else:
                raise NotImplementedError(f"No visual token found ({vlm_image_token}) in the text: {text}")
            # image_exists = True
            # inputs = processor(text=text, images=image, return_tensors="np", max_length=max_length, truncation=True)
            input_ids.append(inputs["input_ids"].squeeze().tolist())
            pixel_values.append(inputs['pixel_values'])
            if 'image_sizes' in inputs:
                image_sizes.append(inputs['image_sizes'])
            if 'image_grid_thw' in inputs:
                image_grid_thw.append(inputs['image_grid_thw'])

    # 2. padding inputs
    batch_encoding = processor.tokenizer.pad({'input_ids': input_ids}, return_tensors="pt")
    input_ids, attention_mask = batch_encoding['input_ids'], batch_encoding['attention_mask']
    inputs = {
        'input_ids': input_ids.long(),
        'attention_mask': attention_mask.long(),
        'texts': texts,
        'images': visual_inputs,
    }
    inputs['pixel_values'] = pixel_values
    inputs['image_sizes'] = image_sizes

    return inputs


def Qwen2_VL_process_fn(model_inputs: dict, processor, max_length=None):
    # TODO: set separate max_len for text/visual inputs, currently max_length is only applied to text-only data
    input_ids, pixel_values, image_grid_thw, pixel_values_videos, video_grid_thw = [], [], [], [], []
    raw_pixel_values = []
    texts, visual_inputs = model_inputs['text'], model_inputs['images']
    image_exists = False
    vlm_image_token, vlm_video_token = VLM_IMAGE_TOKENS[QWEN2_VL], VLM_VIDEO_TOKENS[QWEN2_VL]

    # 1. iterate each pair and process, since processors do not support processing for mixed batch (contains data w/ and w/o visual inputs)
    for text, images in zip(texts, visual_inputs):
        if images is None or (type(images)==list and any(i is None for i in images)):
            # all images must be valid
            inputs = processor(text=[text], images=None, return_tensors="np", max_length=max_length, truncation=True)
            input_id = inputs["input_ids"].squeeze().tolist()
            if isinstance(input_id, int):
                # in case of empty string, only BOS is included
                input_id = [input_id]
            input_ids.append(input_id)
            pixel_values.append(None)
            image_grid_thw.append(None)
            pixel_values_videos.append(None)
            video_grid_thw.append(None)
            raw_pixel_values.append(None)
        else:
            try:
                if vlm_image_token in text:
                    if isinstance(images, PIL.Image.Image):
                        # images is a single image
                        images = [images]
                    for iid, image in enumerate(images):
                        # rare case in MMEB eval: resize to 28*28 if either w or h is smaller than 28
                        if image.size[0] < 28 or image.size[1] < 28:
                            image = image.resize((56, 56))
                            images[iid] = image
                    inputs = processor(text=[text], images=images, return_tensors="np", max_length=None, truncation=False, input_data_format=ChannelDimension.LAST)
                elif vlm_video_token in text:
                    # TODO: check text/video data validity
                    inputs = processor(text=[text], videos=[images], return_tensors="np", max_length=None, truncation=False, input_data_format=ChannelDimension.LAST)
                else:
                    raise NotImplementedError(f"No visual token found ({vlm_image_token} or {vlm_video_token}) in the text: {text}")
            except Exception as e:
                for i in images:
                    print(i.filename)
                raise e
            input_ids.append(inputs["input_ids"].squeeze().tolist())
            if 'pixel_values' in inputs:
                pixel_values.append(inputs['pixel_values'])
                image_grid_thw.append(inputs['image_grid_thw'])
                pixel_values_videos.append(None)
                video_grid_thw.append(None)
                raw_pixel_values.append(inputs['raw_pixel_values'])
            else:
                pixel_values.append(None)
                image_grid_thw.append(None)
                pixel_values_videos.append(inputs['pixel_values_videos'])
                video_grid_thw.append(inputs['video_grid_thw'])
                raw_pixel_values.append(None)

    # 2. padding inputs
    batch_encoding = processor.tokenizer.pad({'input_ids': input_ids}, return_tensors="pt")
    input_ids, attention_mask = batch_encoding['input_ids'], batch_encoding['attention_mask']
    inputs = {
        'input_ids': input_ids.long(),
        'attention_mask': attention_mask.long(), 
        'texts': texts,
        'images': visual_inputs,
    }
    inputs['pixel_values'] = pixel_values
    inputs['image_grid_thw'] = image_grid_thw
    inputs['pixel_values_videos'] = pixel_values_videos
    inputs['video_grid_thw'] = video_grid_thw
    inputs['raw_pixel_values'] = raw_pixel_values

    return inputs


def Gme_process_fn(model_inputs: dict, processor, max_length=None):
    inputs = {
        'texts': model_inputs['text'],
        'images': model_inputs['images'],
    }
    return inputs

def Qwen2_5_Omni_process_fn(model_inputs: dict, processor, max_length=None):
    """
    Process batch inputs using direct batch processor call (similar to AVContrastiveCollator).
    This replaces the item-by-item processing with efficient batch processing.
    """
    texts, visual_inputs = model_inputs['text'], model_inputs['images']
    batch_audios = model_inputs.get('audios', None)
    use_audio_in_video = model_inputs.get('use_audio_in_video', False)
    batch_size = len(texts)
    
    # Use processor's tokens directly (same as what the processor uses internally)
    vlm_image_token = processor.image_token
    vlm_video_token = processor.video_token
    vlm_audio_token = processor.audio_token
    
    # Step 1: Separate modalities per sample and prepare batch inputs
    batch_images = []
    batch_videos = []
    batch_texts = texts.copy()
    
    # Track which samples have which modalities
    has_image = [False] * batch_size
    has_video = [False] * batch_size
    has_audio = [False] * batch_size
    
    for i, (text, images) in enumerate(zip(texts, visual_inputs)):
        if images is None or (isinstance(images, list) and any(img is None for img in images)):
            batch_images.append(None)
            batch_videos.append(None)
            if vlm_audio_token in text:
                has_audio[i] = True
        else:
            if vlm_image_token in text:
                if isinstance(images, PIL.Image.Image):
                    images = [images]
                processed_images = []
                for img in images:
                    if img.size[0] < 28 or img.size[1] < 28:
                        img = img.resize((56, 56))
                    processed_images.append(img)
                batch_images.append(processed_images if len(processed_images) > 1 else processed_images[0])
                batch_videos.append(None)
                has_image[i] = True
            elif vlm_video_token in text:
                batch_images.append(None)
                batch_videos.append(images)
                has_video[i] = True
            elif vlm_audio_token in text:
                batch_images.append(None)
                batch_videos.append(None)
                has_audio[i] = True
            else:
                raise NotImplementedError(f"No multimodal token found ({vlm_image_token}, {vlm_video_token}, or {vlm_audio_token}) in the text: {text}")
    
    # Step 2: Direct batch processor call
    processor_images = [img for img in batch_images if img is not None] if any(has_image) else None
    processor_videos = [vid for vid in batch_videos if vid is not None] if any(has_video) else None
    
    batch_inputs = processor(
        text=batch_texts,
        images=processor_images if processor_images else None,
        videos=processor_videos if processor_videos else None,
        audio=batch_audios,
        return_tensors="pt",
        padding=True,
        truncation=True if max_length is not None else False,
        max_length=max_length,
        use_audio_in_video=use_audio_in_video,
        input_data_format=ChannelDimension.LAST
    )
    
    # Step 3: Extract per-sample outputs from batch-level processor results
    # Processor returns batch-level tensors, but model expects per-sample lists
    pixel_values_list = [None] * batch_size
    image_grid_thw_list = [None] * batch_size
    pixel_values_videos_list = [None] * batch_size
    video_grid_thw_list = [None] * batch_size
    
    # Extract image outputs per sample
    if 'pixel_values' in batch_inputs and batch_inputs['pixel_values'] is not None:
        pixel_values_tensor = batch_inputs['pixel_values']
        image_grid_thw_tensor = batch_inputs.get('image_grid_thw', None)
        
        if image_grid_thw_tensor is not None:
            # Convert to tensor if numpy
            if isinstance(image_grid_thw_tensor, np.ndarray):
                image_grid_thw_tensor = torch.from_numpy(image_grid_thw_tensor)
            elif not isinstance(image_grid_thw_tensor, torch.Tensor):
                image_grid_thw_tensor = torch.tensor(image_grid_thw_tensor)
            
            # Convert pixel_values to tensor if needed
            if not isinstance(pixel_values_tensor, torch.Tensor):
                pixel_values_tensor = torch.from_numpy(pixel_values_tensor)
            
            # Split pixel_values based on image_grid_thw
            # Each image has t*h*w patches, so we need to split accordingly
            image_idx = 0
            patch_idx = 0
            for i in range(batch_size):
                if has_image[i]:
                    # Get grid dimensions for this image
                    if image_grid_thw_tensor.ndim == 2:
                        grid_thw = image_grid_thw_tensor[image_idx]
                    else:
                        grid_thw = image_grid_thw_tensor[image_idx]
                    num_patches = int(grid_thw[0] * grid_thw[1] * grid_thw[2])
                    
                    # Extract patches for this image
                    image_patches = pixel_values_tensor[patch_idx:patch_idx + num_patches]
                    
                    pixel_values_list[i] = image_patches
                    # Ensure grid_thw is a tensor with shape (1, 3) or (3,)
                    if grid_thw.ndim == 0:
                        grid_thw = grid_thw.unsqueeze(0)
                    if grid_thw.ndim == 1:
                        grid_thw = grid_thw.unsqueeze(0)
                    image_grid_thw_list[i] = grid_thw
                    
                    patch_idx += num_patches
                    image_idx += 1
    
    # Extract video outputs per sample
    if 'pixel_values_videos' in batch_inputs and batch_inputs['pixel_values_videos'] is not None:
        pixel_values_videos_tensor = batch_inputs['pixel_values_videos']
        video_grid_thw_tensor = batch_inputs.get('video_grid_thw', None)
        
        if video_grid_thw_tensor is not None:
            # Convert to tensor if numpy
            if isinstance(video_grid_thw_tensor, np.ndarray):
                video_grid_thw_tensor = torch.from_numpy(video_grid_thw_tensor)
            elif not isinstance(video_grid_thw_tensor, torch.Tensor):
                video_grid_thw_tensor = torch.tensor(video_grid_thw_tensor)
            
            # Convert pixel_values_videos to tensor if needed
            if not isinstance(pixel_values_videos_tensor, torch.Tensor):
                pixel_values_videos_tensor = torch.from_numpy(pixel_values_videos_tensor)
            
            # Split pixel_values_videos based on video_grid_thw
            video_idx = 0
            patch_idx = 0
            for i in range(batch_size):
                if has_video[i]:
                    # Get grid dimensions for this video
                    if video_grid_thw_tensor.ndim == 2:
                        grid_thw = video_grid_thw_tensor[video_idx]
                    else:
                        grid_thw = video_grid_thw_tensor[video_idx]
                    num_patches = int(grid_thw[0] * grid_thw[1] * grid_thw[2])
                    
                    # Extract patches for this video
                    video_patches = pixel_values_videos_tensor[patch_idx:patch_idx + num_patches]
                    
                    pixel_values_videos_list[i] = video_patches
                    # Ensure grid_thw is a tensor with shape (1, 3) or (3,)
                    if grid_thw.ndim == 0:
                        grid_thw = grid_thw.unsqueeze(0)
                    if grid_thw.ndim == 1:
                        grid_thw = grid_thw.unsqueeze(0)
                    video_grid_thw_list[i] = grid_thw
                    
                    patch_idx += num_patches
                    video_idx += 1
    
    # Extract audio outputs per sample
    input_features_list = [None] * batch_size
    feature_attention_mask_list = [None] * batch_size

    if 'input_features' in batch_inputs and batch_inputs['input_features'] is not None:
        input_features_tensor = batch_inputs['input_features']
        feature_attn_mask_tensor = batch_inputs.get('feature_attention_mask', None)

        if not isinstance(input_features_tensor, torch.Tensor):
            input_features_tensor = torch.from_numpy(input_features_tensor) if isinstance(input_features_tensor, np.ndarray) else torch.tensor(input_features_tensor)
        if feature_attn_mask_tensor is not None and not isinstance(feature_attn_mask_tensor, torch.Tensor):
            feature_attn_mask_tensor = torch.from_numpy(feature_attn_mask_tensor) if isinstance(feature_attn_mask_tensor, np.ndarray) else torch.tensor(feature_attn_mask_tensor)

        audio_idx = 0
        for i in range(batch_size):
            if has_audio[i] or (has_video[i] and use_audio_in_video):
                input_features_list[i] = input_features_tensor[audio_idx]
                if feature_attn_mask_tensor is not None:
                    feature_attention_mask_list[i] = feature_attn_mask_tensor[audio_idx]
                audio_idx += 1

    # Step 4: Build output dict matching expected format
    inputs = {
        'input_ids': batch_inputs['input_ids'].long(),
        'attention_mask': batch_inputs['attention_mask'].long(),
    }
    
    if any(pv is not None for pv in pixel_values_list):
        inputs['pixel_values'] = pixel_values_list
    if any(igt is not None for igt in image_grid_thw_list):
        inputs['image_grid_thw'] = image_grid_thw_list
    if any(pvv is not None for pvv in pixel_values_videos_list):
        inputs['pixel_values_videos'] = pixel_values_videos_list
    if any(vgt is not None for vgt in video_grid_thw_list):
        inputs['video_grid_thw'] = video_grid_thw_list
    if any(af is not None for af in input_features_list):
        ref_feature = next(x for x in input_features_list if x is not None)
        stacked = [x if x is not None else torch.zeros_like(ref_feature) for x in input_features_list]
        inputs['input_features'] = torch.stack(stacked)
    if any(fam is not None for fam in feature_attention_mask_list):
        ref_mask = next(x for x in feature_attention_mask_list if x is not None)
        stacked = [x if x is not None else torch.zeros_like(ref_mask) for x in feature_attention_mask_list]
        inputs['feature_attention_mask'] = torch.stack(stacked)
    if use_audio_in_video:
        inputs['use_audio_in_video'] = True
    
    # Include all other keys from batch_inputs that we haven't already handled
    handled_keys = {'pixel_values', 'image_grid_thw', 'pixel_values_videos', 'video_grid_thw',
                    'input_features', 'feature_attention_mask'}
    for key, value in batch_inputs.items():
        if key not in handled_keys:
            inputs[key] = value
    
    return inputs


def process_input_text(instruction, model_backbone, text=None, add_video_token=False, add_image_token=False):
    # Formulate input text based on text, special token and instruction.
    # TBD: Reorganize the hard-code part for baselines such as internvideo2
    if model_backbone == "internvideo2":
        return text
    elif model_backbone in [GME, LamRA, LamRA_QWEN2_5]:
        if text:
            return instruction + " " + text # GME and LamRA do not need special tokens
        else:
            return instruction + " "
    prompt = instruction
    if text:
        prompt = prompt + " " + text
    if add_video_token:
        video_token = VLM_VIDEO_TOKENS[model_backbone]
        prompt = video_token + " " + prompt
    if add_image_token:
        image_token = VLM_IMAGE_TOKENS[model_backbone]
        prompt = image_token + " " + prompt

    return prompt


process_vlm_inputs_fns = {
    PHI3V: Phi3V_process_fn,
    QWEN2_VL: Qwen2_VL_process_fn,
    QWEN2_5_VL: Qwen2_VL_process_fn,
    QWEN2_5_OMNI: Qwen2_5_Omni_process_fn,
    WAVE: Qwen2_5_Omni_process_fn,
    GME: Gme_process_fn,
    LamRA: Gme_process_fn,
    LamRA_QWEN2_5: Gme_process_fn,
}