#! /usr/bin/env python3
import datetime
import json
import os
import socket
import subprocess
import time
import urllib.request
import xml.etree.ElementTree as ET
from http.client import RemoteDisconnected
from io import StringIO
from os.path import join as pjoin
from xml.dom import minidom

import docker
import yaml
from docker import errors
from docker.models.containers import Container

ARCHIVE_PATH = 'archive_path'
APPS = 'apps'
APP_PREFIX = 'app-'
CONSUMERS = 'consumers'
CONSUMER_PREFIX = 'consumer-'
# local directory to clone clams-galaxy on the host
GALAXY_LOCAL_PATH = 'clams-galaxy'
# hostname given to the container instantiated from the cloned clams-galaxy
GALAXY_CONTNAME = 'clams-galaxy'
# fixed by galaxy-stable image
GALAXY_CONTPORT = '80'
PRIMARY_HOSTPORT = '8080'
DOCKER_NETWORK_NAME = 'clams-appliance'
CONTAINER_DATA_PATH = '/data'
HOSTNAME = socket.gethostname()
DEVELOP = False
GALAXY_CATEGORIES = {
    'AudioDocument': 'Audio',
    'VideoDocument': 'Video',
    'TextDocument': 'Text',
    'Alignment': 'Alignment',
}

docker_engine = docker.from_env()


#####
# for xml processing
def CDATA(text=None):
    element = ET.Element('![CDATA[')
    element.text = text
    return element

ET._original_serialize_xml = ET._serialize_xml


def _serialize_xml(write, elem, qnames, namespaces, short_empty_elements, **kwargs):
    if elem.tag == '![CDATA[':
        write("\n<{}{}]]>\n".format(elem.tag, elem.text))
        if elem.tail:
            write(_escape_cdata(elem.tail))
    else:
        return ET._original_serialize_xml(write, elem, qnames, namespaces,short_empty_elements, **kwargs)

ET._serialize_xml = ET._serialize['xml'] = _serialize_xml
# for xml processing
#####


def get_docker_image_name(app_name):
    return f"clams-{app_name}"


def clean(directory):
    import shutil
    for f in os.listdir(directory):
        if f.startswith(APP_PREFIX) or f.startswith(CONSUMER_PREFIX) or f == GALAXY_LOCAL_PATH:
            d = pjoin('.', f)
            if os.path.islink(d):
                os.unlink(d)
            elif os.path.isdir(d):
                shutil.rmtree(d)
    try:
        os.remove('docker-compose.yml', )
    except OSError:
        pass


def read_config(config_file_path):
    configs = yaml.load(open(config_file_path).read())
    configs[ARCHIVE_PATH] = os.path.expandvars(os.path.expanduser(configs[ARCHIVE_PATH]))
    return configs


def download_galaxy_mods():
    if not os.path.exists(GALAXY_LOCAL_PATH):
        download('https://github.com/clamsproject/clams-galaxy.git', GALAXY_LOCAL_PATH, [])


def create_docker_compose(config, export_volumename, rebuild=False, develop=False):
    global DEVELOP
    DEVELOP = develop
    if rebuild:
        clean('.')
    prefixed_app_names = list(map(
        lambda x: APP_PREFIX + x,
        [appname for appname in config[APPS].keys() if config[APPS][appname]['enabled']]
    ))

    docker_compose = prep_galaxy(export_volumename, prefixed_app_names, config[ARCHIVE_PATH])
    process_all_apps(config[APPS], docker_compose, config[ARCHIVE_PATH])
    process_all_consumers(config[CONSUMERS], docker_compose, config[ARCHIVE_PATH])
    gen_db_loc_files(config[ARCHIVE_PATH])
    docker_engine.images.build(path=GALAXY_LOCAL_PATH, dockerfile='Dockerfile', tag=get_docker_image_name(GALAXY_LOCAL_PATH), nocache=True)
    with open('docker-compose.yml', 'w') as compose_file:
        yaml.SafeDumper.add_representer(
            type(None),
            lambda dumper, value: dumper.represent_scalar(u'tag:yaml.org,2002:null', '')
        )
        print(yaml.safe_dump(docker_compose))
        compose_file.write((yaml.dump(docker_compose)))


def create_base_compose_obj():
    return {'version': '3', 'services': {}, 'networks': {DOCKER_NETWORK_NAME: None}, 'volumes': {}}


