# Copyright 2021 The IREE Authors
#
# Licensed under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

# This tool handles mirroring tflite testing files from their source to the
# the iree-model-artifacts test bucket. This avoids taking dependency on
# external test files that may change or no longer be available.
#
# To update all files:
#   python update_tflite_models.py --file all
#
# To update a specific file:
#   python update_tflite_models.py --file posenet_i8_input.jpg
#
# Note you must have write permission to the iree-model-artifacts GCS bucket
# with local gcloud authentication.

from absl import app
from absl import flags
from google.cloud import storage
from google_auth_oauthlib import flow

import tempfile
import urllib

FLAGS = flags.FLAGS
flags.DEFINE_string('file', '', 'file to update')

file_dict = dict({
    "mobilenet_v1.tflite":
        "https://tfhub.dev/tensorflow/lite-model/mobilenet_v1_1.0_160/1/default/1?lite-format=tflite",
    "posenet_i8.tflite":
        "https://tfhub.dev/google/lite-model/movenet/singlepose/lightning/tflite/int8/4?lite-format=tflite",
    "posenet_i8_input.jpg":
        "https://github.com/tensorflow/examples/raw/master/lite/examples/pose_estimation/raspberry_pi/test_data/image3.jpeg"
})

BUCKET_NAME = "iree-model-artifacts"
FOLDER_NAME = "tflite-integration-tests"


def upload_model(source, destination, tmpfile):
  """Uploads a file to the bucket."""
  urllib.request.urlretrieve(source, tmpfile)

  storage_client = storage.Client()
  bucket = storage_client.get_bucket(BUCKET_NAME)
  blob = bucket.blob("/".join([FOLDER_NAME, destination]))
  blob.upload_from_filename(tmpfile)


def main(argv):
  tf = tempfile.NamedTemporaryFile()

  items = file_dict.items()

  if FLAGS.file in file_dict:
    items = [(FLAGS.file, file_dict[FLAGS.file])]
  elif FLAGS.file != "all":
    print('Unknown file to upload: ', "\"" + FLAGS.file + "\"")
    exit()

  for dst, src in items:
    upload_model(src, dst, tf.name)


if __name__ == '__main__':
  app.run(main)