from subprocess import check_output
import pandas as pd
from datetime import datetime, timedelta
from io import StringIO
from json import loads
from os.path import join

from qiita_db.util import MaxRSS_helper
from qiita_db.exceptions import QiitaDBUnknownIDError
from qiita_db.processing_job import ProcessingJob
from qiita_db.software import Software


all_commands = [c for s in Software.iter(False) for c in s.commands]

# retrieving only the numerice external_id means that we are only focusing
# on barnacle2/slurm jobs
main_jobs = [j for c in all_commands for j in c.processing_jobs
             if j.status == 'success' and j.external_id.isnumeric()]

sacct = ['sacct', '-p', '--format=JobID,ElapsedRaw,MaxRSS,Submit,Start,MaxRSS,'
         'CPUTimeRAW,ReqMem,AllocCPUs,AveVMSize', '-j']

data = []
for i, j in enumerate(main_jobs):
    if i % 1000 == 0:
        print(f'{i}/{len(main_jobs)}')
    eid = j.external_id
    extra_info = ''
    rvals = StringIO(check_output(sacct + [eid]).decode('ascii'))
    _d = pd.read_csv(rvals, sep='|')
    _d['QiitaID'] = j.id
    cmd = j.command
    s = j.command.software
    try:
        samples, columns, input_size = j.shape
    except QiitaDBUnknownIDError:
        # this will be raised if the study or the analysis has been deleted;
        # in other words, the processing_job was ran but the details about it
        # were erased when the user deleted them - however, we keep the job for
        # the record
        continue
    except TypeError as e:
        # similar to the except above, exept that for these 2 commands, we have
        # the study_id as None
        if cmd.name in {'create_sample_template', 'delete_sample_template',
                        'list_remote_files'}:
            continue
        else:
            raise e

    sname = s.name

    if cmd.name == 'release_validators':
        ej = ProcessingJob(j.parameters.values['job'])
        extra_info = ej.command.name
        samples, columns, input_size = ej.shape
    elif cmd.name == 'complete_job':
        artifacts = loads(j.parameters.values['payload'])['artifacts']
        if artifacts is not None:
            extra_info = ','.join({
                x['artifact_type'] for x in artifacts.values()
                if 'artifact_type' in x})
    elif cmd.name == 'Validate':
        input_size = sum([len(x) for x in loads(
            j.parameters.values['files']).values()])
        sname = f"{sname} - {j.parameters.values['artifact_type']}"
    elif cmd.name == 'Alpha rarefaction curves [alpha_rarefaction]':
        extra_info = j.parameters.values[
            ('The number of rarefaction depths to include between min_depth '
             'and max_depth. (steps)')]

    _d['external_id'] = eid
    _d['sId'] = s.id
    _d['sName'] = sname
    _d['sVersion'] = s.version
    _d['cId'] = cmd.id
    _d['cName'] = cmd.name
    _d['samples'] = samples
    _d['columns'] = columns
    _d['input_size'] = input_size
    _d['extra_info'] = extra_info
    _d.drop(columns=['Unnamed: 10'], inplace=True)
    data.append(_d)

data = pd.concat(data)

# In slurm, each JobID is represented by 3 rows in the dataframe:
#   - external_id:  overall container for the job and its associated
#                   requests. When the Timelimit is hit, the container
#                   would take care of completing/stopping the
#                   external_id.batch job.
#   - external_id.batch: it's a container job, it provides how
#                        much memory it uses and cpus allocated, etc.
#   - external_id.extern: takes into account anything that happens
#                         outside processing but yet is included in
#                         the container resources. As in, if you ssh
#                         to the node and do something additional or run
#                         a prolog script, that processing would be under
#                         external_id but separate from external_id.batch.
# Here we are going to merge all this info into a single row + some
# other columns
date_fmt = '%Y-%m-%dT%H:%M:%S'

df = []
for eid, __df in data.groupby('external_id'):
    tmp = __df.iloc[1].copy()
    # Calculating WaitTime, basically how long did the job took to start
    # this is useful for some general profiling
    tmp['WaitTime'] = datetime.strptime(
        __df.iloc[0].Start, date_fmt) - datetime.strptime(
        __df.iloc[0].Submit, date_fmt)
    df.append(tmp)
df = pd.DataFrame(df)

# This is important as we are transforming the MaxRSS to raw value
# so we need to confirm that there is no other suffixes
print('Make sure that only 0/K/M exist', set(
    df.MaxRSS.apply(lambda x: str(x)[-1])))

# Generating new columns
df['MaxRSSRaw'] = df.MaxRSS.apply(lambda x: MaxRSS_helper(str(x)))
df['ElapsedRawTime'] = df.ElapsedRaw.apply(
    lambda x: timedelta(seconds=float(x)))

# Thu, Apr 27, 2023 was the first time Jeff and I changed the old allocations
# (from barnacle) to a better allocation so using job 1265533 as the
# before/after so we only use the latests for the newest version
df['updated'] = df.external_id.apply(
    lambda x: 'after' if int(x) >= 1265533 else 'before')

fn = join('/panfs', 'qiita', f'jobs_{df.Start.max()[:10]}.tsv.gz')
df.to_csv(fn, sep='\t', index=False)