replace.py

# this is the main file
import shutil
import tempfile
from time import sleep
import tools.check_necessary_files as file_checker
import tools.index_folder as index_folder
import tools.extract_zip as extract_zip
import tools.index_wads as index_wads
import tools.extract_wad as extract_wad
import tools.index_bin as index_bin
import tools.run_ritobin as run_ritobin
import tools.index_py as index_py
import tools.search_replace as search_replace
import tools.config as config
import tools.search_wad_directories as search_wad_directories
import tools.compress_to_wad as compress_to_wad
import multiprocessing
from pathlib import Path
import argparse
import os

def setup_argparse():
    parser = argparse.ArgumentParser(description='Process game files and perform search/replace operations')
    parser.add_argument('--config', default='replace_config.json', help='Path to config file')
    parser.add_argument('--dry-run', action='store_true', help='Show matches without making changes')
    return parser

def extract_wad_worker(wad_file):
    return extract_wad.run_wad_extractor(wad_file[1])

def run_ritobin_worker(bin_file):
    return run_ritobin.run_ritobin(bin_file)

def run_parallel_process(items, worker_func, process_name):
    if len(items) > 1:
        num_cores = max(1, multiprocessing.cpu_count() - 1)
        print(f"Processing {len(items)} {process_name} using {num_cores} cores")
        
        try:
            with multiprocessing.Pool(processes=num_cores) as pool:
                results = pool.map(worker_func, items)
                return results
        except Exception as e:
            print(f"Error during parallel {process_name} processing: {str(e)}")
            return None
    else:
        return [worker_func(items[0])]

def clean_wad_directories(directories, modified_files_set=None):
    """
    Delete all files except skin0.bin in each WAD directory.
    Also remove any .bak backup files.
    Only process directories from modified_files_set if provided.
    """
    print("\nCleaning WAD directories - keeping only skin0.bin files...")
    processed_count = 0
    skipped_count = 0
    
    for wad_dir in directories:
        try:
            wad_path = Path(wad_dir)
            
            # Skip if this directory is not in modified_files_set (when provided)
            if modified_files_set is not None:
                # Find the skin directory this WAD belongs to
                current = wad_path
                skin_dir = None
                while current.name != "extracted_zips" and current.parent != current:
                    if current.parent.name == "extracted_zips":
                        skin_dir = current
                        break
                    current = current.parent
                
                if skin_dir and str(skin_dir) not in modified_files_set:
                    skipped_count += 1
                    print(f"Skipping unmodified WAD directory: {wad_dir}")
                    continue
            
            if wad_path.exists():
                file_count = 0
                deleted_count = 0
                
                for file_path in wad_path.glob('**/*'):
                    if file_path.is_file():
                        file_count += 1
                        # Keep only skin0.bin files, delete everything else including .bak files
                        if file_path.name != 'skin0.bin':
                            file_path.unlink()
                            deleted_count += 1
                
                processed_count += 1
                print(f"Cleaned {wad_dir}: Kept 'skin0.bin', deleted {deleted_count} of {file_count} files")
        except Exception as e:
            print(f"Error cleaning {wad_dir}: {str(e)}")
    
    print(f"Total: {processed_count} WAD directories processed, {skipped_count} directories skipped")

