#!/usr/bin/env python

import argparse
import json
import os
import sys
import tarfile
from urllib.request import Request, urlopen
from urllib.parse import urlparse


def url_download(url, work_dir):
    url_parts = urlparse(url)
    file_path = os.path.abspath(os.path.join(work_dir, os.path.basename(url_parts.path)))
    src = None
    dst = None
    try:
        req = Request(url)
        src = urlopen(req)
        with open(file_path, 'wb') as dst:
            while True:
                chunk = src.read(2**10)
                if chunk:
                    dst.write(chunk)
                else:
                    break
    except Exception as e:
        sys.exit(str(e))
    finally:
        if src:
            src.close()
    if tarfile.is_tarfile(file_path):
        fh = tarfile.open(file_path, 'r:*')
    else:
        return file_path
    fh.extractall(work_dir)
    os.remove(file_path)
    return work_dir


def download(database_id, database_name, url, out_file):

    with open(out_file) as fh:
        params = json.load(fh)

    work_dir = params['output_data'][0]['extra_files_path']
    os.makedirs(work_dir)
    file_path = url_download(url, work_dir)

    data_manager_json = {"data_tables": {}}
    data_manager_entry = {}
    data_manager_entry['value'] = database_id
    data_manager_entry['name'] = database_name
    data_manager_entry['path'] = file_path
    data_manager_json["data_tables"]["gtdbtk_database"] = data_manager_entry

    with open(out_file, 'w') as fh:
        json.dump(data_manager_json, fh, sort_keys=True)


parser = argparse.ArgumentParser()

parser.add_argument('--database_name', dest='database_name', help='GTDB-Tk database display name')
parser.add_argument('--database_id', dest='database_id', help='Unique GTDB-Tk database id')
parser.add_argument('--url', dest='url', help='URL to download GTDB-Tk databse version')
parser.add_argument('--out_file', dest='out_file', help='JSON output file')

args = parser.parse_args()

download(args.database_id, args.database_name, args.url, args.out_file)