#!python

import sys
from unicodedata import name
from SuperfacilityAPI import (
    SuperfacilityAPI,
    SuperfacilityAccessToken,
    SuperfacilityErrors
)
from SuperfacilityAPI.nersc_systems import NERSC_DEFAULT_COMPUTE

import click
from pathlib import Path
import json
import logging

runasync = False


def click_json(*args, **kwargs):
    click.echo(json.dumps(*args))


def check_file_and_open(file_path: str = "") -> str:
    contents = None
    pth = Path(file_path)
    if pth.is_file():
        with open(pth.absolute()) as f:
            contents = f.read()
    return contents


@click.group()
@click.option('--clientid', '-id', default=None,
              help='Client ID for your key. Can be used to specify which key to look for in $HOME/.superfacility.')
@click.option('--client', '-c', default=None,
              help='Client ID for your key. Can be used to specify which key to look for in $HOME/.superfacility.')
@click.option('--debug', '-d', is_flag=True, default=False,
              help='Print debug messages from sfapi and SuperfacilityConnector')
@click.option('--sync', is_flag=True, default=False, help='Run async')
@click.pass_context
def cli(ctx, client, clientid, debug, sync):
    # Entrypoint for all the cli subcommands
    # Basically an __init__ function that sets up the sfapi
    ctx.ensure_object(dict)
    if debug:
        logging.basicConfig(encoding='utf-8', level=logging.DEBUG)
    else:
        logging.basicConfig(encoding='utf-8', level=logging.ERROR)

    if sync:
        runasync = True

    try:
        access_token = SuperfacilityAccessToken(
            name=client, client_id=clientid)
        sfapi = SuperfacilityAPI(token=access_token.token)
    except:
        sfapi = SuperfacilityAPI()

    ctx.obj['sfapi'] = sfapi


@cli.command()
@click.argument('site', default=NERSC_DEFAULT_COMPUTE)
@click.pass_context
def status(ctx, site):
    sfapi = ctx.obj['sfapi']

    if site in ['compute', 'computes']:
        site = 'cori,perlmutter'
    elif site in ['filesystem', 'filesystems']:
        site = 'dna,dtns,global_homes,projectb,global_common,community_filesystem'
    elif site in ['login', 'logins']:
        site = 'cori,perlmutter,jupyter,dtns'
    try:
        if site == 'all':
            ret = sfapi.status(None)
        else:
            ret = [sfapi.status(site) for site in site.split(",")]

        click_json(ret)
    except Exception as err:
        click.echo(f"{type(err).__name__}: {err}")


@cli.command()
@click.argument('site', default=NERSC_DEFAULT_COMPUTE)
@click.pass_context
def outages(ctx, site):
    sfapi = ctx.obj['sfapi']

    if site in ['compute', 'computes']:
        site = 'cori,perlmutter'
    elif site in ['filesystem', 'filesystems']:
        site = 'dna,dtns,global_homes,projectb,global_common,community_filesystem'
    elif site in ['login', 'logins']:
        site = 'cori,perlmutter,jupyter,dtns'

    try:
        if site == 'all':
            ret = sfapi.status(None, outages=True)
        else:
            ret = [sfapi.status(site, outages=True)
                   for site in site.split(",")]

        click_json(ret)
    except SuperfacilityErrors.InternalServerError as err:
        click.echo(f"{type(err).__name__}: {err}")


@cli.command()
@click.argument('site', default=NERSC_DEFAULT_COMPUTE)
@click.pass_context
def system_status(ctx, site):
    sfapi = ctx.obj['sfapi']
    status_code = sfapi.system_status(name=site)

    output = f"{site}@NERSC is currently {status_code}"
    click.echo(output)


@cli.command()
@click.pass_context
def token(ctx):
    sfapi = ctx.obj['sfapi']
    click.echo(sfapi.access_token)


@cli.command()
@click.pass_context
def systems(ctx):
    sfapi = ctx.obj['sfapi']
    click_json(sfapi.system_names())


@cli.command()
@click.pass_context
def roles(ctx):
    sfapi = ctx.obj['sfapi']
    ret = sfapi.roles()
    try:
        ret = [{k: oj[k] for k in ['repo_name', 'id',
                                   'iris_role', 'description']} for oj in ret]
        click_json(ret)
    except Exception as err:
        click.echo(f"{type(err).__name__}: {err}")


@cli.command()
@click.pass_context
def projects(ctx):
    sfapi = ctx.obj['sfapi']

    ret = sfapi.projects()
    click_json(ret)


