import argparse

import requests

WORKFLOW_RUN_ENDPOINT = (
    "https://api.github.com/repos/{owner}/{repo}/actions/runs/{run_id}/artifacts"
)
ARTIFACT_DOWNLOAD_ENDPOINT = (
    "https://api.github.com/repos/{owner}/{repo}/actions/artifacts/{artifact_id}/zip"
)


def download_artifact(
    owner: str, repo: str, run_id: str, artifact_name: str, gh_token: str
):
    """Download artifact by name from a run id
    Args:
        owner: Owner of the repo.
        repo: The name of the repository.
        run_id: The id of the action run that created the artifact.
        artifact_name: The name of the artifact to download.
        gh_token: The token used to download artifacts
    """

    headers = {
        "Accept": "application/vnd.github+json",
        "Authorization": f"Bearer {gh_token}",
        "X-GitHub-Api-Version": "2022-11-28",
    }

    artifacts = requests.get(
        WORKFLOW_RUN_ENDPOINT.format(owner=owner, repo=repo, run_id=run_id),
        headers=headers,
    ).json()
    artifact_id = next(
        (
            artifact["id"]
            for artifact in artifacts["artifacts"]
            if artifact["name"] == artifact_name
        ),
        None,
    )

    if not artifact_id:
        raise ValueError(
            f"Cannot find {artifact_name}! {artifacts} {owner} {repo} {run_id}"
        )

    download = requests.get(
        ARTIFACT_DOWNLOAD_ENDPOINT.format(
            artifact_id=artifact_id,
            owner=owner,
            repo=repo,
        ),
        headers=headers,
    )
    with open(f"{artifact_name}.zip", "wb") as f:
        f.write(download.content)

    return f"{artifact_name}.zip"


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Download an artifact")
    parser.add_argument("run_id", type=str, help="Run id")
    parser.add_argument("artifact_name", type=str)
    parser.add_argument("gh_token", type=str)
    parser.add_argument(
        "--owner", type=str, help="Owner of repository", default="gradio-app"
    )
    parser.add_argument("--repo", type=str, help="Name of repository", default="gradio")
    args = parser.parse_args()
    new_space = download_artifact(
        args.owner, args.repo, args.run_id, args.artifact_name, args.gh_token
    )
    print(new_space)