if __name__ == "__main__":
    # Clean up previous run directories
    print("Cleaning up previous run directories...")
    extracted_zips_dir = Path("./extracted_zips")
    output_dir = Path("./output")
    
    if extracted_zips_dir.exists():
        print(f"Removing existing {extracted_zips_dir} directory...")
        try:
            shutil.rmtree(extracted_zips_dir)
            print(f"Successfully removed {extracted_zips_dir}")
        except Exception as e:
            print(f"Error removing {extracted_zips_dir}: {str(e)}")
    
    if output_dir.exists():
        print(f"Removing existing {output_dir} directory...")
        try:
            shutil.rmtree(output_dir)
            print(f"Successfully removed {output_dir}")
        except Exception as e:
            print(f"Error removing {output_dir}: {str(e)}")
    
    checker = file_checker.FileChecker()
    if not checker.verify_and_setup():
        print("Missing required files. Please ensure all necessary files are present.")
        exit(1)
    parser = setup_argparse()
    args = parser.parse_args()

    # Index all archives in the current directory
    archives = index_folder.index_archives("./skins")
    
    print("Archives found:")
    print(archives)
    

    # Extract all archives
    for archive in archives:
        extract_zip.extract_zip_file(archive[1])

    # Index all .wad.client files in the extracted_zips directory
    wad_files = index_wads.find_wad_client_files("extracted_zips")

    # Process WAD files in parallel
    wad_results = run_parallel_process(
        [wf[1] for wf in wad_files], 
        extract_wad.run_wad_extractor, 
        "WAD files"
    )

    # Index all skin0.bin files in the extracted_zips directory
    print("Searching for skin0.bin files...")
    bin_files = index_bin.find_skin_bin_files("extracted_zips")
    
    print("skin0.bin files found:")
    print(bin_files)
    print("Quantity of skin0.bin files found:")
    print(len(bin_files))

    # Process bin files in parallel
    bin_results = run_parallel_process(bin_files, run_ritobin_worker, "skin0.bin files")
    
    if bin_results:
        for bin_file, (success, message) in zip(bin_files, bin_results):
            print(f"Processing {bin_file}: {message}")
            if not success:
                break

    # Index all .py files in the extracted_zips directory
    py_files = index_py.find_py_files("extracted_zips")
    print("py_files found:")
    print(py_files)
    print("Quantity of py_files found:")
    print(len(py_files))

    # Search and Replace operation
    modified_files = set()  # Keep track of modified files
    if py_files:
        try:
            # Load configuration
            cfg = config.load_config(args.config)
            
            # Validate configuration
            if not cfg['search_pattern'] or not cfg['replacement']:
                print("Please configure search_pattern and replacement in replace_config.json")
                print("Example config:")
                print("""
                {
                    "search_pattern": "your_pattern_here",
                    "replacement": "your_replacement_here",
                    "backup": true,
                    "encoding": "utf-8"
                }
                """)
                exit(1)
            
            # Perform dry run first
            print("\nPerforming dry run to show matches...")
            results = search_replace.search_and_replace_in_files(
                py_files, 
                cfg,
                dry_run=True
            )
            
            # Display results
            print("\nMatches found:")
            for filepath, matched_lines, count, _ in results:
                if count > 0:
                    print(f"\nFile: {filepath}")
                    print(f"Matches: {count}")
                    print("Context:")
                    for line in matched_lines:
                        print(f"  {line}\n")
            
            # Ask for confirmation if not in dry-run mode
            if not args.dry_run:
                total_matches = sum(count for _, _, count, _ in results)
                if total_matches > 0:
                    response = input(f"\nFound {total_matches} matches. Proceed with replacement? (y/n): ")
                    if response.lower() == 'y':
                        results = search_replace.search_and_replace_in_files(
                            py_files,
                            cfg,
                            dry_run=False
                        )
                        print("\nReplacement complete. Backup files created.")
                        for filepath, _, count, backup_path in results:
                            if count > 0:
                                print(f"Modified: {filepath}")
                                print(f"Backup: {backup_path}")
                                
                                # Add to modified files set - Store the base skin directory
                                # Extract the skin directory from the filepath
                                file_path = Path(filepath)
                                # Navigate up to the extracted_zips directory
                                skin_dir = None
                                current = file_path
                                while current.name != "extracted_zips" and current.parent != current:
                                    if current.parent.name == "extracted_zips":
                                        skin_dir = current
                                        break
                                    current = current.parent
                                
                                if skin_dir:
                                    modified_files.add(str(skin_dir))
                                    print(f"Added to processing list: {skin_dir}")
                else:
                    print("No matches found.")
            
        except Exception as e:
            print(f"Error during search and replace: {str(e)}")
        if py_files:  # Only if we processed files
            print("\n" + "="*50)
            print("All operations completed.")
            print("="*50)

    # Process .py files with ritobin
    print("\nProcessing .py files with ritobin...")
    py_file_paths = [filepath for _, filepath in py_files]  # Extract just the filepaths
    
    # Process py files in parallel
    py_ritobin_results = run_parallel_process(
        py_file_paths,
        run_ritobin.run_ritobin,
        "Python files"
    )
    
    if py_ritobin_results:
        for py_file, (success, message) in zip(py_file_paths, py_ritobin_results):
            print(f"Processing {py_file}: {message}")
            if not success:
                break
    
    print("Searching for WAD directories...")
    wad_directories = search_wad_directories.find_wad_directories("extracted_zips")
    print("\nWAD directories found:")
    print(wad_directories)
    print("Quantity of WAD directories found:")
    print(len(wad_directories))

    # NEW: Clean WAD directories - keep only skin0.bin files
    clean_wad_directories(wad_directories, modified_files if modified_files else None)

    # Filter WAD directories to only process modified ones
    if modified_files:
        filtered_wad_directories = []
        for wad_dir in wad_directories:
            wad_path = Path(wad_dir)
            current = wad_path
            skin_dir = None
            while current.name != "extracted_zips" and current.parent != current:
                if current.parent.name == "extracted_zips":
                    skin_dir = current
                    break
                current = current.parent
            
            if skin_dir and str(skin_dir) in modified_files:
                filtered_wad_directories.append(wad_dir)
        
        print(f"Filtered: Processing {len(filtered_wad_directories)} of {len(wad_directories)} WAD directories")
        wad_directories = filtered_wad_directories

    # Process WAD directories in parallel
    wad_results = run_parallel_process(
        wad_directories,
        compress_to_wad.compress_to_wad,
        "WAD directories"
    )

    # delete all .wad directorties and everyting inside
    print("\nRemoving processed WAD directories...")
    for wad_dir in wad_directories:
        try:
            wad_path = Path(wad_dir)
            if wad_path.exists():
                shutil.rmtree(wad_path)
                print(f"Removed: {wad_dir}")
        except Exception as e:
            print(f"Error removing {wad_dir}: {str(e)}")

    
    # Create output directory if it doesn't exist
    output_dir = Path("./output")
    output_dir.mkdir(exist_ok=True)
    
    # First, organize archives by their parent directory
    archive_groups = {}
    for original_archive, original_path in archives:
        # Debug print
        print(f"\nProcessing archive: {original_path}")
        
        # Normalize path handling
        original_path = Path(original_path).resolve()
        parent_dir = original_path.parent.name
        
        if parent_dir not in archive_groups:
            archive_groups[parent_dir] = []
        archive_groups[parent_dir].append((original_archive, str(original_path)))
    
    # Process each group
    for parent_dir, group_archives in archive_groups.items():
        print(f"\nProcessing group: {parent_dir}")
        
        # Create parent directory in output if needed
        parent_output_dir = output_dir / parent_dir
        parent_output_dir.mkdir(exist_ok=True)
        
        for original_archive, original_path in group_archives:
            try:
                # Convert and normalize paths
                original_path = Path(original_path)
                is_fantome = original_path.suffix.lower() == '.fantome'
                
                # Get the base name without extension
                base_name = original_path.stem
                if is_fantome:
                    base_name = Path(base_name).stem
                
                # Create new archive path
                new_archive_path = parent_output_dir / original_archive
                
                # Construct and verify source directory path
                source_dir = Path("extracted_zips") / base_name
                source_dir = source_dir.resolve()  # Resolve to absolute path
                
                print(f"Checking source directory: {source_dir}")
                print(f"Directory exists: {source_dir.exists()}")
                
                if not source_dir.exists():
                    # Try alternate path construction
                    alt_source_dir = Path("extracted_zips") / parent_dir / base_name
                    alt_source_dir = alt_source_dir.resolve()
                    print(f"Trying alternate path: {alt_source_dir}")
                    print(f"Alternate exists: {alt_source_dir.exists()}")
                    
                    if alt_source_dir.exists():
                        source_dir = alt_source_dir
                    else:
                        print(f"Warning: Neither path exists for: {original_archive}")
                        print(f"Tried:\n1. {source_dir}\n2. {alt_source_dir}")
                        continue
                
                # Skip if not in modified_files (when we have modifications)
                if modified_files and str(source_dir) not in modified_files:
                    print(f"Skipping unmodified archive: {original_archive}")
                    continue
                
                print(f"Creating archive: {new_archive_path}")
                
                if is_fantome:
                    # Handle .fantome files
                    temp_zip_path = parent_output_dir / f"{base_name}.zip"
                    
                    with tempfile.TemporaryDirectory() as temp_dir:
                        temp_dir_path = Path(temp_dir)
                        
                        # Copy directories with explicit existence checks
                        meta_dir = source_dir / "META"
                        wad_dir = source_dir / "WAD"
                        
                        if meta_dir.exists():
                            shutil.copytree(meta_dir, temp_dir_path / "META")
                        if wad_dir.exists():
                            shutil.copytree(wad_dir, temp_dir_path / "WAD")
                        
                        # Create zip archive
                        shutil.make_archive(
                            str(temp_zip_path.with_suffix('')),
                            'zip',
                            root_dir=str(temp_dir_path),
                            base_dir='.'
                        )
                    
                    # Rename to .fantome if zip was created
                    if temp_zip_path.exists():
                        temp_zip_path.rename(new_archive_path)
                else:
                    # Handle regular zip files
                    with tempfile.TemporaryDirectory() as temp_dir:
                        temp_dir_path = Path(temp_dir)
                        
                        # Copy directories with explicit existence checks
                        meta_dir = source_dir / "META"
                        wad_dir = source_dir / "WAD"
                        
                        if meta_dir.exists():
                            shutil.copytree(meta_dir, temp_dir_path / "META")
                        if wad_dir.exists():
                            shutil.copytree(wad_dir, temp_dir_path / "WAD")
                        
                        # Create zip archive
                        shutil.make_archive(
                            str(new_archive_path.with_suffix('')),
                            'zip',
                            root_dir=str(temp_dir_path),
                            base_dir='.'
                        )
                
                print(f"Successfully created: {new_archive_path}")
                
            except Exception as e:
                print(f"Error processing {original_archive}: {str(e)}")
                import traceback
                traceback.print_exc()
        
    print("\nCompression completed. Output files are in the ./output directory")
    # ...existing code...print("All operations completed.")
    print('='*50)
    print("Made by @koobzaar")
    print("Any issues or feedback, please contact support@darkseal.org")