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

def download_file(url: str, destination: Path) -> None:
    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_root: Path) -> None:
    zip_url = f"https://github.com/{repo_owner}/{repo_name}/archive/refs/heads/main.zip"
    response = requests.get(zip_url)
    response.raise_for_status()

    target_folder_path_in_repo = f"{repo_name}-main/{folder_path}/"
    local_target_folder = destination_root / folder_path

    with ZipFile(BytesIO(response.content)) as zip_file:
        for member_info in zip_file.infolist():
            if member_info.filename.startswith(target_folder_path_in_repo) and not member_info.is_dir():
                relative_path = member_info.filename.replace(target_folder_path_in_repo, "", 1)
                output_path = local_target_folder / relative_path
                output_path.parent.mkdir(parents=True, exist_ok=True)
                output_path.write_bytes(zip_file.read(member_info.filename))

def setup_imagebind():
    current_dir = Path(".")
    
    imagebind_code_dir = current_dir / "imagebind"
    if imagebind_code_dir.exists():
        shutil.rmtree(imagebind_code_dir)
    
    download_github_folder(
        repo_owner="facebookresearch",
        repo_name="ImageBind",
        folder_path="imagebind",
        destination_root=current_dir
    )
    
    setup_py_url = "https://raw.githubusercontent.com/facebookresearch/ImageBind/main/setup.py"
    setup_py_local_path = current_dir / "setup.py"
    download_file(setup_py_url, setup_py_local_path)

    readme_url = "https://raw.githubusercontent.com/facebookresearch/ImageBind/main/README.md"
    readme_local_path = current_dir / "README.md" # This will overwrite your README.md in the build context
    download_file(readme_url, readme_local_path)


if __name__ == "__main__":
    setup_imagebind()