#!/usr/bin/env python3
#
# Copyright (C) 2013-2016 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 json
import argparse
import os
import fcntl
from dxpy.utils.resolver import *
from dxpy.utils.printing import *

from contextlib import contextmanager

parser = argparse.ArgumentParser(formatter_class=argparse.RawTextHelpFormatter,
                                 description=fill('Reads and modifies job_output.json in your home directory to be a JSON hash with key *name* and value  *value*.') + '\n\n' + fill('If --class is not provided or is set to "auto", auto-detection of the output format will occur.  In particular, it will treat it as a number, hash, or boolean if it can be successfully parsed as such.  If it is a string which matches the pattern for a data object ID, it will encapsulate it in a DNAnexus link hash; otherwise it is treated as a simple string.') + '\n\n' + fill('Use --array to append to an array of values or prepend "array:" to the --class argument.') + '\n\n' + fill('To use the output of another job as part of your output, use --class=jobref (which will throw an error if it is not formatted correctly), or use the automatic parsing which will recognize anything starting with a job ID as a job-based object reference.  You should format the value as follows:') + '''

  Format: <job ID>:<output field name>
  Example: dx-jobutil-add-output myoutput --class=jobref \\
             job-B2JKYqK4Zg2K915yQxPQ0024:other_output

'''
                                 + fill('Analysis references can be specified similarly with --class=analysisref and formatted as:') + '''

  Format: <analysis ID>:<stage ID>.<output field name>
          <analysis ID>:<exported output field name>
  Example: dx-jobutil-add-output myoutput --class=analysisref \\
             analysis-B2JKYqK4Zg2K915yQxPQ0024:some_output
''')
parser.add_argument('name', help='Name of the output field name')
parser.add_argument('value', help='Value of the output field')
parser.add_argument('--class', dest='classname', metavar='CLASSNAME',
                    choices=['int', 'float', 'string', 'boolean', 'hash',
                             'array:int', 'array:float', 'array:string', 'array:boolean',
                             'file', 'record', 'applet', 'workflow',
                             'array:file', 'array:record', 'array:applet',
                             'array:workflow', 'jobref', 'array:jobref', 'analysisref',
                             'array:analysisref', 'auto', 'array:auto'],
                    help=fill('Class of output for formatting purposes, e.g. "int"; default "auto"',
                              width_adjustment=-24),
                    nargs='?',
                    default='auto')
parser.add_argument('--array', help='Output field is an array', action='store_true')
# replacement output filename to use instead of "~/job_output.json" (primarily for testing)
parser.add_argument('-o', '--output', help=argparse.SUPPRESS,
                    default=os.path.expanduser(os.path.join('~', 'job_output.json')))
args = parser.parse_args()

value = None

if args.classname.startswith('array:'):
    args.array = True
    args.classname = args.classname[6:]

if args.classname in ['file', 'record', 'applet', 'workflow']:
    args.classname = 'dxobject'

if args.classname == 'auto':
    if is_data_obj_id(args.value):
        value = {'$dnanexus_link': args.value}
    else:
        colon_substrings = split_unescaped(':', args.value)
        if len(colon_substrings) == 2 and (is_job_id(colon_substrings[0]) or \
                                           is_localjob_id(colon_substrings[0]) or \
                                           is_analysis_id(colon_substrings[0])):
            if is_analysis_id(colon_substrings[0]):
                value = {"$dnanexus_link": {"analysis": colon_substrings[0],
                                            "field": colon_substrings[1]}}
            else:
                value = {"job": colon_substrings[0],
                         "field": unescape_name_str(colon_substrings[1])}
        else:
            try:
                # works for int, float, boolean, hash, dxobject (when already in the hash)
                value = json.loads(args.value)
            except:
                # string
                value = args.value
elif args.classname == 'int':
    try:
        value = int(args.value)
    except:
        parser.exit(3, 'Value could not be parsed as an int\n')
elif args.classname == 'float':
    try:
        value = float(args.value)
    except:
        parser.exit(3, 'Value could not be parsed as a float\n')
elif args.classname == 'string':
    value = args.value
elif args.classname == 'boolean':
    if args.value.lower().startswith('t') or args.value == '1':
        value = True
    elif args.value.lower().startswith('f') or args.value == '0':
        value = False
    else:
        parser.exit(3, 'Value could not be parsed as a boolean\n')
elif args.classname == 'hash':
    try:
        value = json.loads(args.value)
    except:
        parser.exit(3, 'Value could not be parsed as a hash\n')
elif args.classname == 'dxobject':
    if is_data_obj_id(args.value):
        value = {'$dnanexus_link': args.value}
    else:
        try:
            value = json.loads(args.value)
        except:
            parser.exit(3, 'Value could not be parsed as a data object ID or hash with key "$dnanexus_link"\n')
        if not isinstance(value, dict) or value.get('$dnanexus_link', None) is None:
            parser.exit(3, 'Value could not be parsed as a data object ID or hash with key "$dnanexus_link"\n')
elif args.classname == 'jobref':
    colon_substrings = split_unescaped(':', args.value)
    if len(colon_substrings) == 2 and (is_job_id(colon_substrings[0]) or \
                                       is_localjob_id(colon_substrings[0])):
        value = {"job": colon_substrings[0],
                 "field": unescape_name_str(colon_substrings[1])}
    else:
        parser.exit(3, fill('Value could not be parsed as a job-based object reference with format <job ID>:<job output field name>') + '\n')
elif args.classname == 'analysisref':
    colon_substrings = split_unescaped(':', args.value)
    if len(colon_substrings) == 2 and is_analysis_id(colon_substrings[0]):
        value = {"$dnanexus_link": {"analysis": colon_substrings[0],
                                    "field": colon_substrings[1]}}
    else:
        parser.exit(3, fill('Value could not be parsed as an analysis reference with format <analysis ID>:<analysis output field name>') + '\n')


# Note, this operation is not immune to race conditions, see: https://linux.die.net/man/2/fcntl
@contextmanager
def flock(fd):
    fcntl.flock(fd, fcntl.LOCK_EX)
    try:
        yield
    finally:
        fcntl.flock(fd, fcntl.LOCK_UN)


def update_output_json(output_json):
    output_spec = json.loads(output_json)
    if args.array:
        if args.name in output_spec:
            if isinstance(output_spec[args.name], list):
                output_spec[args.name].append(value)
            else:
                parser.exit(3, 'Key ' + args.name + ' was found in existing output but was not an array\n')
        else:
            output_spec[args.name] = [value]
    else:
        output_spec[args.name] = value
    return json.dumps(output_spec, indent=4) + "\n"


try:
    # Atomically test for file existence and create it if it doesn't
    # exist. Below we branch on whether the file existed or not at the time
    # that this was tried.
    output_fd = os.open(args.output, os.O_WRONLY | os.O_CREAT | os.O_EXCL)
except OSError as e:
    # File exists, but we don't have a handle to it yet. Open it and lock for
    # exclusive use
    with open(args.output, 'r+') as output_file:
        with flock(output_file):
            output_json = output_file.read()
            output_file.seek(0)
            output_file.write(update_output_json(output_json))
            output_file.truncate()
else:
    # File didn't exist before; we now have an FD on it. Lock it for exclusive
    # use
    with os.fdopen(output_fd, 'w') as output_file:
        with flock(output_file):
            output_file.write(update_output_json("{}"))
