#!/usr/bin/env python
#
# fetch a buildkite artifact
#

from __future__ import print_function

import argparse
import os
import requests

parser = argparse.ArgumentParser(
    "fetch an artifact from buildkite"
)

parser.add_argument(
    "--access-token",
    help=(
        "Buildkite access token. If unset, will be read from the "
        "BUILDKITE_ACCESS_TOKEN environment variable."
    ),
)

parser.add_argument(
    "--build-filter",
    default="branch=master&state=passed",
    help="Query string for the build filter. Default: %(default)s",
)

parser.add_argument("org_slug")
parser.add_argument("build_slug")
parser.add_argument("artifact_name")

args=parser.parse_args()

access_token=args.access_token
if access_token is None:
    access_token = os.environ["BUILDKITE_ACCESS_TOKEN"]
artifact_name = args.artifact_name

session = requests.Session()
session.headers["Authorization"] = "Bearer " + access_token

# we start by fetching the REST url for the latest matching build
print("Fetching build list")
r = session.get(
    "https://api.buildkite.com/v2/organizations/%s/pipelines/%s/builds?%s" % (
        args.org_slug, args.build_slug, args.build_filter,
    ),
)
r.raise_for_status()
build_url = r.json()[0]["url"]

# then we need to get the download_url for the artifact
print("Fetching artifact list")
r = session.get(build_url + "/artifacts")
r.raise_for_status()
artifact = next(a for a in r.json() if a["filename"] == artifact_name)
download_url = artifact["download_url"]

# the download_url returns a 302 to the S3 URL.
print("Downloading artifact %s -> %s" % (download_url, artifact_name))
r = session.get(download_url)
r.raise_for_status()
with open(artifact_name, "wb") as f:
    f.write(r.content)
    print("done")
