import os
import subprocess
from concurrent.futures import ThreadPoolExecutor

# --- Configuration ---
MANIFEST_ROOT = "/root/hf_repo/weights"
FSX_ROOT = "/fsx_scanline/from_eyeline/koichi/motion-aware-vid-stylization/train_ckpt/Netflix65K"
DEST_ROOT = "/hf_repo/weights"

def find_copy_tasks():
    """
    1. Scans MANIFEST_ROOT for model folders.
    2. Scans inside those folders for specific .safetensors files.
    Returns: list of (model_alias, ckpt_folder, filename)
    """
    tasks = []
    
    if not os.path.exists(MANIFEST_ROOT):
        print(f"Error: Manifest root {MANIFEST_ROOT} does not exist.")
        return []

    print(f"Scanning {MANIFEST_ROOT} for specific safetensor files...")
    
    try:
        model_aliases = [d for d in os.listdir(MANIFEST_ROOT) if os.path.isdir(os.path.join(MANIFEST_ROOT, d))]
    except OSError as e:
        print(f"Error reading manifest root: {e}")
        return []

    for alias in model_aliases:
        alias_path = os.path.join(MANIFEST_ROOT, alias)
        
        try:
            # Find checkpoint folders (e.g. AllShrink_...)
            checkpoints = [d for d in os.listdir(alias_path) if os.path.isdir(os.path.join(alias_path, d))]
        except OSError:
            continue
            
        for ckpt_folder in checkpoints:
            ckpt_full_path = os.path.join(alias_path, ckpt_folder)
            
            # --- CRITICAL CHANGE ---
            # We look inside the folder to see WHICH safetensors file is there
            try:
                files_in_ckpt = os.listdir(ckpt_full_path)
                target_files = [f for f in files_in_ckpt if f.endswith('.safetensors')]
                
                if not target_files:
                    # Optional: warning if a folder is empty
                    # print(f"Warning: No safetensors found in {alias}/{ckpt_folder}")
                    pass
                
                for filename in target_files:
                    # We create a task for this SPECIFIC file
                    tasks.append((alias, ckpt_folder, filename))
                    
            except OSError:
                continue
            
    print(f"Found {len(tasks)} specific files to copy.")
    return tasks

def process_copy(task):
    """
    Copies a SINGLE file from FSX to Destination using rclone copyto
    """
    alias, ckpt_folder, filename = task
    
    # 1. Source File (Flat FSX structure)
    # /fsx/.../Netflix65K/{ckpt_folder}/{filename}
    src_file_path = os.path.join(FSX_ROOT, ckpt_folder, filename)
    
    # 2. Dest File (Nested structure)
    # /hf_repo/weights/{alias}/{ckpt_folder}/{filename}
    dest_file_path = os.path.join(DEST_ROOT, alias, ckpt_folder, filename)
    
    # Check if the specific file exists on FSX source
    if not os.path.exists(src_file_path):
        print(f"[{alias}] MISSING SOURCE: {filename} (Expected at {src_file_path})")
        return

    # 3. Use 'copyto' for single files
    # This is more precise than 'copy'
    cmd = [
        "rclone", "copyto",
        src_file_path,
        dest_file_path,
        "--transfers", "4",
        "--progress" 
    ]

    try:
        result = subprocess.run(cmd, capture_output=True, text=True)
        
        if result.returncode == 0:
            print(f"[{alias}] SUCCESS: {filename}")
        else:
            print(f"[{alias}] ERROR copying {filename}: {result.stderr.strip()}")
            
    except Exception as e:
        print(f"[{alias}] CRASHED on {filename}: {e}")

if __name__ == "__main__":
    all_tasks = find_copy_tasks()
    
    # Parallel execution
    # Since we are copying individual large files, 4-8 workers is usually optimal
    if all_tasks:
        print(f"Starting parallel copy with {min(len(all_tasks), 8)} workers...")
        with ThreadPoolExecutor(max_workers=8) as executor:
            executor.map(process_copy, all_tasks)
            
    print("All tasks completed.")