import os
import shutil
from huggingface_hub import hf_hub_download, upload_folder, HfApi
#from loras import loras

# Your Hugging Face model repo
HF_REPO = "K00B404/LoraStack"


def download_lora(repo_id, save_dir="loras"):
    """Download all files from a LoRA repository."""
    print(f"lora repo {repo_id}")
    model_name = repo_id.split("/")[-1]
    target_dir = os.path.join(save_dir, model_name)
    os.makedirs(target_dir, exist_ok=True)
    
    # Get file list from repo
    api = HfApi()
    files = api.list_repo_files(repo_id, repo_type="model")
    
    for file in files:
        if not file.endswith('.png') and not file.endswith('.jpg') and not file.endswith('.webp'):
            print(f"Downloading {file} from {repo_id}...")
            hf_hub_download(repo_id=repo_id, filename=file, local_dir=target_dir)
    
    return target_dir

def upload_to_hf(local_dir, hf_repo):
    """Upload the downloaded LoRA folder to Hugging Face."""
    print(f"Uploading {local_dir} to {hf_repo}...")
    upload_folder(folder_path=local_dir, repo_id=hf_repo, repo_type="model")

def main(loras):
    """Fetch, download, and upload LoRAs to Hugging Face."""
    for repo_id in loras:
        local_dir = download_lora(repo_id)
        upload_to_hf(local_dir, HF_REPO)
        print(f"Successfully uploaded {repo_id} to {HF_REPO}\n")
        
    print("All LoRAs processed!")

if __name__ == "__main__":
    main()