def prep_galaxy(export_volumename, dependencies, host_data_path):
    download_galaxy_mods()
    compose_obj = create_base_compose_obj()
    galaxy_service = get_service_def(GALAXY_CONTNAME, int(PRIMARY_HOSTPORT))
    # even though we configure `http` address to be bound to 0.0.0.0:5000 in galaxy.yml, 
    # external communication still needs to be using port 80 for underlying wsgi modules to work
    galaxy_service[GALAXY_CONTNAME].update({'privileged': True, 'depends_on': dependencies, 'ports': [f'{PRIMARY_HOSTPORT}:{GALAXY_CONTPORT}']})
    compose_obj['services'].update(galaxy_service)
    add_data_volume(GALAXY_CONTNAME, compose_obj, host_data_path)
    add_galaxy_export_volume(compose_obj, GALAXY_CONTNAME, export_volumename)
    return compose_obj


def get_service_def(cont_hostname, port):
    service_def = {cont_hostname: {
        'image': get_docker_image_name(cont_hostname),
        'container_name': cont_hostname,
        'networks': [DOCKER_NETWORK_NAME],
    }}
    if port != 5000:
        service_def[cont_hostname]['ports'] = [f'{port}:5000']
    return service_def


def process_all_apps(apps_config, docker_compose_obj, host_data_path):
    tool_conf_path = pjoin(GALAXY_LOCAL_PATH, 'config', 'tool_conf.xml')
    tool_conf_tree = ET.parse(tool_conf_path)
    cont_port = 5000
    for host_port, (app_name, app_config) in enumerate(apps_config.items(), 8001):
        app_name = f'{APP_PREFIX}{app_name}'
        if app_config['enabled']:
            download_app(app_name, app_config)
            build_docker_image(app_name)
            add_to_docker_compose(app_name, docker_compose_obj, host_port)
            add_data_volume(app_name, docker_compose_obj, host_data_path)
            config_xml_tree = gen_app_config_xml(app_name, app_config, cont_port)
            add_to_tool_conf_xml(tool_conf_tree, config_xml_tree, app_name)
    tool_conf_tree.write(tool_conf_path, encoding='utf-8', xml_declaration=True)


def get_tool_config_xml_fullpath(app_name):
    return pjoin(GALAXY_LOCAL_PATH, 'tools', get_tool_config_xml_filename(app_name))


def get_tool_config_xml_filename(app_name):
    return get_docker_image_name(app_name) + '.xml'


def process_all_consumers(consumers_config, docker_compose_obj, host_data_path):
    datatypes_conf_path = pjoin(GALAXY_LOCAL_PATH, 'config', 'datatypes_conf.xml')
    datatypes_conf_tree = ET.parse(datatypes_conf_path)
    if not consumers_config:
        consumers_config = {}
    for port, (consumer_name, consumer_config) in enumerate(consumers_config.items(), 9001):
        consumer_name = f'{CONSUMER_PREFIX}{consumer_name}'
        if consumer_config['enabled']:
            download_app(consumer_name, consumer_config)
            build_docker_image(consumer_name)
            add_to_docker_compose(consumer_name, docker_compose_obj, port)
            add_data_volume(consumer_name, docker_compose_obj, host_data_path, flask_static=True)
            gen_display_app_xml(consumer_name, port, consumer_config['description'])
            add_to_datatypes_conf_xml(datatypes_conf_tree, consumer_name)
    datatypes_conf_tree.write(datatypes_conf_path, encoding='utf-8')


def get_display_app_xml_fullpath(consumer_name):
    return pjoin(GALAXY_LOCAL_PATH, 'display_applications', get_display_app_xml_filename(consumer_name))


def get_display_app_xml_filename(consumer_name):
    return get_docker_image_name(consumer_name) + '.xml'


def add_to_docker_compose(cont_hostname, docker_compose_obj, port):
    docker_compose_obj['services'].update(get_service_def(cont_hostname, port))


def add_data_volume(cont_hostname, docker_compose_obj, host_data_path, flask_static=False):
    # 'ro' for read-only
    docker_compose_obj['services'][cont_hostname].update({'volumes': [f'{host_data_path}:{CONTAINER_DATA_PATH}:ro']})
    if flask_static: 
        docker_compose_obj['services'][cont_hostname]['volumes'].append(f'{host_data_path}:/app/static{CONTAINER_DATA_PATH}:ro')


