"""
Setup script to copy and adapt source files from hg-research-hub to infinite-world.
This creates a standalone project without external dependencies.
"""

import os
import re
import shutil

# Source and target directories
SRC_BASE = '/mnt/dolphinfs/ssd_pool/docker/user/hadoop-videogen-hl/hadoop-camera3d/wuruiqi/hg-research-hub/source'
DST_BASE = '/mnt/dolphinfs/ssd_pool/docker/user/hadoop-videogen-hl/hadoop-camera3d/wuruiqi/infinite-world/infworld'

# Files to copy and their destination subdirectories
FILES_TO_COPY = {
    # Models
    'meigen/model_wanx_multi_action_v2v_convenc_locmem_slidewindow_temp_sample_mask_attn_real_checkpointing.py': 'models/dit_model.py',
    'meigen/rectified_flow_wanx_t2v_action.py': 'models/scheduler.py',
    'meigen/checkpoint.py': 'models/checkpoint.py',
    'meigen/umt5.py': 'models/umt5.py',
    'meigen/t5.py': 'models/t5.py',
    
    # VAE
    'vae/__init__.py': 'vae/__init__.py',
    'vae/wan/vae.py': 'vae/vae.py',
    
    # CLIP
    'clip/clip.py': 'clip/clip.py',
    'clip/tokenizers.py': 'clip/tokenizers.py',
    'clip/xlm_roberta.py': 'clip/xlm_roberta.py',
    
    # Context Parallel
    'context_parallel/context_parallel_util.py': 'context_parallel/context_parallel_util.py',
    
    # Utils
    'dataset/utils.py': 'utils/data_utils.py',
    'dataset/prepare_dataloader.py': 'utils/prepare_dataloader.py',
    
    # OpenSora (for registry and dataset utils)
    'opensora/utils/dataset_utils.py': 'utils/dataset_utils.py',
    'opensora/registry.py': 'utils/registry.py',
}

# Import replacements (old pattern -> new pattern)
IMPORT_REPLACEMENTS = [
    # Models
    (r'from source\.meigen\.checkpoint', 'from infworld.models.checkpoint'),
    (r'from source\.meigen\.model_wanx_multi_action', 'from infworld.models.dit_model'),
    (r'from source\.meigen\.rectified_flow_wanx_t2v_action', 'from infworld.models.scheduler'),
    (r'from source\.meigen\.umt5', 'from infworld.models.umt5'),
    (r'from source\.meigen\.t5', 'from infworld.models.t5'),
    (r'from source\.meigen', 'from infworld.models'),
    
    # Context Parallel
    (r'from source\.context_parallel\.context_parallel_util', 'from infworld.context_parallel.context_parallel_util'),
    (r'from source\.context_parallel import context_parallel_util', 'from infworld.context_parallel import context_parallel_util'),
    
    # VAE
    (r'from source\.vae\.wan\.vae', 'from infworld.vae.vae'),
    (r'from source\.vae\.cogvideo\.autoencoder_kl_cogvideox', 'from infworld.vae.vae'),
    (r'from source\.vae', 'from infworld.vae'),
    (r'from source\.opensora\.registry import MODELS', '# Registry disabled for standalone'),
    
    # CLIP
    (r'from source\.clip\.clip', 'from infworld.clip.clip'),
    (r'from source\.clip\.tokenizers', 'from infworld.clip.tokenizers'),
    (r'from source\.clip\.xlm_roberta', 'from infworld.clip.xlm_roberta'),
    (r'from source\.clip', 'from infworld.clip'),
    
    # Dataset utils
    (r'from source\.dataset\.utils', 'from infworld.utils.data_utils'),
    (r'from source\.dataset\.prepare_dataloader', 'from infworld.utils.prepare_dataloader'),
    (r'from source\.opensora\.utils\.dataset_utils', 'from infworld.utils.dataset_utils'),
    (r'from source\.opensora\.registry', 'from infworld.utils.registry'),
]


def ensure_dir(path):
    """Create directory if it doesn't exist."""
    os.makedirs(os.path.dirname(path), exist_ok=True)


def copy_and_transform(src_path, dst_path):
    """Copy file and transform imports."""
    print(f"Copying: {src_path} -> {dst_path}")
    
    ensure_dir(dst_path)
    
    with open(src_path, 'r', encoding='utf-8') as f:
        content = f.read()
    
    # Apply import replacements
    for old_pattern, new_pattern in IMPORT_REPLACEMENTS:
        content = re.sub(old_pattern, new_pattern, content)
    
    with open(dst_path, 'w', encoding='utf-8') as f:
        f.write(content)


def create_init_files():
    """Create __init__.py files for all packages."""
    packages = ['infworld', 'infworld/models', 'infworld/vae', 'infworld/clip', 
                'infworld/context_parallel', 'infworld/utils', 'infworld/configs']
    
    for pkg in packages:
        init_path = os.path.join(DST_BASE, '..', pkg, '__init__.py')
        init_path = os.path.normpath(init_path)
        ensure_dir(init_path)
        
        if not os.path.exists(init_path):
            with open(init_path, 'w') as f:
                f.write(f'# {pkg} package\n')
            print(f"Created: {init_path}")


def main():
    print("=" * 60)
    print("Setting up Infinite World standalone project")
    print("=" * 60)
    
    # Create package directories
    create_init_files()
    
    # Copy and transform files
    for src_rel, dst_rel in FILES_TO_COPY.items():
        src_path = os.path.join(SRC_BASE, src_rel)
        dst_path = os.path.join(DST_BASE, dst_rel)
        
        if os.path.exists(src_path):
            copy_and_transform(src_path, dst_path)
        else:
            print(f"WARNING: Source file not found: {src_path}")
    
    print("\n" + "=" * 60)
    print("Setup complete!")
    print("=" * 60)


if __name__ == '__main__':
    main()