#!/usr/bin/env python3

import os
import sys
import logging
import json
import time
import pwd

from meeseeks import Client
from meeseeks.config import Config
from meeseeks.util import *

def usage():

    print("""
    
    meeseeks-client [client-options] <command> [args...]

    or use q-symlinks: q{stat|sub|job|del|mod|conf} [args]

    commands are:
        sub [submit-options] <pool[@node]> <executable> [args....] (submits job, returns job info)
            submit-options:
                pool= sets pool for job to run in
                node= sets node(s) for job to run on. Can be comma-seperated list or end with * for all/wildcard
                stdin= stdout= stderr= (redirect job in/out/err to files named)
                restart= (1=restart when done)
                retries= (times to retry if job fails)
                resubmit= (1=resubmit when job is done or fails)
                runtime= (max runtime of job)
                id= (set new job's id or submit changes to existing job)
                state= (change existing job state, 'new' will restart finished job)
                hold= (1=queue but do not start job)
                tag= list of tags, can be matched in query with tag=

        job|get [jobids|filter] (get all or specified job info as JSON)
            jobids are any number of job id arguments
            filter is any job attributes such as node= pool= state= tag=

        ls [filter] (list job ids)

        del|kill <jobids|filter> (kill job)

        mod|set <jobids|filter : > key=value ... (set key=value in jobs matching jobids or filter, return new job info)
            if a filter is provided, ':' is used to delimit filter key=value from job key=value
            set a job in any finished state (done,failed,killed) to state='new' to restart job
            if node is not specified when restarting job, node will be cleared.

        stat|show [filter] {nodes pools jobs active tree} 
            (prints flat or tree cluster status, 
             specify which elements to show, defaults to flat/all)
             Job Flags: A=active H=hold E=error R=repeating

        nodes (prints full node status JSON)

        pools (prints full pool status JSON)

        conf [key=value] [node]
            get/sends config to directly connected or specified node

    client-options can be set by the environment var MEESEEKS_CONF and default to: 
        address= (default localhost)
        port= (defult 49463)
        refresh= (interval to continuously refresh status until no jobs left)

""")

def pretty_print(d): print (json.dumps(d, sort_keys=True, indent=4))

def show(client,args,**kwargs):     
    args=[arg[0] for arg in args]
    status=client.request({'nodes':None,'pools':None}) 
    nodes=status.get('nodes',{})
    pools=status.get('pools',{})
    kwpool=kwargs.get('pool')
    if kwpool: del kwargs['pool']
    kwnode=kwargs.get('node')
    if kwnode: del kwargs['node']
    if 'n' not in args and 'p' not in args and 'j' not in args:
        args.extend(['n','p','j'])
    #put pool and job status on nodes
    for n,node in nodes.items():
        if kwnode and n != kwnode: continue
        for p,pnodes in pools.items():
            if kwpool and p != kwpool: continue
            if n in pnodes: node.setdefault('free',{})[p]=pnodes[n]
        jobs=client.query(node=n,**kwargs)
        for jid,job in jobs.items(): 
            if job['pool'] == '__config': continue
            node.setdefault('jobs',{}).setdefault(job['pool'],{})[jid]=job
    #now we put nodes under other nodes
    def _tree(s):
        for n in list(s.keys()):
            if n in s: #already moved?
                node=s[n]
                if not node.get('online'): continue
                for dn in node.get('routing',[]):
                    if dn in s and dn != n: #we can move this node
                        s[n].setdefault('next',{})[dn]=s[dn]
                        del s[dn]
                if s[n].get('next'): _tree(s[n]['next']) #walk downstream
    if 't' in args: _tree(nodes)
    #now we print the tree
    def _print(s,i=0):
        for n,node in sorted(s.items()):
            if 'n' in args:
                l='\t'*i + n+'\t'
                if node.get('online'): l+='up'
                else: l+='down'
                l+='\t%s'%node.get('loadavg')
                if n: print(l)
            pn=set(node.get('free',{}).keys())|set(node.get('jobs',{}).keys())
            for p in sorted(pn):
                if 'n' in args: l='\t'*i + '\t'
                else: l=''
                l+='%s@%s'%(p,n)
                jobs=node.get('jobs',{}).get(p,{})
                jc=len(jobs)
                l+='\t%s jobs'%jc
                free=node.get('free',{}).get(p)
                if type(free) is int: l+='\t%s free'%max(free,0)
                if 'p' in args and ('a' not in args or jc): print (l)
                l=''
                if 'n' in args: l+='\t'*i+'\t'
                if 'p' in args: l+='\t'
                #else: l+='%s\t'%pn
                if 'j' in args:
                    for (jid,job) in sorted(jobs.items(),key=lambda j: (j[1].get('state')) ):
                        flags=''
                        rc=''
                        if job.get('active'): 
                            flags+='A'
                            if 'start_ts' in job: rc='+'+str(int(time.time()-job['start_ts']))
                        elif job.get('rc'): 
                            rc='(%s'%job['rc']
                            if job.get('retries') and job.get('fail_count'): rc+=' %s/%s'%(job['fail_count'],job['retries'])
                            rc+=')'
                        if job.get('error'): flags+='E'
                        if job.get('hold'): flags+='H'
                        if job.get('restart'): flags+='R'
                        if job.get('resubmit'): flags+='S'
                        try: user=pwd.getpwuid(job.get('uid')).pw_name
                        except: user=job.get('uid')
                        print(l+'%s\t%s\t[%s]\t%s\t%s\t%s\t%s'%(
                            jid,user,','.join(str(tag) for tag in job['tags']),flags,job['state'],rc,' '.join(job.get('args',[]))) )
            if 'next' in node:
                _print(node['next'],i+1)
    _print(nodes,0)

