from huggingface_hub import HfApi, RepoUrl


async def launch_substra_space(
        hf_api: HfApi, repo_id: str,
        hospital_a: int, hospital_b: int,
) -> RepoUrl:
    repo_id = "owkin/trainer-" + repo_id

    BASE_IMAGE = "ghcr.io/substra/substrafl:substra-hf-demo"

    DOCKERFILE = f"""
    FROM {BASE_IMAGE}
    
    ENV SUBSTRA_ORG1_DISTR={hospital_a / 100}
    ENV SUBSTRA_ORG2_DISTR={hospital_b / 100} 
    
    CMD ["bash", "docker-run.sh"]
    """
        
    repo_url = hf_api.create_repo(
        repo_id=repo_id,
        repo_type="space",
        space_sdk="docker",
    )

    hf_api.upload_file(
        repo_id=repo_id,
        path_or_fileobj=DOCKERFILE.encode(),
        path_in_repo="Dockerfile",
        repo_type="space",
    )
    hf_api.upload_file(
        repo_id=repo_id,
        path_or_fileobj="./substra_template/README.md",
        path_in_repo="README.md",
        repo_type="space",
    )
    return repo_url