def add_galaxy_export_volume(docker_compose_obj, cont_hostname, volumename):
    try:
        docker_engine.volumes.get(volumename)
    except docker.errors.NotFound:
        docker_engine.volumes.create(volumename)
    docker_compose_obj['services'][cont_hostname].update({'volumes': [f'{volumename}:/export']})
    docker_compose_obj['volumes'][volumename] = {'external': True}


def build_docker_image(dir_name):
    print(f"Building a Docker image from {dir_name}...")
    image = docker_engine.images.build(path=dir_name, dockerfile='Dockerfile', tag=get_docker_image_name(dir_name), nocache=False)[0]
    print(f"Built: {image.id}")


def gen_app_config_xml(app_name, app_config, port):
    image_name = get_docker_image_name(app_name)
    if os.path.exists(pjoin(app_name, 'config.xml')):
        config_xml_tree = ET.parse(pjoin(app_name, 'config.xml'))
    elif docker_engine.images.list(name=image_name):
        image_name = get_docker_image_name(app_name)
        container: Container = docker_engine.containers.run(image=image_name, remove=True, detach=True,
                                                            ports={port: PRIMARY_HOSTPORT})
        attempts = 0
        appmetadata = None
        try:
            while attempts < 20 and appmetadata is None:
                try:
                    time.sleep(1)  # give a second for server to start update 
                    appmetadata = urllib.request.urlopen(f'http://127.0.0.1:{PRIMARY_HOSTPORT}').read()
                    config_xml_tree = appmetadata_to_config_xml_tree(appmetadata)
                    break
                    # generate xmltree
                except (RemoteDisconnected, ConnectionResetError):
                    attempts += 1
            else:
                # TODO (krim @ 7/13/21): should I just use app_config and continue? 
                raise ConnectionError(f"app \"{app_name}\" is not responding, and galaxy config.xml cannot be generated")
        finally:
            container.stop()
    else:
        config_xml_tree = configyml_to_config_xml_tree(app_name, app_config)

    command_tag = config_xml_tree.find('command')
    params_collector = '#set $params = "".join([f"{k}={v}&" for k, v in $runtime_params.items()]) if $runtime_params else ""\n'
    curl_cmd = CDATA(f"{params_collector}curl -s -X POST -H 'Content-Type: application/json' -d @$input \"{app_name}:{port}?$params\" > $output")
    command_tag.append(curl_cmd)
    try: 
        del command_tag.attrib['interpreter']
    except KeyError as ignored:
        pass
    config_xml_string = ET.tostring(config_xml_tree.getroot(), 'utf8')
    reparsed = minidom.parseString(config_xml_string)
    with open(get_tool_config_xml_fullpath(app_name), 'w', encoding='utf8') as config_xml_file:
        config_xml_file.write(reparsed.toprettyxml(indent="\t"))
    return config_xml_tree


def configyml_to_config_xml_tree(app_name, app_config):
    # create from scratch
    tool_tag = ET.Element('tool')
    tool_tag.set('id', app_name)
    tool_tag.set('name', app_config['description'])
    tool_tag.set('version', app_config['branch'] if 'branch' in app_config else 'upstream')
    tool_tag.append(ET.Comment('Generated by CLAMS appliance maker'))

    ET.SubElement(tool_tag, 'description')
    ET.SubElement(tool_tag, 'command')

    inputs_tag = ET.SubElement(tool_tag, 'inputs')
    input_param_tag = ET.SubElement(inputs_tag, 'param')
    input_param_tag.set('name', 'input')
    input_param_tag.set('type', 'data')
    input_param_tag.set('format', 'json')
    input_param_tag.set('label', 'Input MMIF')

    outputs_tag = ET.SubElement(tool_tag, 'outputs')
    output_data_tag = ET.SubElement(outputs_tag, 'data')
    output_data_tag.set('name', 'output')
    output_data_tag.set('format', 'json')
    output_data_tag.set('label', f'{app_name} annotations')

    help_tag = ET.SubElement(tool_tag, 'help')
    help_tag.text = f'{app_config["description"]} (This tool is automatically configured by CLAMS appliance maker)'
    category_tag = ET.SubElement(tool_tag, 'categories')
    category_tag.text = app_config['type'].title()

    return ET.ElementTree(tool_tag)


