#!/usr/bin/python3
import requests
import json
import argparse
import os
import time

SERVER_URL = 'https://editor.beey.io' # or your on-premise installation
API_TOKEN = 'your_api_token'

def create_project(name, token):
    url = f'{SERVER_URL}/API/Project'
    payload = {
        'Name': name,
        'CustomPath': ''
    }
    payload = json.dumps(payload)
    headers = {
        'Authorization': API_TOKEN,
        'Content-Type': 'application/json'
    }
    response = requests.request('POST', url, headers=headers, data=payload)
    if response.status_code != 200:
        print('creating project failed' + str(response.status_code) + ' ' + response.text)
        raise Exception('create project failed')
    return response.json()['Id'], response.json()['AccessToken']

def upload_file(fname, project_id):
    file_size = os.path.getsize(fname)
    url = f'{SERVER_URL}/API/Project/{project_id}/Files/UploadMediaFile?fileSize={file_size}'

    payload = {}
    files = [
        ('name', open(fname, 'rb'))
    ]
    headers = {
        'Authorization': API_TOKEN
    }

    response = requests.request('POST', url, headers=headers, data=payload, files=files)
    if response.status_code != 200:
        print('file upload failed' + str(response.status_code) + ' ' + response.text)
        raise Exception('file upload failed')

def transcribe_file(project_id, language):
    url = f'{SERVER_URL}/API/Project/Queue/Enqueue?projectId={project_id}&lang={language}&transcriptionProfile=default'
    print(url)
    payload = {}
    files = {}
    headers = {
      'Authorization': API_TOKEN
    }
    response = requests.request('GET', url, headers=headers, data=payload, files=files)
    if response.status_code != 200:
        # something unexpected happened
        print('enqueue failed ' + str(response.status_code) + ' ' + response.text)
        raise Exception('failed to transcribe')

# acquires trsx (xml output of the ASR including transcription with timestamps)
def download_original_trsx(project_id):
    url = f'{SERVER_URL}/API/Project/{project_id}/Files/OriginalTrsx'
    payload  = {}
    headers = {
        'Authorization': API_TOKEN,
        'Content-Type': 'application/json' 
    }

    response = requests.request('GET', url, headers=headers, data = payload)
    print(response)
    return response.text

def get_subtitles(project_id):
    url = f'{SERVER_URL}/API/Project/{project_id}/Export?formatId=srt'

    payload = {}
    headers = {
      'Authorization': API_TOKEN,
      'Content-Type': 'application/json; charset=utf-8'
    }

    response = requests.request('GET', url, headers=headers, data = payload)
    response.encoding = 'utf-8'
    print(response)
    return response.text


# Wait until transcription is complete
def wait_for_trsx(project_id):
    print('poll for trsx:')
    url = f'{SERVER_URL}/API/Project/{project_id}'

    payload = {}
    headers = {
        'Authorization': API_TOKEN,
        'Content-Type': 'application/json'
    }
    originalTrsxId = None
    for i in range(240): # give up after 2 hours
        response = requests.request('GET', url, headers=headers, data = payload)
        originalTrsxId = response.json()['OriginalTrsxId']
        if originalTrsxId != None:
            print('trsx available - transcription ended')
            return True
        print('trsx not yet available, retrying in 30 s...')
        time.sleep(30)
    print('transcription timed out')
    return False

def save_file(content, name):
    with open(name, 'w', encoding='utf-8') as output_file:
        output_file.write(content)


if (__name__ == '__main__'):
    parser = argparse.ArgumentParser()
    parser.add_argument('recording')
    parser.add_argument('output')
    parser.add_argument('--language', default='cs-CZ')
    parser.add_argument('--subtitles', action='store_true')
    args = parser.parse_args()

    print('creating project...')
    project_id, access_token = create_project('api ' + args.recording)

    print('uploading file...')
    upload_file(args.recording, project_id)

    print('requesting transcription...')
    transcribe_file(project_id, args.language)

    print('waiting for trsx...')
    wait_for_trsx(project_id)
    
    print('saving...')
    if (args.subtitles):
      subtitles = get_subtitles(project_id)
      save_file(subtitles, args.output)
      print('subtitles saved to: ' + args.output)
    else:
      trsx = download_original_trsx(project_id)
      save_file(trsx, args.output)
      print('trsx saved to: ' + args.output)
    print('done.')