import modal
import os
from pathlib import Path

# Create app for downloading files
download_app = modal.App("openlipsync-downloader")
download_image = (
    modal.Image.debian_slim(python_version="3.11")
    .pip_install("huggingface_hub[hf_transfer]==0.26.2")
    .env({"HF_HUB_ENABLE_HF_TRANSFER": "1"})
)

# Create app for listing files
list_app = modal.App("openlipsync-lister")

# Create a shared volume
volume = modal.Volume.from_name("openlipsync-volume", create_if_missing=True)

# List of files to download
files_to_download = [
    "checkpoints/whisper/tiny.pt",
    "checkpoints/latentsync/latentsync_syncnet.pt",
    "checkpoints/latentsync/README.md",
    "checkpoints/latentsync/latentsync_unet.pt",
    "checkpoints/latentsync/config.json",
    "checkpoints/sd-vae-ft-mse/README.md",
    "checkpoints/sd-vae-ft-mse/diffusion_pytorch_model.safetensors",
    "checkpoints/sd-vae-ft-mse/config.json",
    "checkpoints/sd-vae-ft-mse/diffusion_pytorch_model.bin",
    "checkpoints/auxiliary/i3d_torchscript.pt",
    "checkpoints/auxiliary/sfd_face.pth",
    "checkpoints/auxiliary/vit_g_hybrid_pt_1200e_ssv2_ft.pth",
    "checkpoints/auxiliary/koniq_pretrained.pkl",
    "checkpoints/auxiliary/s3fd-619a316812.pth",
    "checkpoints/auxiliary/2DFAN4-cd938726ad.zip",
    "checkpoints/auxiliary/vgg16-397923af.pth",
    "checkpoints/auxiliary/syncnet_v2.model",
    "configs/syncnet/syncnet_16_pixel.yaml",
    "configs/syncnet/syncnet_25_pixel.yaml",
    "configs/syncnet/syncnet_16_latent.yaml",
    "configs/unet/first_stage.yaml",
    "configs/unet/second_stage.yaml",
    "configs/scheduler_config.json",
    "configs/audio.yaml",
]

# Approach using individual file downloads
@download_app.function(
    image=download_image,
    volumes={"/data": volume},
    timeout=3600
)
def download_files():
    from huggingface_hub import hf_hub_download
    import os
    
    repo_id = "miguelamendez/openlipsync"
    
    # Download each file individually
    for file_path in files_to_download:
        try:
            # Download file to a temporary location
            local_file = hf_hub_download(
                repo_id=repo_id,
                filename=file_path,
                revision="main"
            )
            
            # Copy to volume
            with volume.batch_upload() as batch:
                batch.put_file(local_file, f"/data/{file_path}")
            
            # Clean up temporary file
            os.remove(local_file)
            
            print(f"Downloaded: {file_path}")
        except Exception as e:
            print(f"Error downloading {file_path}: {e}")
    
    # Commit changes to volume
    volume.commit()
    print("All files downloaded and committed to volume")

@download_app.local_entrypoint()
def main_download():
    print("Starting download...")
    download_files.remote()
    print("Download completed!")