#!/usr/bin/env python3
#
# python exec entry-point, called directly from RiskScape, for marshalling tuples
# to and from
import os
import sys
import logging
import traceback
import serializer
from serializer import Boolean, RiskscapePyException, Text
import adapter


# comment out this next line for debugging to file

# first arg is the filepath containing the user's python function
filepath = sys.argv[1]
# tuple type coming in
input_type = sys.argv[2]
# stated type coming out
output_type = sys.argv[3]
# optional parameters arg
has_parameters = len(sys.argv) > 4
parameters = parameter_type = None
if has_parameters:
    parameter_type = sys.argv[4]

def wrong_number_args(actual, actual_str, expected):
    error = "Your python() pipeline step is passing a different number of arguments to what your Python script expects. "
    if actual > 2 or actual == 0:
        # user is way off, they clearly have a fundamental misunderstanding of how the python() step works
        return error + ("Your Python function takes {0} arguments, but the python() pipeline step can only pass "
            "1-2 arguments to your Python script. The first argument is the rows of pipeline data (via a generator function). "
            "A second optional argument can pass through parameters from your pipeline python() step. ").format(actual_str)
    elif expected == 2:
        # python() step has a parameters argument
        return error + ("The python() pipeline step uses parameters, but the Python function is missing "
            "a second 'parameters' argument").format(actual)
    else:
        return error + ("Your Python function takes {0} arguments, but the python() pipeline step is currently only passing "
            "1 argument (which is the rows of pipeline data, via a generator function). "
            "Your Python function may be expecting parameters that are missing from the python() pipeline step").format(actual_str)

adapter.stop_sdterr_blocking()
logger = adapter.logger

(input, output) = adapter.initialize_io()

try:
    # the user's Python function can accept an optional 2nd 'parameters' argument
    expected_args = 2 if has_parameters else 1
    user_function = adapter.extract_callable(filepath, expected_args, wrong_number_args)

    # Assign magic 'model_output' function to the scope
    user_function.__globals__['model_output'] = adapter.model_output

    adapter.ready(output)

    def tuple_generator():
        while True:
            # for now this will be one or zero, but I'm going to make this work with a
            # generator function so the function can control the consumption of tuples
            row_count = serializer.read('Integer', input)

            sys.stderr.flush()

            # done
            if row_count == 0:
                return

            ctr = 0
            while ctr < row_count:
                # read the input args that the main Java process passed to us via stdin
                yield serializer.read(input_type, input)
                ctr += 1

    try:
        # when optional parameters are specified, the parameter values get serialized first. Read them
        # before we create the tuple generator (otherwise they'll get treated as pipeline data)
        if has_parameters:
            parameters = serializer.read(parameter_type, input)

        tuples = tuple_generator()
        args = [tuples]
        if has_parameters:
          args.append(parameters)

        result_generator = user_function(*args)

        output_expected = output_type != 'Struct[]'

        if result_generator is not None:
            if not output_expected:
                raise Exception("No result-type was specified, but your function yielded values")

            for result in result_generator:
                serializer.write('Integer', output, 1)
                serializer.write(output_type, output, result)
        elif output_expected:
            # hopefully the rs code will spot this and report a nice error
            raise Exception("Output expected: {0}".format(output_type))

        # Loop through any tuples that weren't already consumed by the user function so that riskscape doesn't
        # try and write to a stream that's closed
        for tuple in tuples:
            pass

    except RiskscapePyException as rse:
        # we hit an error in the serialization code. The traceback is probably
        # completely meaningless to the user, so just display the error message
        print('Error: %s' % rse, file=sys.stderr)
        exit(1)
    except Exception as e:
        adapter.handle_function_call_failure(output, e)

    # tell riskscape we are done.
    serializer.write('Integer', output, 0)
    output.flush()

    adapter.send_outputs(output)

except RiskscapePyException as rse:
    # Something went wrong in extract_callable() - there's probably no `def function()`
    # As above, the traceback won't be very meaningful, so just display the error.
    print('Error: %s' % rse, file=sys.stderr)
    exit(1)
except Exception as e:
    # we hit an error outside of the user's code that wasn't a serialization failure.
    # This is very likely to be an internal error that only a dev will be able to solve.
    # They will probably see the last thing that was printed to stderr. They can then
    # put that in a bug report - result!
    traceback.print_tb(e.__traceback__.tb_next, file=sys.stderr)
    # Now print out the exception without any stack trace.
    # Note Python works the exception here from the current 'except' case.
    # We could use other traceback APIs here, but they've changed a bit
    # between 3.5 and 3.10, so might cause compatability problems
    traceback.print_exc(limit=0, file=sys.stderr)
    exit(1)
finally:
    output.close()