def appmetadata_to_config_xml_tree(appmetadata):
    if not isinstance(appmetadata, dict):
        appmetadata = json.loads(appmetadata)

    tool_tag = ET.Element('tool')
    
    def escape_to_galaxy_tool_id(s: str):
        escaped = ""
        for c in s:
            if not c.isalpha() and not c.isdigit() and c not in "-_":
                escaped += "_"
            else:
                escaped += c
        return escaped
                
    tool_tag.set('id', escape_to_galaxy_tool_id(appmetadata['identifier']))
    tool_tag.set('name', appmetadata['name'])
    tool_tag.set('version', appmetadata['app_version'])
    tool_tag.append(ET.Comment('Generated by CLAMS appliance maker'))

    ET.SubElement(tool_tag, 'description')
    ET.SubElement(tool_tag, 'command')

    inputs_tag = ET.SubElement(tool_tag, 'inputs')
    input_mmif_tag = ET.SubElement(inputs_tag, 'param')
    input_mmif_tag.set('name', 'input')
    input_mmif_tag.set('type', 'data')
    input_mmif_tag.set('format', 'json')
    input_mmif_tag.set('label', 'Input MMIF')
    runtime_params_section = ET.SubElement(inputs_tag, 'section')
    runtime_params_section.set('name', 'runtime_params')
    runtime_params_section.set('title', 'Configurations')
    runtime_params_section.set('expanded', 'true')
    if 'parameters' in appmetadata:
        for parameter in appmetadata['parameters']:
            runtime_param_tag = ET.SubElement(runtime_params_section, 'param')
            runtime_param_tag.set('name', parameter['name'])
            runtime_param_tag.set('help', parameter['description'])
            if parameter['type'] == 'integer':
                runtime_param_tag.set('type', 'integer')
                if 'default' in parameter:
                    runtime_param_tag.set('value', str(parameter['default']))
            elif parameter['type'] == 'number':
                runtime_param_tag.set('type', 'float')
                if 'default' in parameter:
                    runtime_param_tag.set('value', str(parameter['default']))
            elif parameter['type'] == 'string':
                runtime_param_tag.set('type', 'text')
                if 'choices' in parameter:
                    for choice in parameter['choices']:
                        option_tag = ET.SubElement(runtime_param_tag, 'option')
                        option_tag.set('value', choice)
                        option_tag.text = choice
                        if 'default' in parameter and choice == parameter['choices']:
                            option_tag.set('selected', 'true')
                else:
                    if 'default' in parameter:
                        runtime_param_tag.set('value', parameter['default'])
            elif parameter['type'] == 'boolean':
                runtime_param_tag.set('type', 'boolean')
                if 'default' in parameter and parameter['default'] == 'true':
                    runtime_param_tag.set('checked', 'true')
        
    outputs_tag = ET.SubElement(tool_tag, 'outputs')
    output_data_tag = ET.SubElement(outputs_tag, 'data')
    output_data_tag.set('name', 'output')
    output_data_tag.set('format', 'json')
    output_data_tag.set('label', f'{appmetadata["name"]} annotations')

    help_tag = ET.SubElement(tool_tag, 'help')
    help_msg = StringIO()
    help_msg.write(f"{appmetadata['name']}\n{'=' * len(appmetadata['name'])}\n")
    help_msg.write(f"{appmetadata['description']}\n\n")
    help_msg.write(f"\nLicense\n=======\n")
    help_msg.write(f"This app is licensed under {appmetadata['license']}.\n")
    help_msg.write(f"\nInput Requirements\n==================\n")
    help_msg.write(f"This app can take a MMIF with these annotations:\n\n")
    for inp in appmetadata['input']:
        help_msg.write(f"* {inp['@type']}{'' if inp['required'] else ' (optional)'}\n")
    help_msg.write(f"\nOutput Annotations\n==================\n")
    help_msg.write(f"And the app can add these annotations:\n\n")
    for outp in appmetadata['output']:
        help_msg.write(f"* {outp['@type']}\n")
    help_tag.append(CDATA(help_msg.getvalue()))
    
    category_tag = ET.SubElement(tool_tag, 'categories')
    categories = ""
    for inp in appmetadata['input']:
        at_type = inp['@type'].split('/')[-1] 
        if at_type in GALAXY_CATEGORIES:
            categories += f'{GALAXY_CATEGORIES[at_type]},'
    category_tag.text = categories

    return ET.ElementTree(tool_tag)


