#!/cvmfs/icecube.opensciencegrid.org/iceprod/v2.7.1/env-shell.sh python
from rest_tools.client import SavedDeviceGrantAuth
from collections import Counter
from pprint import pprint

api = SavedDeviceGrantAuth(
    address='https://iceprod2-api.icecube.wisc.edu',
    token_url='https://keycloak.icecube.wisc.edu/auth/realms/IceCube',
    filename='.iceprod-auth',
    client_id='iceprod-public'
)

# get a list of datasets
#result = api.request_seq('GET', '/datasets', {'keys': 'dataset|jobs_submitted'})
#for dataset_id, metadata in result.items():
#    #print(metadata)
#    if metadata['dataset'] == 22794:
#        break

dataset_id = 'afd13fa236bc11ee9cd300505684797b'

# for the last dataset, get a list of tasks
result = api.request_seq('GET', f'/datasets/{dataset_id}/tasks', {'keys': 'job_index|name|status|site'})
log_task_ids = []
bad_site = Counter()
for task_id, metadata in result.items():
    #print(metadata)
    if metadata['status'] == 'suspended':
        log_task_ids.append(task_id)
        bad_site[metadata['site']] += 1

print('failed tasks', log_task_ids)

pprint(bad_site)

# for the last completed task, get the last logs
errors = Counter()
for log_task_id in log_task_ids[:100]:
    print('.',end='', flush=True)
    try:
        result = api.request_seq('GET', f'/datasets/{dataset_id}/tasks/{log_task_id}/logs', {'group': True})
    except Exception:
        continue
    for log in result['logs']:
        if log['name'] == 'stderr':
            last_lines = log['data'].strip('\n').split('\n')[-1]
            if 'dCorsika' in last_lines:
                last_lines = last_lines.split('!')[0]
            errors[last_lines] += 1

pprint(errors)