import shutil
import requests
from pathlib import Path
from zipfile import ZipFile
from io import BytesIO

def download_file(url: str, destination: Path) -> None:
    """Download a file from URL to the specified destination."""
    response = requests.get(url)
    response.raise_for_status()
    destination.write_bytes(response.content)

def download_github_folder(repo_owner: str, repo_name: str, folder_path: str, destination: Path) -> None:
    """Download a specific folder from a GitHub repository using the ZIP download feature."""
    # Download the whole repository as a ZIP file
    zip_url = f"https://github.com/{repo_owner}/{repo_name}/archive/refs/heads/main.zip"
    response = requests.get(zip_url)
    response.raise_for_status()

    # Extract only the needed folder from the ZIP
    with ZipFile(BytesIO(response.content)) as zip_file:
        folder_prefix = f"{repo_name}-main/{folder_path}"
        # Extract only files from the specified folder
        for file in zip_file.namelist():
            if file.startswith(folder_prefix):
                # Remove the repository name and branch prefix from the path
                relative_path = file.replace(f"{repo_name}-main/", "", 1)
                if relative_path.endswith('/'):  # Skip directory entries
                    continue
                
                # Read the file content from ZIP
                content = zip_file.read(file)
                
                # Create the file path
                output_path = destination / relative_path
                output_path.parent.mkdir(parents=True, exist_ok=True)
                
                # Write the file
                output_path.write_bytes(content)

def setup_imagebind():
    """Setup ImageBind by downloading only required files."""
    # Create clean imagebind directory if needed
    imagebind_dir = Path("imagebind")
    if imagebind_dir.exists():
        shutil.rmtree(imagebind_dir)
    
    # Download the imagebind folder
    download_github_folder(
        repo_owner="facebookresearch",
        repo_name="ImageBind",
        folder_path="imagebind",
        destination=Path(".")
    )
    
    # Download setup.py file
    setup_py_url = "https://raw.githubusercontent.com/facebookresearch/ImageBind/main/setup.py"
    setup_py_path = Path("setup.py")
    download_file(setup_py_url, setup_py_path)

if __name__ == "__main__":
    setup_imagebind()