def add_to_tool_conf_xml(tool_conf_tree: ET.ElementTree, config_xml_tree: ET.ElementTree, app_name):
    categories = config_xml_tree.find('categories').text.split(',')
    tool_tag = ET.Element('tool', {'file': get_tool_config_xml_filename(app_name)})
    for category in categories:
        category_exists = False
        for section in tool_conf_tree.findall('section'):
            if section.attrib['id'] == category:
                section.append(tool_tag)
                category_exists = True
        if not category_exists:
            section_tag = ET.Element('section', {'id': category, 'name': f"{category} Apps"})
            section_tag.append(tool_tag)
            tool_conf_tree.getroot().append(section_tag)


def gen_display_app_xml(consumer_name, port, description):
    display_tag = ET.Element('display', {'id': consumer_name, 'version': '1.0.0', 'name': description})
    link_tag = ET.SubElement(display_tag, 'link', {'id': 'open', 'name': 'open'})
    url_tag = ET.SubElement(link_tag, 'url')
    url_tag.text = f'http://{HOSTNAME}:{port}/display?file=${{qp("/".join($txt_file.url.split("/")[:2] + ["{GALAXY_CONTNAME}:{GALAXY_CONTPORT}"] + $txt_file.url.split("/")[3:]))}}'
    param_tag = ET.SubElement(link_tag, 'param', {'type': 'data', 'name': 'txt_file', 'url': 'galaxy.txt'})
    ET.ElementTree(display_tag).write(pjoin(get_display_app_xml_fullpath(consumer_name)), encoding='utf-8')


def add_to_datatypes_conf_xml(datatypes_conf_tree: ET.ElementTree, consumer_name):
    display_tag = ET.Element('display', {'file': get_display_app_xml_filename(consumer_name)})
    for dt in datatypes_conf_tree.find('registration').findall('datatype'):
        if dt.attrib['extension'] == 'json':
            dt.append(display_tag)


def gen_db_loc_files(host_data_path):
    for mtype in ['text', 'video', 'image', 'audio']:
        type_path = pjoin(host_data_path, mtype)
        if os.path.exists(type_path) and os.path.isdir(type_path):
            with open(pjoin(GALAXY_LOCAL_PATH, 'tool-data', f'{mtype}db.loc'), 'w') as loc_file:
                for f_name in os.listdir(type_path):
                    if os.path.isfile(pjoin(type_path, f_name)) and not f_name.startswith('.'):
                        loc_file.write(f'{f_name}\t{pjoin(CONTAINER_DATA_PATH, mtype, f_name)}\n')


def download_app(app_name, app_config):
    if app_config['enabled'] and not os.path.exists(app_name):
        if 'branch' in app_config:
            more_params = ['--branch', app_config['branch']]
        else:
            more_params = []
        download(app_config['repository'], app_name, more_params)


def download(repo_url, clone_dir, more_params=[]):
    if DEVELOP:
        os.symlink(pjoin('..', clone_dir), clone_dir)
    else:
        subprocess.run(['git', 'clone', '--depth', '1'] + more_params + [ repo_url, clone_dir], check=True)


def docker_run(image_name, container_name):
    subprocess.run(['docker', 'run', '--rm', '--name', container_name, '-d', image_name], check=True)


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser(
        description="Make a CLAMS-Galaxy appliance using docker-compose"
    )
    parser.add_argument(
        '-v', '--volumename',
        default='clams_appliance_volume',
        action='store',
        nargs='?',
        help=''
    )
    parser.add_argument(
        '-d', '--develop',
        action='store_true',
        help='Run the script in *develop* mode. In develop mode Galaxy, apps, and consumers are copied from local file system, instead of being downloaded from github.'
    )
    parser.add_argument(
        '-f', '--force-rebuild',
        action='store_true',
        help='Delete existing CLAMS Apps and Galaxy. Then download all and re-build docker images.'
    )
    args = parser.parse_args()
    galaxy_export_volumename = args.volumename
    if args.develop:
        now = datetime.datetime.now().strftime("%y%m%d_%H%M%S")
        galaxy_export_volumename = f'{galaxy_export_volumename}_{now}' 

    create_docker_compose(read_config('config.yaml'), galaxy_export_volumename, args.force_rebuild, args.develop)
    # subprocess.run(['docker-compose', 'up'])
