#!/usr/bin/env python3

import datetime
from tracemalloc import start
from unittest import mock
import argparse
import enum
import json
import json.decoder
import natsort
import os
import requests
import shutil
import subprocess
import tempfile
import time
import webbrowser

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_simulators_endpoint(biosimulators_deployment):

    if biosimulators_deployment == 'local':
        return 'http://localhost:3333/simulators/latest'
    elif biosimulators_deployment == 'dev':
        return 'https://api.biosimulators.dev/simulators/latest'
    else:
        return 'https://api.biosimulators.org/simulators/latest'


def get_submit_run_endpoint(runbiosimulations_deployment):

    if runbiosimulations_deployment == 'local':
        return 'http://localhost:3333/runs'
    elif runbiosimulations_deployment == 'dev':
        return 'https://api.biosimulations.dev/runs'
    else:
        return 'https://api.biosimulations.org/runs'


def get_simulation_run_view_url(runbiosimulations_deployment, run_id):

    if runbiosimulations_deployment == 'local':
        return 'http://localhost:4200/runs/{}'.format(run_id)
    elif runbiosimulations_deployment == 'dev':
        return 'https://run.biosimulations.dev/runs/{}'.format(run_id)
    else:
        return 'https://run.biosimulations.org/runs/{}'.format(run_id)


def main(runbiosimulations_deployment='dev', biosimulators_deployment='prod',
         biosimulators_test_suite_branch='deploy', example_names=None,
         test_mode=False, browser=False, dry_run=False, timeout=999, stress=1):
    """ 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_deployment (:obj:`str`): which deployment of the runBioSimulations deployment to use (``dev``, ``prod``, or ``local``)
        biosimulators_deployment (:obj:`str`): which deployment of the BioSimulators deployment to use (``dev`` or ``prod``)
        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(biosimulators_deployment))
    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))))

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

    else:
        simulation_runs = {}

    simulations = simulations * stress
    print(BColors.HEADER + "Submitting {} simulations ...".format(len(simulations)) + BColors.ENDC)
    for i_simulation, simulation in enumerate(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/OMEX 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
        start_time = time.time()
        submit_endpoint = get_submit_run_endpoint(runbiosimulations_deployment)
        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',
            })))

        stdout = process.stdout.decode()
        try:
            simulation_run = json.loads(stdout)
        except json.decoder.JSONDecodeError as exception:
            simulation_run = None
            print(BColors.FAIL + "Failed to submit simulation: " + name + BColors.ENDC)
            raise RuntimeError('Simulation could not be run on {}: {}'.format(
                submit_endpoint,
                stdout))

        if 'id' in simulation_run:
            print("Submitted simulation {}: {} ({})".format(i_simulation + 1, name, simulation_run['id']))
            get_simulation_run_view_url
            sim_url = get_simulation_run_view_url(runbiosimulations_deployment, 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:
        passed = monitor_runs(new_simulation_runs, runbiosimulations_deployment, start_time, timeout)
        if not passed:
            exit(1)
        else:
            exit(0)

    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_filename, 'w') as file:
            json.dump(sorted_simulation_runs, file, indent=2)

    shutil.rmtree(temp_dir)


def monitor_runs(simulation_runs, runbiosimulations_deployment, start_time, timeout=9999):
    total = len(simulation_runs)
    pending_runs = simulation_runs
    failed_runs = []
    passed_runs = []
    timed_out_runs = []
    start_time = start_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_run_endpoint(runbiosimulations_deployment)
            print("Checking status of {} ({})".format(run['name'], run['id']))
            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()
        total_failed_runs = len(failed_runs)
        total_passed = len(passed_runs)
        total_timedout = len(timed_out_runs)
        total_pending = len(pending_runs)
        if(current_time-start_time > timeout):
            timeout_occurred = True
            for run in pending_runs:
                timed_out_runs.append(run)
        print()
        print(BColors.HEADER + "Total Number of runs: " + str(total) + BColors.ENDC)
        print(BColors.HEADER + "Total time: " + str(datetime.timedelta(seconds=(current_time-start_time))) + 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()
        if(total_pending == 0):
            break

        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:
            api_host = get_submit_run_endpoint(runbiosimulations_deployment)
            response = requests.get(api_host + "/" + run['id'])
            sim = response.json()

            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)
        print(BColors.OKBLUE + "Starting validation" + BColors.ENDC)
        invalid_run_count = validate_runs(passed_runs, runbiosimulations_deployment)
        if invalid_run_count > 0:
            total_passed-=invalid_run_count
            total_failed_runs+=invalid_run_count
    print()
    print(BColors.HEADER + "Total Number of runs: " + str(total) + BColors.ENDC)
    print(BColors.HEADER + "Total time: " + str(datetime.timedelta(seconds=(current_time-start_time))) + 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()
    return passed


def validate_runs(runs,  deployment):
    api_host = get_submit_run_endpoint(deployment)
    valid_runs = []
    invalid_runs = []
    for run in runs:
        print("Checking status of {} ({})".format(run['name'], run['id']))
        response = requests.get(api_host + "/" + run['id'] + "/validate")
        try:
            response.raise_for_status()
            print(BColors.OKGREEN + "VALID " + BColors.ENDC + run['name'] + " (" + run['id'] + ")")
            valid_runs.append(run)
        except requests.exceptions.HTTPError as e:
            print(BColors.FAIL + "INVALID " + BColors.ENDC + run['name'] + " (" + run['id'] + ")")
            invalid_runs.append(run)
    if len(invalid_runs) > 0:
        print(BColors.FAIL + "One or more simulations failed" + BColors.ENDC)
        for run in invalid_runs:
            print(BColors.FAIL + "FAILED {} --- {}".format(run['id'], run['name']) + BColors.ENDC)
        return False
    else:
        print(BColors.OKGREEN + "All runs passed" + BColors.ENDC)
        return len(invalid_runs)


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-deployment', type=str, default='dev',
        help='runBioSimulations deployment which simulations should be submitted to (`dev`, `prod`, `local`). Default: `dev`.')
    parser.add_argument(
        '--biosimulators-deployment', type=str, default='prod',
        help=('BioSimulators deployment which should be used to select the version of each simulation tool used '
              'to execute simulations (`dev`, `prod`). Default: `prod`.'))
    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',

    )
    parser.add_argument(
        '--timeout',
        type=int,
        default=999,
        help='Timeout in minutes for waiting for simulations to complete. Default: 999',
        dest='timeout',
    )
    parser.add_argument(
        '--stress',
        type=int,
        default=1,
        help='Number of times to run each simulation. Default: 1',
        dest='stress',
    )

    args = parser.parse_args()

    main(runbiosimulations_deployment=args.runbiosimulations_deployment,
         biosimulators_deployment=args.biosimulators_deployment,
         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,
         timeout=args.timeout,
         stress=args.stress)
