from huggingface_hub import HfApi, HfFileSystem
import urllib.parse
import os
import sys, io
sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8')

api = HfApi()
fs = HfFileSystem()

def get_language(filename):
    name = urllib.parse.unquote(filename).lower()
    
    if 'english' in name or 'american' in name or 'british' in name or 'patois' in name or 'pidgin' in name or 'south african' in name:
        return 'English'
    if 'korean' in name:
        return 'Korean'
    if 'igbo' in name:
        return 'Asusu Igbo'
    if 'tagalog' in name or 'filipino' in name:
        return 'Filipino'
    if 'creole' in name:
        return 'Creole'
    if 'indonesian' in name:
        return 'Indonesian'
    
    return 'Other'

repo_id = "toecm/IEDID"

def migrate():
    files = fs.ls(f"datasets/{repo_id}", detail=False)
    # Strip prefix 'datasets/toecm/IEDID/'
    files = [f.replace(f"datasets/{repo_id}/", "") for f in files]
    
    operations = []
    
    for f in files:
        if f in ['.gitattributes', 'README.md', 'minted_history.csv', 'audio', 'audio_files']:
            continue
            
        lang = get_language(f)
        
        # Decide if it goes to /data/ or /personas/
        if f.endswith('.csv'):
            new_path = f"{lang}/data/{f}"
        elif f.endswith('.json'):
            new_path = f"{lang}/personas/{f}"
        else:
            continue
            
        print(f"Moving {f} -> {new_path}")
        operations.append({
            "path_in_repo": f,
            "path_in_repo_new": new_path,
        })
        
    if not operations:
        print("No files to migrate.")
        return
        
    print("Executing move operations...")
    for op in operations:
        try:
            # We use fs.mv for moving files on the Hugging Face filesystem
            src = f"datasets/{repo_id}/{op['path_in_repo']}"
            dest = f"datasets/{repo_id}/{op['path_in_repo_new']}"
            fs.mv(src, dest)
            print(f"Success: {op['path_in_repo']}")
        except Exception as e:
            print(f"Failed to move {op['path_in_repo']}: {e}")
            
if __name__ == "__main__":
    migrate()