#!/usr/bin/env python3
#
# Copyright 2020 New Vector Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import argparse
import hashlib
import json
import os
# Run `pip3 install requests` if not installed yet
import requests

# This script downloads artifacts from buildkite.
# Ref: https://buildkite.com/docs/apis/rest-api/artifacts#download-an-artifact

# Those two variable are specific to the RiotX project
ORG_SLUG = "matrix-dot-org"
PIPELINE_SLUG = "riotx-android"

### Arguments

parser = argparse.ArgumentParser(description='Download artifacts from Buildkite.')
parser.add_argument('-t',
                    '--token',
                    required=True,
                    help='The buildkite token.')
parser.add_argument('-b',
                    '--build',
                    type=int,
                    required=True,
                    help='the buildkite build number.')
parser.add_argument('-e',
                    '--expecting',
                    type=int,
                    default=-1,
                    help='the expected number of artifacts. If omitted, no check will be done.')
parser.add_argument('-d',
                    '--directory',
                    default="",
                    help='the target directory, where files will be downloaded. If not provided the build number will be used to create a directory.')
parser.add_argument('-v',
                    '--verbose',
                    help="increase output verbosity.",
                    action="store_true")
parser.add_argument('-s',
                    '--simulate',
                    help="simulate action, do not create folder or download any file.",
                    action="store_true")

args = parser.parse_args()

# parser has checked that the build was an int, convert to String for the rest of the script
build_str = str(args.build)

if args.verbose:
    print("Argument:")
    print(args)

headers = {'Authorization': "Bearer %s" % args.token}
base_url = "https://api.buildkite.com/v2/organizations/%s/pipelines/%s/builds/%s" % (ORG_SLUG, PIPELINE_SLUG, build_str)

### Fetch build state

buildkite_build_state_url = base_url

print("Getting build state of project %s/%s build %s" % (ORG_SLUG, PIPELINE_SLUG, build_str))

if args.verbose:
    print("Url: %s" % buildkite_build_state_url)

r0 = requests.get(buildkite_build_state_url, headers=headers)
data0 = json.loads(r0.content.decode())

if args.verbose:
    print("Json data:")
    print(data0)

print("   git branch         : %s" % data0.get('branch'))
print("   git commit         : \"%s\"" % data0.get('commit'))
print("   git commit message : \"%s\"" % data0.get('message'))
print("   build state        : %s" % data0.get('state'))

if data0.get('state') != 'passed':
    print("❌ Error, the build failed (state: %s)" % data0.get('state'))
    exit(0)

### Fetch artifacts list

buildkite_artifacts_url = base_url + "/artifacts"

print("Getting artifacts list of project %s/%s build %s" % (ORG_SLUG, PIPELINE_SLUG, build_str))

if args.verbose:
    print("Url: %s" % buildkite_artifacts_url)

r = requests.get(buildkite_artifacts_url, headers=headers)
data = json.loads(r.content.decode())

print("   %d artifact(s) found." % len(data))

if args.expecting != -1 and args.expecting != len(data):
    print("Error, expecting %d artifacts and found %d." % (args.expecting, len(data)))
    exit(1)

if args.verbose:
    print("Json data:")
    print(data)

if args.verbose:
    print("Create subfolder %s to download artifacts..." % build_str)

if args.directory == "":
    targetDir = build_str
else:
    targetDir = args.directory

if not args.simulate:
    os.mkdir(targetDir)

error = False

for elt in data:
    if args.verbose:
        print()
        print("Artifact info:")
        for key, value in elt.items():
            print("   %s: %s" % (key, str(value)))
    url = elt.get("download_url")
    filename = elt.get("filename")
    target = targetDir + "/" + filename
    print("Downloading %s to '%s'..." % (filename, targetDir))
    if not args.simulate:
        # open file to write in binary mode
        with open(target, "wb") as file:
            # get request
            response = requests.get(url, headers=headers)
            # write to file
            file.write(response.content)
        print("Verifying checksum...")
        # open file to read in binary mode
        with open(target, "rb") as file:
            data = file.read()
            hash = hashlib.sha1(data).hexdigest()
        if elt.get("sha1sum") != hash:
            error = True
            print("❌ Checksum mismatch: expecting %s and get %s" % (elt.get("sha1sum"), hash))

if error:
    print("❌ Error(s) occurred, check the log")
    exit(1)
else:
    print("Done!")