#!/usr/bin/env python3

from unittest import mock
import argparse
import enum
import json
import natsort
import os
import requests
import shutil
import subprocess
import tempfile
import time
import webbrowser

GET_SIMULATORS_ENDPOINT = 'https://api.biosimulators.{}/simulators/latest'

EXAMPLE_COMBINE_ARCHIVES_BASE_URL = 'https://github.com/biosimulators/Biosimulators_test_suite/raw/{}/examples/'

EXAMPLE_SIMULATIONS_FILENAME = os.path.join(os.path.dirname(__file__), 'example-projects.json')

EXAMPLE_SIMULATIONS_RUNS_FILENAME = os.path.join(os.path.dirname(__file__),
                                                 '..', 'apps', 'dispatch', 'src', 'app', 'components',
                                                 'simulations', 'browse', 'example-simulations.{}.json')


class BColors(str, enum.Enum):
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKCYAN = '\033[96m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'


def get_submit_endpoint(runbiosimulations_api):

    if(runbiosimulations_api == 'local'):
        return 'http://localhost:3333/runs'

    return 'https://api.biosimulations.{}/runs'.format(runbiosimulations_api)


def main(runbiosimulations_api='dev', biosimulators_api='org',
         biosimulators_test_suite_branch='deploy', example_names=None,
         test_mode=False, browser=False, dry_run=False):
    """ Submit example simulations from the BioSimulators test suite to the runBioSimulations API and
    record their runs to ``example-simulations.json`` within the browse simulations module of the
    dispatch app so that users can load runs of these simulations as examples.

    Args:
        runbiosimulations_api (:obj:`str`): which deployment of the runBioSimulations API to use (``dev``, ``org``, or ``local``)
        biosimulators_api (:obj:`str`): which deployment of the BioSimulators API to use (``dev`` or ``org``)
        biosimulators_test_suite_branch (:obj:`str`): branch of the BioSimulators test suite to use (e.g., ``deploy`` or ``dev``).
        example_names (:obj:`list` of :obj:`str`): names of example simulations to execute. Default: execute all examples.
        test_mode (:obj:`bool`, optional): whether to run this program in test mode
        browser (:obj:`bool`, optional): whether to open the simulations in a browser
        dry_run (:obj:`bool`, optional): If :obj:`True`, do not submit simulations to runBioSimulations
    """

    # get latest version of each simulator
    response = requests.get(GET_SIMULATORS_ENDPOINT.format(biosimulators_api))
    response.raise_for_status()
    simulator_latest_versions = {
        simulator['id']: simulator['version'] for simulator in response.json()}
    simulator_specs = {
        simulator['id']: simulator for simulator in response.json()}

    # read simulations
    with open(EXAMPLE_SIMULATIONS_FILENAME, 'r') as file:
        simulations = json.load(file)

    # filter out disabled runs
    simulations = list(filter(
        lambda simulation:
        not simulation.get('disabled', False),
        simulations))

    # filter out runs for simulators that aren't enabled
    simulations = list(filter(
        lambda simulation:
        simulator_specs[simulation['simulator']]['biosimulators']['validated'],
        simulations))

    # filter to selected simulations
    if example_names:
        simulations = list(
            filter(lambda simulation: simulation['name'] in example_names, simulations))

        missing_example_names = set(example_names).difference(
            set(simulation['name'] for simulation in simulations))
        if missing_example_names:
            raise ValueError('No examples have the following names:\n  - {}'.format(
                '\n  - '.join(sorted(missing_example_names))))

    # execution simulations
    temp_dir = tempfile.mkdtemp()
    new_simulation_runs = []
    if example_names:
        with open(EXAMPLE_SIMULATIONS_RUNS_FILENAME.format(runbiosimulations_api), 'r') as file:
            simulation_runs = {run['name']: run for run in json.load(file)}

    else:
        simulation_runs = {}

    print(BColors.HEADER + "Submitting Simulations" + BColors.ENDC)
    for simulation in simulations:
        name = simulation['name']
        simulator = simulation['simulator']
        simulator_version = simulator_latest_versions[simulator]
        url = EXAMPLE_COMBINE_ARCHIVES_BASE_URL.format(
            biosimulators_test_suite_branch) + simulation['filename']
        # get COMBINE archive
        response = requests.get(url)
        response.raise_for_status()

        archive_filename = os.path.join(
            temp_dir, os.path.basename(simulation['filename']))
        combine_archive = response.content

        with open(archive_filename, 'wb') as file:
            file.write(combine_archive)
        project_size = len(combine_archive)

        # submit simulation
        submit_endpoint = get_submit_endpoint(runbiosimulations_api)
        if not dry_run:
            process = subprocess.run(['curl',
                                      '-X', 'POST', submit_endpoint,
                                      '-H', "accept: application/json",
                                      '-H', "Content-Type: application/json",
                                      '-d', json.dumps({
                                          "name": name,
                                          'simulator': simulator,
                                          "simulatorVersion": simulator_version,
                                          "cpus": 1,
                                          "memory": 8,
                                          "maxTime": 20,
                                          "envVars": [],
                                          "purpose": "academic",
                                          "email": None,
                                          "public": simulation['publishToBioSimulations'],
                                          "url": url,
                                      }),
                                      ], capture_output=True)
        else:
            process = mock.Mock(stdout=mock.Mock(decode=lambda: json.dumps({
                'id': 'XXXXXXXXX',
                'submitted': 'XXXXXXXXX',
                'updated': 'XXXXXXXXX',
            })))

        simulation_run = json.loads(process.stdout.decode())
        if 'id' in simulation_run:
            print("Submitted simulation: " + simulation_run['id'])
            sim_url = "https://run.biosimulations.{}/simulations/{}".format(
                runbiosimulations_api if not runbiosimulations_api == "local" else "dev", simulation_run['id'])
            print("View: " + sim_url)
            if browser:
                webbrowser.open(sim_url)
        else:
            print(BColors.FAIL + "Failed to submit simulation: " + name + BColors.ENDC)
            raise RuntimeError('Simulation could not be run on {}: {}'.format(
                submit_endpoint,
                simulation_run['error']))

        # log run
        simulation_runs[name] = {
            "id": simulation_run['id'],
            "name": name,
            'simulator': simulator,
            "simulatorVersion": simulator_version,
            "cpus": 1,
            "memory": 8,
            "maxTime": 20,
            "envVars": [],
            "purpose": "academic",
            "submittedLocally": False,
            "status": "CREATED",
            "submitted": simulation_run['submitted'],
            "updated": simulation_run['updated'],
            "projectSize": project_size,
        }
        new_simulation_runs.append(simulation_runs[name])

    if test_mode:
        monitor_runs(new_simulation_runs, runbiosimulations_api)

    else:
        sorted_simulation_runs = natsort.natsorted(list(simulation_runs.values()),
                                                   key=lambda run: (run['name'], run['simulator'], run['simulatorVersion']),
                                                   alg=natsort.IGNORECASE)
        with open(EXAMPLE_SIMULATIONS_RUNS_FILENAME.format(runbiosimulations_api), 'w') as file:
            json.dump(sorted_simulation_runs, file, indent=2)

    shutil.rmtree(temp_dir)


