#!/usr/bin/env python3
"""
Script to modify replicate.py for downloading data at runtime.

This script adds data downloading capability to the Replicate predictor,
allowing it to fetch data files during initialization rather than bundling them.
"""

import os
import re
import sys
import shutil

# The file to modify
TARGET_FILE = "replicate.py"
BACKUP_FILE = "replicate.py.bak"

# The download code to insert
DOWNLOAD_CODE = """
        # Download data files if needed
        try:
            import subprocess
            
            print("Checking data files...")
            
            # Run the download script
            download_result = subprocess.run(
                [sys.executable, "download_data.py", "--skip-verify"],
                capture_output=True,
                text=True
            )
            
            if download_result.returncode != 0:
                print(f"Warning: Data download script returned non-zero exit code: {download_result.returncode}")
                print(f"Error output: {download_result.stderr}")
            else:
                print("Data files are ready")
                
        except Exception as e:
            print(f"Warning: Failed to download data files: {e}")
            # Continue anyway - we'll check for files below
"""

def backup_file(file_path):
    """Create a backup of the file."""
    if os.path.exists(file_path):
        shutil.copy2(file_path, BACKUP_FILE)
        print(f"Created backup: {BACKUP_FILE}")
        return True
    else:
        print(f"Error: File {file_path} not found")
        return False

def modify_replicate_file():
    """Modify the replicate.py file to add data downloading."""
    if not os.path.exists(TARGET_FILE):
        print(f"Error: {TARGET_FILE} not found")
        return False
    
    # Make a backup first
    if not backup_file(TARGET_FILE):
        return False
    
    # Read the file
    with open(TARGET_FILE, "r") as f:
        content = f.read()
    
    # Look for the setup method
    setup_pattern = r"(def setup\(self\):.*?# Check data files exist.*?for file_path in required_files:.*?raise FileNotFoundError.*?)\n"
    
    # Replace with modified code - keep the original check but add download before it
    if re.search(setup_pattern, content, re.DOTALL):
        modified_content = re.sub(
            setup_pattern,
            r"\1\n" + DOWNLOAD_CODE + "\n",
            content,
            flags=re.DOTALL
        )
        
        # Write the modified content
        with open(TARGET_FILE, "w") as f:
            f.write(modified_content)
        
        print(f"Successfully modified {TARGET_FILE} to include data downloading")
        return True
    else:
        print(f"Error: Could not find the setup method pattern in {TARGET_FILE}")
        return False

def main():
    """Main function."""
    print(f"Modifying {TARGET_FILE} to add data downloading capabilities...")
    success = modify_replicate_file()
    
    if success:
        print("Modification complete!")
        return 0
    else:
        print("Modification failed. See errors above.")
        return 1

if __name__ == "__main__":
    sys.exit(main())