#!/usr/bin/env python

"""
Create a webpage with the dependency graph and status for DESI jobs on a night
"""

import os, sys, glob
import numpy as np
from astropy.table import Table

import argparse

from desispec.workflow.proctable import get_processing_table_path
from desispec.workflow.proctable import get_processing_table_name
from desispec.workflow.queue import queue_info_from_qids
from desispec.io.meta import specprod_root
from desispec.workflow.redshifts import get_ztile_script_pathname

p = argparse.ArgumentParser()
p.add_argument('-n', '--night', type=int, required=True,
        help='night to process')
p.add_argument('-o', '--output', type=str, required=False,
        help='output HTML file (default run/jobgraph/jobgraph-NIGHT.html in specproddir)')
p.add_argument('-s', '--specprod', type=str, required=False,
        help=('override $SPECPROD, or full path '
              'to override $DESI_SPECTRO_REDUX/$SPECPROD'))

args = p.parse_args()

#- Default output file into specprod directory
if args.specprod is not None and os.path.isdir(args.specprod):
    specproddir = args.specprod
else:
    specproddir = specprod_root(args.specprod)

#- override environment variables for get_processing_table_* calls later
os.environ['DESI_SPECTRO_REDUX'] = os.path.dirname(specproddir)
os.environ['SPECPROD'] = os.path.basename(specproddir)

if args.output is None:
    args.output = f'{specproddir}/run/jobgraph/jobgraph-{args.night}.html'
    os.makedirs(os.path.dirname(args.output), exist_ok=True)

#- avoid logging from get_processing_table...
os.environ['DESI_LOGLEVEL'] = 'warning'

#- Read processing table
procpath = get_processing_table_path(None)
procfile = get_processing_table_name(None, str(args.night))
procfile = os.path.join(procpath, procfile)
proctable = Table.read(procfile)

#- Get the job state from rows that are actually slurm jobs
ii = proctable['LATEST_QID']>0
qinfo = queue_info_from_qids(proctable['LATEST_QID'][ii])
jobinfo = {row['JOBID']:row for row in qinfo}

#- Map production directory -> URL for log links
specprodurl = specproddir.replace(os.environ['DESI_ROOT'], 'https://data.desi.lbl.gov/desi')
outputurl = os.path.abspath(args.output).replace(os.environ['DESI_ROOT'], 'https://data.desi.lbl.gov/desi')
outputdir = os.path.dirname(os.path.abspath(args.output))

#- Write the mermaid graph
fx = open(args.output, 'w')
fx.write(f"""<!DOCTYPE html>
<html lang="en">
<head>
  <meta charset="utf-8">
  <style>
    h1 {{font-family: Helvetica, sans-serif; margin-bottom: 0.1em;}}
    p {{font-family: Menlo, Monospace; margin-top: 0.1em;}}
  </style>
</head>
<body>
  <h1>Night {args.night}</h1>
  <p><a href="{specprodurl}">{specproddir}</a></p>
  <div class="mermaid">
  graph TD
    classDef COMPLETED fill:#a8ddb5;
    classDef PENDING   fill:#bdbdbd;
    classDef RUNNING   fill:#bdd7e7;
    classDef FAILED        fill:#d95f0e;
    classDef OUT_OF_MEMORY fill:#d95f0e;
    classDef TIMEOUT       fill:#d95f0e;
    classDef CANCELLED     fill:#fed98e;
    classDef NOTSUBMITTED  fill:#fcae1e;
    classDef UNKNOWN       fill:#ffffcc;
""")

jobs = dict()
state_counter = dict()
for row in proctable:
    qid = row['LATEST_QID']
    intid = row['INTID']        #- internal ID
    jobdesc = row['JOBDESC']

    #- nightly bias job doesn't have expids associated with it
    if row['EXPID'] == '|':
        expid = ''
        e1 = -1
    else:
        expid = row['EXPID'][0:-1].replace('|', ',')  #- comma separated expids
        e1 = int(expid.split(',')[0])  #- the first expid

    tileid = row['TILEID']
    proccamword = row['PROCCAMWORD']

    if qid in jobinfo:
        state = jobinfo[qid]['STATE'].split()[0]
    elif qid == 1:
        state = 'NOTSUBMITTED'
    else:
        state = 'UNKNOWN'

    if state not in state_counter:
        state_counter[state] = 1
    else:
        state_counter[state] += 1

    if jobdesc == 'tilenight':
        logfile = f'{specproddir}/run/scripts/night/{args.night}/{jobdesc}-{args.night}-{tileid}-{qid}.log'
    elif jobdesc in ('cumulative', 'pernight'):
        logfile = get_ztile_script_pathname(tileid, jobdesc, night=args.night)
        logfile = logfile.replace('.slurm', f'-{qid}.log')
    elif jobdesc == 'perexp':
        logfile = get_ztile_script_pathname(tileid, jobdesc, expid=e1)
        logfile = logfile.replace('.slurm', f'-{qid}.log')
    else:
        logfile = f'{specproddir}/run/scripts/night/{args.night}/{jobdesc}-{args.night}-{e1:08d}-{proccamword}-{qid}.log'

    logurl = os.path.relpath(logfile, outputdir)

    if jobdesc in ('tilenight', 'pernight', 'perexp', 'cumulative'):
        description = f'{jobdesc} {tileid}'
        if jobdesc == 'perexp':
            description += f'-{expid}'
    elif expid.count(',') == 0:
        #- single expid job, e.g. arc, flat, prestdstar, ...
        description = f'{jobdesc} {expid}'
    else:
        #- multi-expid job, e.g. stdstar fits
        description = f'{jobdesc}'

    timing = ''
    if state == 'COMPLETED':
        hh, mm, ss = jobinfo[qid]['ELAPSED'].split(':')
        minutes = int(hh)*60 + int(mm)
        timing = f'<br>{minutes}m{ss}s'

    description += f'<br><small><it>Job {qid}<br>{state}{timing}</it></small>'

    tooltip = f'{jobdesc} Exp {expid} Job {qid} {state}'
    if tileid > 0:
        tooltip = f'Tile {tileid} {tooltip}'

    fx.write(f'    {intid}({description}):::{state}\n')
    fx.write(f'    click {intid} "{logurl}" "{tooltip}"\n')
    for depid in row['INT_DEP_IDS'].split('|'):
        if depid != '':
            fx.write(f'    {depid} --> {intid}\n')

fx.write(""" </div>

<script src="https://unpkg.com/mermaid@9.1.7/dist/mermaid.min.js">
    mermaid.initialize({
        startOnLoad:true
    });
</script>
</body>
</html>
""")
fx.close()

print(f'Wrote {args.output}')
if outputurl != args.output:
    print(outputurl)

for state, n in state_counter.items():
    print(f"{state:12s} {n:2}")