@cli.command()
@click.option('--group', '-g', default=None, help='NERSC Group')
@click.pass_context
def group(ctx, group):
    sfapi = ctx.obj['sfapi']

    ret = sfapi.get_groups(groups=group)

    click_json(ret)


@cli.command()
@click.argument('site', default=NERSC_DEFAULT_COMPUTE)
@click.option('--path', '-p', default=None, help='Path to slurm submit file at NERSC.')
@click.option('--local', '-l', default=None, help='Path to local file to submit.')
@click.pass_context
def sbatch(ctx, site, path, local):
    sfapi = ctx.obj['sfapi']
    script = None
    isPath = False

    if path is not None:
        isPath = True
        script = path
    elif local is not None:
        script = check_file_and_open(local)

    ret = sfapi.post_job(site=site, script=script,
                         isPath=isPath, run_async=runasync)
    if runasync:
        click_json(ret)
    else:
        try:
            click.echo(ret['jobid'])
        except Exception as err:
            click.echo(f"{type(err).__name__}: {err}")


@cli.command()
@click.argument('site', default=NERSC_DEFAULT_COMPUTE)
@click.option('--path', '-p', default=None, help='Path to slurm submit file at NERSC.')
@click.pass_context
def ls(ctx, site, path):
    sfapi = ctx.obj['sfapi']

    ret = sfapi.ls(site=site, remote_path=path)
    try:
        click_json(ret['entries'])
    except Exception as err:
        click.echo(f"{type(err).__name__}: {err}")


@cli.command()
@click.argument('site', default=NERSC_DEFAULT_COMPUTE)
@click.option('--path', '-p', default=None, help='Path to slurm submit file at NERSC.')
@click.pass_context
def cat(ctx, site, path):
    sfapi = ctx.obj['sfapi']

    ret = sfapi.download(site=site, remote_path=path, save=False)
    logging.debug(ret)

    try:
        click.echo_via_pager(ret['file'])
    except Exception as err:
        click.echo(f"{type(err).__name__}: {err}")


@cli.command()
@click.argument('site', default=NERSC_DEFAULT_COMPUTE)
@click.option('--sacct/--no-sacct', default=False)
@click.option('--user', '-u', default=None, help='User to to get queue info for.')
@click.option('--jobid', '-j', default=None, help='Specific jobid to get queue info for.')
@click.pass_context
def squeue(ctx, site, sacct, user, jobid):
    sfapi = ctx.obj['sfapi']

    ret = sfapi.get_jobs(site=site, sacct=sacct, user=user, jobid=jobid)

    cols = ['jobid',
            'name',
            'account',
            'cpus',
            'features',
            'partition',
            'reason',
            'start_time',
            'state',
            'submit_time',
            'time',
            'time_left',
            'time_limit',
            ]

    try:
        logging.debug(ret)
        if sacct:
            outputs = ret['output']
        else:
            outputs = [{k: oj[k] for k in cols} for oj in ret['output']]
    except Exception as err:
        click.echo(f"{type(err).__name__}: {err}")
        exit(1)

    for output in outputs:
        click_json(output)


@cli.command()
@click.argument('jobid')
@click.argument('site', default=NERSC_DEFAULT_COMPUTE)
@click.pass_context
def scancel(ctx, site, jobid):
    sfapi = ctx.obj['sfapi']
    logging.info(f"Running delete on {jobid}@{site}")

    ret = sfapi.delete_job(site=site, jobid=jobid)
    click.echo(ret)


@cli.command()
@click.argument('taskid')
@click.pass_context
def task(ctx, taskid):
    sfapi = ctx.obj['sfapi']
    # Waits (up to {timeout} seconds) for the job to be submited before returning
    timeout = 40
    sleeptime = 1
    from time import sleep
    for i in range(timeout):
        if i > 0:
            sleep(sleeptime)

        logging.debug(f"Running {i}")
        task = sfapi.tasks(task_id=taskid)
        logging.debug(f"task = {task}")
        if task is not None and task['status'] == 'completed':
            jobinfo = json.loads(task['result'])
            click.echo(
                {'error': jobinfo['error'], 'jobid': jobinfo['jobid'], 'task_id': taskid})
            return


@cli.command()
@click.option('--client', '-c', default="sfpai", help='Name the sfapi json file')
def manage_keys(client):
    SuperfacilityAccessToken.save_token(tag=client)


if __name__ == '__main__':
    cli(obj={})