#parse env vars
cfg=cmdline_parser(os.getenv('MEESEEKS_CONF','').split())[0]
try:
    # cmd=first non cfg argument
    # args=remaining arguments 
        
    #handle cmd from symlink to meeseeks-client
    EXEC_NAME=os.path.basename(sys.argv[0])
    if EXEC_NAME.startswith('q'):cmd=EXEC_NAME[1:]
    else: cmd=None
    
    #get command line options/args
    args=sys.argv[1:]
    if not cmd:
        c,args=cmdline_parser(args)
        if not args: sys.exit(usage())
        cfg.update(c)
        cmd,args=args[0],args[1:] #end of options
    
    #load cfg files
    if 'file' in cfg: cfg.update(read_cfg_files(cfg['file']))
    
    #parse post-cmd args
    kwargs,args=cmdline_parser(args)

    #set logging
    logging.basicConfig(level=cfg.get('log',logging.INFO))
    logging.debug('%s %s %s %s'%(cfg,cmd,kwargs,args))

    client=Client(**cfg.copy()) #create client and connect
    
    #if we will be refreshing, wait until state populated
    if cfg.get('refresh'): client.wait()
    jids=None

    if cmd.startswith('sub'):
        #get pool and node selection from kwargs
        pool=jobargs=node=None
        if 'node' in kwargs:
            node=kwargs['node']
            del kwargs['node']
        if 'pool' in kwargs:
            pool=kwargs['pool']
            del kwargs['pool']
        # if no pool in kwargs, get pool@node arg
        elif args: 
            pool=args[0]
            args=args[1:]
        if pool and '@' in pool: 
            pool,node=pool.split('@',1)
            if ',' in node: node=node.split(',')
        #set submit user and environ in kwargs
        kwargs.update(uid=os.geteuid(),env=dict(os.environ))
        #if we are refreshing submit locally and let it dispatch
        if cfg.get('refresh'): jids=[client.submit_job(pool=pool,args=args,node=node,**kwargs)]
        #submit directly
        else: pretty_print(client.submit(pool=pool,args=args,node=node,**kwargs))

    elif cmd == 'get' or cmd == 'job': #get jobs by ids or filter on key=value
        if cfg.get('refresh'):
            if args: jids=args
            else: jids=True #get all
        elif args: pretty_print(client.query(ids=args))
        else: pretty_print(client.query(**kwargs))

    elif cmd == 'set' or cmd == 'mod': #set k=v in matching jobs
        #kwargs will be set if we have initial k=v args like 'pool=... : key=...' 
        if kwargs: jids=client.ls(**kwargs) #we were provided a filter
        else: jids=[arg for arg in args if '=' not in arg] #no filter, we have 'jobid jobid key=...'
        kwargs,args=cmdline_parser([arg for arg in args if '=' in arg]) #parse following k=v args
        for jid in jids: client.submit(id=jid,**kwargs) #submit changes 
        if jids: pretty_print(client.query(ids=jids)) #print changes 

    elif cmd == 'ls': pretty_print(client.ls(**kwargs))

    elif cmd == 'kill' or cmd == 'del': pretty_print(client.kill(args,**kwargs))

    elif cmd == 'nodes': 
        if cfg.get('refresh'): 
            while True:
                pretty_print(client.get_nodes())
                time.sleep(cfg.get('refresh'))
        else: pretty_print(client.nodes())

    elif cmd == 'pools': 
        if cfg.get('refresh'): 
            while True:
                pretty_print(client.get_pools())
                time.sleep(cfg.get('refresh'))
        else: pretty_print(client.pools())

    elif cmd.startswith('conf'): 
        if 'file' in kwargs: kwargs.update(read_cfg_files(kwargs['file']))
        if args:
            for node in args: 
                pretty_print(client.submit(pool='__config',node=node,args=kwargs.copy()))
        else: pretty_print(client.request(dict(config=kwargs.copy()))) #peer config push

    elif cmd == 'show' or cmd.startswith('stat'): show(client,args,**kwargs) #pretty print full cluster

    else: sys.exit(usage())
    
    #refreshing on jobs
    if jids and cfg.get('refresh'):
        seq={} #sequence numbers for printing updates
        while True: 
            if jids is True: jobs=client.get(**kwargs)
            else: jobs=dict((jid,client.get_job(jid)) for jid in jids)
            if not any(jobs.values()): break #exit when no jobs left
            #print updates (when job seq != saved seq)
            updates=dict((jid,job) for (jid,job) in jobs.items() if job['seq'] != seq.get(jid))
            if updates: pretty_print(updates)
            seq=dict((jid,job['seq']) for (jid,job) in jobs.items())
            time.sleep(cfg.get('refresh'))
except Exception as e: raise e#print (e,file=sys.stderr)

#if we were polling state, stop gracefully
if cfg.get('refresh'): client.close()
