import requests
import json
import posixpath
import os

from modules.posenet import MOBILENET_V1_CHECKPOINTS

GOOGLE_CLOUD_STORAGE_DIR = "https://storage.googleapis.com/tfjs-models/weights/posenet/"


def download_json(checkpoint, filename, base_dir):
    url = posixpath.join(GOOGLE_CLOUD_STORAGE_DIR, checkpoint, filename)
    response = requests.get(url)
    data = json.loads(response.content)

    with open(os.path.join(base_dir, checkpoint, filename), "w") as outfile:
        json.dump(data, outfile)


def download_file(checkpoint, filename, base_dir):
    url = posixpath.join(GOOGLE_CLOUD_STORAGE_DIR, checkpoint, filename)
    response = requests.get(url)
    f = open(os.path.join(base_dir, checkpoint, filename), "wb")
    f.write(response.content)
    f.close()


def download(checkpoint, base_dir="./weights/"):
    save_dir = os.path.join(base_dir, checkpoint)
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)

    download_json(checkpoint, "manifest.json", base_dir)

    f = open(os.path.join(save_dir, "manifest.json"), "r")
    json_dict = json.load(f)

    for x in json_dict:
        filename = json_dict[x]["filename"]
        print("Downloading", filename)
        download_file(checkpoint, filename, base_dir)


def main():
    checkpoint = MOBILENET_V1_CHECKPOINTS[101]
    download(checkpoint)


if __name__ == "__main__":
    main()