def monitor_runs(simulation_runs, runbiosimulations_api, timeout=240):
    total = len(simulation_runs)
    pending_runs = simulation_runs
    failed_runs = []
    passed_runs = []
    timed_out_runs = []
    start_time = time.time()
    current_time = time.time()
    timeout_occurred = False
    print(BColors.HEADER + "Monitoring Simulation Runs with timeout" + BColors.BOLD + str(timeout) + BColors.ENDC)
    while(not timeout_occurred):
        print(BColors.HEADER + "Getting Runs from API" + BColors.ENDC)
        if len(pending_runs) == 0:
            break
        for run in pending_runs[:]:

            api_host = get_submit_endpoint(runbiosimulations_api)
            print("Checking status of {}".format(run['name']))
            response = requests.get(api_host + "/" + run['id'])
            try:
                response.raise_for_status()
            except requests.exceptions.HTTPError as e:
                pending_runs.remove(run)
                failed_runs.append(run)

            sim = response.json()
            if sim['status'] == 'SUCCEEDED':
                print(BColors.OKGREEN + "SUCCEEDED " + BColors.ENDC + run['name'])
                pending_runs.remove(run)
                passed_runs.append(run)
            elif sim['status'] == 'FAILED':
                print(BColors.FAIL + "FAILED " + BColors.ENDC + run['name'])
                pending_runs.remove(run)
                failed_runs.append(run)
            elif sim['status'] == 'RUNNING':
                print(BColors.OKCYAN + "RUNNING " + BColors.ENDC + run['name'])

            elif sim['status'] == 'PROCESSING':
                print(BColors.OKCYAN + "PROCESSING " + BColors.ENDC + run['name'])

            elif sim['status'] == 'CREATED':
                print(BColors.OKBLUE + " CREATED " + BColors.ENDC + run['name'])

            elif sim['status'] == 'QUEUED':
                print(BColors.WARNING + "QUEUED " + BColors.ENDC + run['name'])
            else:
                print(BColors.WARNING + "UNKNOWN " + BColors.ENDC + run['name'])

        current_time = time.time()
        if(len(pending_runs) == 0):
            break
        if(current_time-start_time > timeout):
            timeout_occurred = True
            for run in pending_runs:
                timed_out_runs.append(run)
        total_failed_runs = len(failed_runs)
        total_passed = len(passed_runs)
        total_timedout = len(timed_out_runs)
        total_pending = len(pending_runs)
        print()
        print(BColors.HEADER + "Total Number of runs: " + str(total) + BColors.ENDC)
        print(BColors.OKBLUE + "Total Number of pending runs: " + str(total_pending) + BColors.ENDC)
        print(BColors.OKGREEN + "Total Number of passed runs: " + str(total_passed) + BColors.ENDC)
        print(BColors.WARNING + "Total Number of timed out runs: " + str(total_timedout) + BColors.ENDC)
        print(BColors.FAIL + "Total Number of failed runs: " + str(total_failed_runs) + BColors.ENDC)
        print()
        time.sleep(5)
    passed = True
    if len(failed_runs) > 0:
        print(BColors.FAIL + "One or more simulations failed" + BColors.ENDC)
        passed = False
        for run in failed_runs:
            print(BColors.FAIL + "FAILED {} --- {}".format(run['id'], run['name']) + BColors.ENDC)
    if len(timed_out_runs) > 0:
        print(BColors.WARNING + "One or more simulations timed-out" + BColors.ENDC)
        passed = False
        for run in timed_out_runs:
            print(BColors.WARNING + "TIMEOUT {} --- {}".format(run['id'], run['name']) + BColors.ENDC)
    if passed:
        print(BColors.OKGREEN + "All runs passed" + BColors.ENDC)

    return passed


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description='Submit the example simulations to the runBioSimulations API and save their runs to the dispatch app.')
    parser.add_argument(
        '--runbiosimulations-api', type=str, default='dev',
        help='runBioSimulations API which simulations should be submitted to (`dev`, `org`, `local`). Default: `dev`.')
    parser.add_argument(
        '--biosimulators-api', type=str, default='org',
        help=('BioSimulators API which should be used to select the version of each simulation tool used '
              'to execute simulations (`dev`, `org`). Default: `org`.'))
    parser.add_argument(
        '--biosimulators-test-suite-branch', type=str, default='deploy',
        help=('Branch of the BioSimulators test suite from which the example COMBINE/OMEX archives should be obtained. '
              'Default: `deploy`.'))
    parser.add_argument(
        '--example', type=str, nargs='*',
        help='Names of the example simulations to execute. Default: execute all simulations.',
        default=None, dest='example_names',

    )
    parser.add_argument(
        '--test', type=bool,
        help='If set, run the script in test mode to test API functionality',
        default=False, dest='test_mode',

    )
    parser.add_argument(
        '--web', type=bool,
        help='If set, open a browser window with simulation runs',
        default=False, dest='browser',

    )
    parser.add_argument(
        '--dry-run',
        help='If set, do not submit simulations to runBioSimulations',
        action='store_true',

    )
    args = parser.parse_args()

    main(runbiosimulations_api=args.runbiosimulations_api,
         biosimulators_api=args.biosimulators_api,
         biosimulators_test_suite_branch=args.biosimulators_test_suite_branch,
         example_names=args.example_names,
         test_mode=args.test_mode,
         browser=args.browser,
         dry_run=args.dry_run)
