#!/usr/bin/env python3
"""
VILA Dataset Path Patcher (Optimized)

Only patches VILA dataset __init__ methods to remap paths.
Does NOT patch builtins.open for better multi-process compatibility.

Usage:
    Add path mappings to PATH_MAPPINGS dict, then import this module before loading datasets.
"""

import os
from pathlib import Path

# Path mappings: hardcoded_path -> actual_path
PATH_MAPPINGS = {
    "/home/jasonlu/workspace/InternVL/internvl_chat/playground": 
        "/lustre/fsw/portfolios/nvr/projects/nvr_elm_llm/dataset/vila-sft/internvl_chat/playground",
    "/home/jasonlu/vlm_datasets2/":
        "/lustre/fsw/portfolios/nvr/projects/nvr_elm_llm/users/jasonlu/vlm_datasets/",
    # Add more mappings as needed:
    # "/home/user/old_path": "/lustre/new_path",
}

# Control verbosity
_VERBOSE = os.environ.get("VILA_PATH_PATCHER_VERBOSE", "0") == "1"
_remap_count = 0
_MAX_LOG_MESSAGES = 5  # Only log first 5 remappings


def remap_path(original_path):
    """Remap a path if it matches any prefix in PATH_MAPPINGS."""
    global _remap_count
    
    if not isinstance(original_path, (str, Path)):
        return original_path
    
    path_str = str(original_path)
    
    # Exact match
    if path_str in PATH_MAPPINGS:
        remapped = PATH_MAPPINGS[path_str]
        if _VERBOSE and _remap_count < _MAX_LOG_MESSAGES:
            print(f"[PathPatcher] Remapped: {path_str} -> {remapped}", flush=True)
            _remap_count += 1
        return remapped
    
    # Prefix match
    for old_prefix, new_prefix in PATH_MAPPINGS.items():
        if path_str.startswith(old_prefix):
            remapped = path_str.replace(old_prefix, new_prefix, 1)
            if _VERBOSE and _remap_count < _MAX_LOG_MESSAGES:
                print(f"[PathPatcher] Remapped: {path_str} -> {remapped}", flush=True)
                _remap_count += 1
            elif _remap_count == _MAX_LOG_MESSAGES:
                print(f"[PathPatcher] (suppressing further remap logs...)", flush=True)
                _remap_count += 1
            return remapped
    
    return original_path


def patch_vila_dataset_classes():
    """Patch VILA dataset __init__ methods to remap paths."""
    patched_count = 0
    
    # Patch LazySupervisedDataset
    try:
        from llava.data.dataset import LazySupervisedDataset
        original_init = LazySupervisedDataset.__init__
        
        def patched_init(self, data_path, image_folder, *args, **kwargs):
            data_path = remap_path(data_path)
            if image_folder:
                image_folder = remap_path(image_folder)
            original_init(self, data_path, image_folder, *args, **kwargs)
        
        LazySupervisedDataset.__init__ = patched_init
        patched_count += 1
        
    except (ImportError, AttributeError) as e:
        print(f"[PathPatcher] Could not patch LazySupervisedDataset: {e}", flush=True)
    
    # Patch LazyWDSDataset
    try:
        from llava.data.dataset import LazyWDSDataset
        original_init = LazyWDSDataset.__init__
        
        def patched_init(self, data_path, image_folder, *args, **kwargs):
            data_path = remap_path(data_path)
            if image_folder:
                image_folder = remap_path(image_folder)
            original_init(self, data_path, image_folder, *args, **kwargs)
        
        LazyWDSDataset.__init__ = patched_init
        patched_count += 1
        
    except (ImportError, AttributeError):
        pass
    
    # Patch LLaVADataset (uses media_dir instead of image_folder)
    try:
        from llava.data.dataset_impl.llava import LLaVADataset
        original_init = LLaVADataset.__init__
        
        def patched_init(self, data_path, media_dir=None, is_video=False, **kwargs):
            # Remap paths before passing to original __init__
            data_path = remap_path(data_path)
            if media_dir:
                media_dir = remap_path(media_dir)
            # Call original with exact signature
            original_init(self, data_path, media_dir=media_dir, is_video=is_video, **kwargs)
        
        LLaVADataset.__init__ = patched_init
        patched_count += 1
        
    except (ImportError, AttributeError) as e:
        print(f"[PathPatcher] Could not patch LLaVADataset: {e}", flush=True)
    
    # Don't patch BaseDataset - it doesn't have data_path/media_dir parameters
    # Those are in subclasses like LLaVADataset
    
    return patched_count


# Auto-apply patches when module is imported
print("[PathPatcher] Optimized path patcher (dataset __init__ only, no builtins.open)", flush=True)
print(f"[PathPatcher] Configured {len(PATH_MAPPINGS)} path mapping(s)", flush=True)

patched = patch_vila_dataset_classes()
print(f"[PathPatcher] ✅ Patched {patched} dataset classes", flush=True)
print("[PathPatcher] Note: builtins.open is NOT patched (better multi-process performance)", flush=True)

