#!/usr/bin/env python
#
# Copyright (C) 2013 DNAnexus, Inc.
#
# This file is part of dx-toolkit (DNAnexus platform client libraries).
#
#   Licensed under the Apache License, Version 2.0 (the "License"); you may not
#   use this file except in compliance with the License. You may obtain a copy
#   of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
#   WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
#   License for the specific language governing permissions and limitations
#   under the License.

import os, sys, json, fileinput, re
import argparse
from collections import defaultdict
import dxpy

parser = argparse.ArgumentParser(description="Get stdout/stderr streams for a job and, optionally, its subjobs.")
parser.add_argument("-1", "--only-this-job", help="Do not recurse into subjobs' output streams", action='store_true')
#parser.add_argument("-f", "--print-to-files", help="Instead of printing a summary, print each stream into a file <job_id>.stdout, etc. in the working directory", action='store_true')
parser.add_argument("-q", "--quiet", help="Do not print banners when using summary format", action='store_true')
parser.add_argument("-n", "--no-tag", help="Do not prepend job tag to every line", action='store_true')
#parser.add_argument("--force", help="Force to run", action='store_true')
parser.add_argument("job_id", help="Job ID")
args = parser.parse_args()

# if not args.force:
#     print >>sys.stderr, "dx-get-job-streams is deprecated. Please use `dx watch` instead."
#     print >>sys.stderr, "To run dx-get-job-streams anyway, add the --force option."
#     exit(1)

if args.job_id == None:
    parser.print_help()
    parser.error("Missing job ID")

job_id = args.job_id
root_job = dxpy.DXJob(job_id)

if args.only_this_job:
    jobs = [root_job]
else:
    jobs = [dxpy.DXJob(job['id']) for job in dxpy.find_jobs(origin_job=job_id)]
    if len(jobs) == 0:
        jobs = [root_job]
        jobs.extend([dxpy.DXJob(job['id']) for job in dxpy.find_jobs(parent_job=job_id)])
        # TODO: recurse? This only pulls in the first level children

job_descriptions = {}; subjobs = {}; stdouts = {}; stderrs = {}

for job in jobs:
    desc = job.describe()
    job_descriptions[job.get_id()] = desc
    subjobs.setdefault(desc['parentJob'], [])
    if desc['parentJob'] != desc['id'] and not args.only_this_job:
        subjobs[desc['parentJob']].append(desc['id'])
    stdouts[job.get_id()] = desc.get('stdout', None)
    stderrs[job.get_id()] = desc.get('stderr', None)

def print_streams(job_id, indent_level):
    INDENT = '    '*indent_level
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKGREEN = '\033[92m'
    WARN = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
    desc = job_descriptions[job_id]
    applet_or_app = desc['applet'] if 'applet' in desc else desc['app']

    status = desc['state']
    if status == 'failed':
        status = FAIL + 'FAILED' + OKGREEN
    if not args.quiet:
        print INDENT + OKGREEN + '==== Job', desc['id'], '(' + applet_or_app + ':' + HEADER + desc['function'] + OKGREEN + ') (' + status + ')', 'stdout', '====', ENDC
    tag = '' if args.no_tag else OKBLUE + '%(id)s:%(function)s> ' % job_descriptions[job_id] + ENDC
    if stdouts[job_id] != None:
        for line in dxpy.DXFile(stdouts[job_id]):
            print INDENT + tag + line.decode("utf-8")

    status = desc['state']
    if status == 'failed':
        status = FAIL + 'FAILED' + WARN
    if not args.quiet:
        print INDENT + WARN + '==== Job', desc['id'], '(' + applet_or_app + ':' + HEADER + desc['function'] + WARN + ') (' + status + ')', 'stderr', '====', ENDC
    tag = '' if args.no_tag else OKBLUE + '%(id)s:%(function)s> ' % job_descriptions[job_id] + ENDC
    if stderrs[job_id] != None:
        for line in dxpy.DXFile(stderrs[job_id]):
            print INDENT + tag + line.decode("utf-8")

    for subjob_id in subjobs.get(job_id, []):
        print_streams(subjob_id, indent_level+1)

print_streams(job_id, 0)
