#!/usr/bin/env python3


import os
import sys
import requests
import time
import io
import zipfile
import json


AUTH_FORM = {
    'grant_type': 'client_credentials',
    'client_id': 'XXX',
    'client_secret': 'XXX',
}


AUTH_FILE = 'oauth.json'
PLAYER_FILE = 'player.json'


def get_auth_header():
    now = time.time()
    access_token = None
    token_type = None

    if os.path.isfile(AUTH_FILE):
        print('Checking previously saved token validity')

        with open(AUTH_FILE) as ifd:
            auth = json.load(ifd)

        expired = now > auth['expires']

        if not expired:
            access_token = auth['access_token']
            token_type = auth['token_type']

    request_new = not (access_token and token_type)

    if request_new:
        print('Getting a token...')
        rsp = requests.post(
            'https://api.avatarsdk.com/o/token/',
            data=AUTH_FORM
        ).json()

        access_token = rsp['access_token']
        token_type = rsp['token_type']

        print('Saving token for future use')
        with open(AUTH_FILE, 'w') as ofd:
            # reserve 1 minute to avoid expiration in the middle of
            # avatar retrieval. Works only for relatively fast
            # connections
            rsp['expires'] = now + rsp['expires_in'] - 60
            json.dump(rsp, ofd)

    headers = {'Authorization': '{0} {1}'.format(token_type, access_token)}

    return headers


def get_player_uid_header(headers):
    player_uid = None

    if os.path.isfile(PLAYER_FILE):
        print('Loading previously saved PlayerUID')

        with open(PLAYER_FILE) as ifd:
            player = json.load(ifd)

        player_uid = player['code']

    if not player_uid:
        print('Creating a player...')

        player_form = {'comment': 'test_py'}
        rsp = requests.post(
            'https://api.avatarsdk.com/players/',
            data=player_form, headers=headers
        ).json()

        player_id = rsp['code']

        print('Saving PlayerUID future use')
        with open(PLAYER_FILE, 'w') as ofd:
            json.dump(rsp, ofd)

    return {'X-PlayerUID': player_uid}


def main(selfie):
    headers = get_auth_header()

    headers.update(
        get_player_uid_header(headers)
    )

    pipeline_type = 'metaperson_2.0'
    pipeline_subtype = 'female'

    parameters = {
        "avatar_modifications": {
            "remove_smile": True,
            "remove_glasses": True,
        },
        "model_info": ["age", "gender", "race"]
    }

    export_parameters = {
        "format": "glb",
        "lod": "LOD1",
        "embed": True,
        "finalize": True,

        "textures": {
            "list": [
                "Color",
                "Normal",
                "GltfMetallicRoughness",
                "BodyVisibilityMask",
                "HeadVisibilityMask"
            ],
            "embed": True,
            "profile": "1K.png"
        },

        "avatar": {
            "list": [
                "AvatarBody",
                "AvatarEyelashes",
                "AvatarHead",
                "AvatarLeftCornea",
                "AvatarLeftEyeball",
                "AvatarRightCornea",
                "AvatarRightEyeball",
                "AvatarTeethLower",
                "AvatarTeethUpper"
            ]
        },

        "blendshapes": {
            "list": ["face_modifications", "body_shape", "mobile_51"]
        },

        "haircuts": {
            "list": ["HaircutGenerated"]
        },

        "outfits": {
            "list": ["LORI"]
        }
    }

    print('Uploading image...')
    data = {
        'name': 'test',
        'pipeline': pipeline_type,
        'pipeline_subtype': pipeline_subtype,
        'parameters': json.dumps(parameters),
        'export_parameters': json.dumps(export_parameters),
    }
    files = {'photo': open(selfie, 'rb')}
    rsp = requests.post(
        'https://api.avatarsdk.com/avatars/',
        headers=headers, data=data, files=files
    ).json()

    avatar_status_url = rsp['url']

    print('Waiting for avatar to compute...', end='', flush=True)
    while True:
        rsp = requests.get(avatar_status_url, headers=headers).json()

        print('.', end='', flush=True)

        if rsp['status'] == 'Completed':
            print('Completed!')
            break

        time.sleep(3)

    print('Retrieve avatar exports list...')

    avatar_exports_url = rsp['exports']
    exports = requests.get(avatar_exports_url, headers=headers).json()

    # given that we requested only one finalized export, use the first file url
    # from the first avatar export
    mesh_url = exports[0]['files'][0]['file']

    print('Downloading avatar...')
    mesh = requests.get(mesh_url, headers=headers)

    print('Saving avatar model to the current directory...')
    with io.BytesIO(mesh.content) as zipmemory:
        with zipfile.ZipFile(zipmemory) as archive:
            archive.extractall()

    print('Ready!')


if __name__ == '__main__':
    if AUTH_FORM['client_id'] == 'XXX' or AUTH_FORM['client_secret'] == 'XXX':
        print((
            'Before running this script, please replace \'XXX\' in '
            'lines 15 and 16 with your `client_id` and '
            '`client_secret` correspondingly. You can get them by '
            'logging to https://accounts.avatarsdk.com/developer/ '
            'with your credentials. Do not forget to set Authorization Grant '
            'to `Client credentials`!'
        ))
        sys.exit(1)

    if len(sys.argv) < 2:
        print('Specify selfie file')
        sys.exit(1)

    main(sys.argv[1])