#!/bin/python3

import argparse
import requests
import json
import time
from datetime import datetime


API_URL = 'https://api.deepva.com/api/v1/jobs'

argparser = argparse.ArgumentParser(description='Test AI Image Tagging | DeepVA')


argparser.add_argument(
    '-cf',
    '--client-file',
    default='',
    help='File with API Key',
)
argparser.add_argument(
    '-ak',
    '--api-key',
    default='',
    help='API Key (if set, ignore --client-file)',
)
argparser.add_argument(
    '-i',
    '--image-url',
    help='Image URL',
)
argparser.add_argument(
    '-api',
    '--api-url',
    default=API_URL,
    help='API URL (default: {})'.format(API_URL),
)
mm = 'object_scene_recognition:general-c,landmark_recognition:europe'
argparser.add_argument(
    '-m',
    '--models',
    default=mm,
    help='Pairs of Module:Model (default: mm)',
)


def parse_modules_models(modules_models: str) -> dict[str]:
    mm = {}
    if not modules_models:
        return mm

    for pair in modules_models.split(','):
        kv = pair.split(':')
        if len(kv) != 2:
            continue
        mm[kv[0].strip()] = {
            'model': kv[1].strip(),
        }

    return mm


def header(api_key: str) -> dict[str]:
    return {
        'Content-Type': 'application/json',
        'Authorization': f'Key {api_key}',
    }


def post_image_get_job(
    api_url: str, api_key: str, image_url: str, modules_models: str
) -> requests.Response:
    return requests.post(
        api_url,
        json={
            'sources': [
                image_url,
            ],
            'modules': parse_modules_models(modules_models),
        },
        headers=header(api_key),
    )


def get_job(api_url: str, api_key: str, job_id: str) -> requests.Response:
    return requests.get(
        f'{api_url}/{job_id}',
        headers=header(api_key),
    )


def parse_job_result(response: requests.Response) -> tuple[dict[str], str, bool]:
    status_ok = [200, 202]
    if response.status_code not in status_ok or False:
        print(
            f'Request failed, expected status codes {status_ok}: {response.status_code}: {response.text}'
        )
        exit(1)

    result = json.loads(response.text)
    status = result.get('state')

    job_id = result.get('id')
    if not job_id or job_id == '':
        print(f'Could not get a Job ID')
        exit(1)

    progress = result.get('progress')
    if not isinstance(progress, float):
        progress = 0.0

    print(f'Progress: {int(100.0*progress)}%\tStatus: {status}')

    if status == 'completed':
        parsed_result = result.get('result')
        if not parsed_result:
            print(f'Could not parse result from completed response')
            exit(1)

        # done: return result
        return parsed_result, job_id, True

    if status in ['stopped', 'failed']:
        print(f'Job could not be completed: {status}')
        exit(1)

    # not done: return job id
    return None, job_id, False


if __name__ == '__main__':
    args = argparser.parse_args()

    api_url = args.api_url.strip()
    if api_url == '':
        print('No valid API Url!')
        exit(1)

    api_key = args.api_key
    if args.client_file != '':
        with open(args.client_file, 'r') as f:
            content = f.read()
            parts = content.split('\n')
            for p in parts:
                p = p.strip()
                if len(p) > 0:
                    api_key = p
                    break
    if api_key == '':
        print('No valid API Key!')
        exit(1)

    image_url = args.image_url
    if api_key == '':
        print('No valid Image Url!')
        exit(1)
    print(f'Image URL: {image_url}')

    # first step: upload image, get job

    start = datetime.now()
    response = post_image_get_job(
        api_url,
        api_key,
        image_url,
        args.models,
    )
    took = datetime.now() - start
    print('Request took', took)

    result, job_id, done = parse_job_result(response)
    if done:
        print(f'Response (Result): {json.dumps(result,indent=4)}')
        exit(0)

    print(f'Job ID: "{job_id}"')

    # second step: poll until job is completed (or failed)

    delay = 3
    while True:

        start = datetime.now()
        response = get_job(
            api_url,
            api_key,
            job_id,
        )
        took = datetime.now() - start
        print('Request took', took)

        result, job_id, done = parse_job_result(response)
        if done:
            print(f'Response (Result):\n{json.dumps(result,indent=4)}')
            break

        print(f'Wait for {delay} second(s) until job "{job_id}" is checked again')
        time.sleep(delay)
